Create doc_searcher.py

doc_searcher.py

@@ -0,0 +1,65 @@
+from qdrant_client import QdrantClient
+from fastembed import SparseTextEmbedding, LateInteractionTextEmbedding
+from qdrant_client import QdrantClient, models
+from sentence_transformers import SentenceTransformer
+from config import DENSE_MODEL, SPARSE_MODEL, LATE_INTERACTION_MODEL, QDRANT_URL, QDRANT_API_KEY,HUGGING_FACE_API_KEY
+
+class DocSearcher:
+    
+    def __init__(self, collection_name):
+        self.collection_name = collection_name
+        self.dense_model = SentenceTransformer(DENSE_MODEL,device="cpu",token=HUGGING_FACE_API_KEY)
+        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)
+    
+    async def search(self, text: str):
+        
+        dense_query = self.dense_model.encode(text).tolist()
+        sparse_query = next(self.sparse_model.query_embed(text))
+
+        prefetch = [
+                models.Prefetch(
+                    query=dense_query,
+                    using=DENSE_MODEL,
+                    params=models.SearchParams(
+                    quantization=models.QuantizationSearchParams(
+                            rescore=False,
+                        ),
+                    ),
+                    limit=200
+                ),
+                models.Prefetch(
+                    query=models.SparseVector(**sparse_query.as_object()),
+                    using=SPARSE_MODEL,
+                    params=models.SearchParams(
+                    quantization=models.QuantizationSearchParams(
+                            rescore=False,
+                        ),
+                    ),
+                    limit=200
+                )
+        ]
+
+        search_result = self.qdrant_client.query_points(
+            collection_name= self.collection_name,
+            search_params=models.SearchParams(
+                hnsw_ef=128,
+                quantization=models.QuantizationSearchParams(
+                        rescore=True,
+                    ),
+                ),
+            prefetch=prefetch,
+            query=models.FusionQuery(
+                fusion=models.Fusion.RRF,
+            ),
+            with_payload=True,
+            limit = 10
+        ).points
+
+        data = []
+
+        for hit in search_result:
+            data.append(hit.payload)
+        
+        return data