adds reranker

doc_searcher.py

@@ -2,6 +2,7 @@ from qdrant_client import QdrantClient
 from qdrant_client.models import Filter, FieldCondition, MatchValue, MatchText
 from fastembed import SparseTextEmbedding, LateInteractionTextEmbedding
 from qdrant_client import QdrantClient, models
+from reranker import Reranker
 from sentence_transformers import SentenceTransformer
 from config import DENSE_MODEL, SPARSE_MODEL, LATE_INTERACTION_MODEL, QDRANT_URL, QDRANT_API_KEY,HUGGING_FACE_API_KEY
 
@@ -9,8 +10,9 @@ class DocSearcher:
     
     def __init__(self, collection_name):
         self.collection_name = collection_name
+        self.reranker = Reranker()
         self.dense_model = SentenceTransformer(DENSE_MODEL,device="cpu",token=HUGGING_FACE_API_KEY)
-        self.model = SentenceTransformer("Qwen/Qwen3-Embedding-4B",device="cpu")
+        self.model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B",device="cpu")
         self.sparse_model = SparseTextEmbedding(SPARSE_MODEL)
         self.late_interaction_model = LateInteractionTextEmbedding(LATE_INTERACTION_MODEL)
         self.qdrant_client = QdrantClient(QDRANT_URL,api_key=QDRANT_API_KEY,timeout=30)
@@ -101,35 +103,30 @@ class DocSearcher:
     
     async def search_temp(self, text: str):
         
+        queries = [text]
         dense_query = self.model.encode(text).tolist()
         # sparse_query = next(self.sparse_model.query_embed(text))
 
         prefetch = [
                 models.Prefetch(
                     query=dense_query,
-                    using="Qwen/Qwen3-Embedding-4B",
-                    limit=10
+                    using="Qwen/Qwen3-Embedding-0.6B",
+                    limit=100
                 ),
-                # models.Prefetch(
-                #     query=models.SparseVector(**sparse_query.as_object()),
-                #     using=SPARSE_MODEL,
-                #     limit=100
-                # )
         ]
             
         search_result = self.qdrant_client.query_points(
             collection_name= "sl-list",
             prefetch=prefetch,
-            query=models.FusionQuery(
-                fusion=models.Fusion.RRF,
-            ),
             with_payload=True,
-            limit = 10,
+            limit = 100,
         ).points
 
         data = []
 
         for hit in search_result:
-            data.append(hit.payload)
+            data.append(hit.payload["tekst"])
+
+        scores = self.reranker.compute_logits(queries,data)
         
-        return data
+        return scores

reranker.py

@@ -0,0 +1,49 @@
+import torch
+from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
+
+class Reranker:
+    
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Reranker-4B", padding_side='left')
+        self.model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-Reranker-4B").eval()
+
+    def format_instruction(instruction, query, doc):
+        if instruction is None:
+            instruction = 'Given a web search query, retrieve relevant passages that answer the query'
+        output = "<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}".format(instruction=instruction,query=query, doc=doc)
+        return output
+      
+    def process_inputs(self,pairs):
+        prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n"
+        suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
+        prefix_tokens = self.tokenizer.encode(prefix, add_special_tokens=False)
+        suffix_tokens = self.tokenizer.encode(suffix, add_special_tokens=False)
+        max_length = 2048
+        inputs = self.tokenizer(
+            pairs, padding=False, truncation='longest_first',
+            return_attention_mask=False, max_length=max_length - len(prefix_tokens) - len(suffix_tokens)
+        )
+        for i, ele in enumerate(inputs['input_ids']):
+            inputs['input_ids'][i] = prefix_tokens + ele + suffix_tokens
+        inputs = self.tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=max_length)
+        for key in inputs:
+            inputs[key] = inputs[key].to(self.model.device)
+        return inputs
+
+    @torch.no_grad
+    def compute_logits(self,queries,documents):
+        token_false_id = self.tokenizer.convert_tokens_to_ids("no")
+        token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
+
+        task = 'Given a web search query, retrieve relevant passages that answer the query'
+        pairs = [self.format_instruction(task, query, doc) for query, doc in zip(queries, documents)]
+        inputs = self.process_inputs(pairs)
+        
+        batch_scores = self.model(**inputs).logits[:, -1, :]
+        true_vector = batch_scores[:, token_true_id]
+        false_vector = batch_scores[:, token_false_id]
+        batch_scores = torch.stack([false_vector, true_vector], dim=1)
+        batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
+        scores = batch_scores[:, 1].exp().tolist()
+
+        return scores