revert changes

app.py

@@ -13,11 +13,11 @@ app = FastAPI()
 
 neural_searcher = NeuralSearcher(collection_name=os.getenv('COLLECTION_NAME'))
 
-REQUEST_TIMEOUT_ERROR = 1
+REQUEST_TIMEOUT_ERROR = 30
 
 @app.get("/api/search")
-def search(q: str):
-    data =  neural_searcher.search(text=q)
+async def search(q: str):
+    data =  await neural_searcher.search(text=q)
     return data
     
 

ner.py

@@ -0,0 +1,97 @@
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained("kalusev/NER4Legal_SRB", use_auth_token=True)
+model = AutoModelForTokenClassification.from_pretrained("kalusev/NER4Legal_SRB", use_auth_token=True).to(device)
+
+id_to_label = {
+    0: 'O',
+    1: 'B-COURT',
+    2: 'B-DATE',
+    3: 'B-DECISION',
+    4: 'B-LAW',
+    5: 'B-MONEY',
+    6: 'B-OFFICIAL GAZZETE',
+    7: 'B-PERSON',
+    8: 'B-REFERENCE',
+    9: 'I-COURT',
+    10: 'I-LAW',
+    11: 'I-MONEY',
+    12: 'I-OFFICIAL GAZZETE',
+    13: 'I-PERSON',
+    14: 'I-REFERENCE'
+}
+
+def perform_ner(text):
+    try:
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        logits = outputs.logits
+        predictions = torch.argmax(logits, dim=2).squeeze().tolist()
+
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            print("Switching to CPU due to memory constraints.")
+            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+            with torch.no_grad():
+                outputs = model.cpu()(**inputs)  # Run model on CPU
+            logits = outputs.logits
+            predictions = torch.argmax(logits, dim=2).squeeze().tolist()
+        else:
+            raise e 
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze())
+    labels = [id_to_label[pred] for pred in predictions]
+
+    results = [
+        (token, label)
+        for token, label in zip(tokens, labels)
+        if token not in tokenizer.all_special_tokens
+    ]
+    return results
+
+text = """1
+osnovni sud u bijelom polju je vrsio veliku nuzdu
+"""
+
+def merge_entities(token_label_pairs):
+    merged_words, merged_labels = [], []
+    current_word, current_label = "", None
+
+    for token, label in token_label_pairs:
+        if token.startswith("##"):
+            current_word += token[2:]
+        else:
+            if current_word:
+                merged_words.append(current_word)
+                merged_labels.append(current_label)
+
+            current_word, current_label = token, label
+
+    if current_word:
+        merged_words.append(current_word)
+        merged_labels.append(current_label)
+
+    final_words, final_labels = [], []
+    
+    for i, (word, label) in enumerate(zip(merged_words, merged_labels)):
+        if final_labels and (
+            label == final_labels[-1] or
+            (label.startswith("I-") and final_labels[-1].endswith(label[2:])) or
+            (label.startswith("B-") and final_labels[-1].endswith(label[2:]))  
+        ):  
+            
+            final_words[-1] += " " + word  
+        else:
+            final_words.append(word)
+            final_labels.append(label)
+
+    return final_words, final_labels
+
+results = perform_ner(text)
+
+words,labels = merge_entities(results)
+
+for i,b in zip(words,labels):
+    print(i + " ### " + b)

neural_searcher.py

@@ -5,45 +5,47 @@ from sentence_transformers import SentenceTransformer
 import os
 
 class NeuralSearcher:
-
+    
     def __init__(self, collection_name):
         self.collection_name = collection_name
         self.dense_model = SentenceTransformer(os.getenv('DENSE_MODEL'),device="cpu")
         self.sparse_model = SparseTextEmbedding(os.getenv('SPARSE_MODEL'))
         self.late_interaction_model = LateInteractionTextEmbedding(os.getenv('LATE_INTERACTION_MODEL'))
-        self.qdrant_client = QdrantClient(os.getenv('QDRANT_URL'),api_key=os.getenv('QDRANT_API_KEY'),https=True)
+        self.qdrant_client = QdrantClient(os.getenv('QDRANT_URL'),api_key=os.getenv('QDRANT_API_KEY'))
     
-    def search(self, text: str):
-  
+    async def search(self, text: str):
+        
         dense_query = self.dense_model.encode(text).tolist()
         sparse_query = next(self.sparse_model.query_embed(text))
-        late_query = next(self.late_interaction_model.query_embed(text))
+        # late_query = next(self.late_interaction_model.query_embed(text))
 
         prefetch = [
                 models.Prefetch(
                     query=dense_query,
                     using=os.getenv('DENSE_MODEL'),
-                    limit=100
+                    limit=200
                 ),
                 models.Prefetch(
                     query=models.SparseVector(**sparse_query.as_object()),
                     using=os.getenv('SPARSE_MODEL'),
-                    limit=100
+                    limit=200
                 )
         ]
 
-        search_result = self.qdrant_client.query_points_groups(
+        search_result = self.qdrant_client.query_points(
             collection_name= self.collection_name,
-            group_by="dbid",
             prefetch=prefetch,
-            group_size=3,
-            query=late_query,
-            using=os.getenv('LATE_INTERACTION_MODEL'),
+            query=models.FusionQuery(
+                fusion=models.Fusion.RRF,
+            ),
+            # using=os.getenv('LATE_INTERACTION_MODEL'),
             with_payload=True,
-            score_threshold=0.8,
             limit = 10
-        ).groups
+        ).points
+
+        data = []
+
+        for hit in search_result:
+            data.append(hit.payload)
         
-        for group in search_result:
-            print(group)
-        return search_result
+        return data