removed rerank part, added hybrid index

app.py

@@ -2,10 +2,9 @@ import os
 import httpx
 import gradio as gr
 from openai import OpenAI
-from qdrant_client import QdrantClient
+from qdrant_client import QdrantClient, models
 from sentence_transformers import SentenceTransformer
-from flashrank import Ranker, RerankRequest
-from types import SimpleNamespace
+from fastembed import SparseTextEmbedding
 
 API_KEY = os.environ.get('DEEPSEEK_API_KEY')
 BASE_URL = "https://api.deepseek.com"
@@ -13,10 +12,12 @@ BASE_URL = "https://api.deepseek.com"
 QDRANT_PATH = "./qdrant_db"
 COLLECTION_NAME = "huggingface_transformers_docs"
 EMBEDDING_MODEL_ID = "fyerfyer/finetune-jina-transformers-v1"
+SPARSE_MODEL_ID = "prithivida/Splade_PP_en_v1"
 
 class HFRAG:
   def __init__(self):
-    self.embed_model = SentenceTransformer(EMBEDDING_MODEL_ID, trust_remote_code=True)
+    self.dense_model = SentenceTransformer(EMBEDDING_MODEL_ID, trust_remote_code=True)
+    self.sparse_model = SparseTextEmbedding(model_name=SPARSE_MODEL_ID)
     
     lock_file = os.path.join(QDRANT_PATH, ".lock")
     if os.path.exists(lock_file):
@@ -42,51 +43,49 @@ class HFRAG:
       http_client=httpx.Client(proxy=None, trust_env=False)
     )
 
-    self.reranker = Ranker(model_name="ms-marco-TinyBERT-L-2-v2", cache_dir="/tmp")
-
-  def retrieve(self, query: str, top_k: int = 5, score_threshold: float = 0.40):
-    query_vector = self.embed_model.encode(query).tolist()
+  def retrieve(self, query: str, top_k: int = 5):
+    # Generate dense vector
+    query_dense_vec = self.dense_model.encode(query).tolist()
+    
+    # Generate sparse vector
+    query_sparse_gen = list(self.sparse_model.embed([query]))[0]
+    query_sparse_vec = models.SparseVector(
+      indices=query_sparse_gen.indices.tolist(),
+      values=query_sparse_gen.values.tolist()
+    )
     
-    if hasattr(self.db_client, 'search'):
-      results = self.db_client.search(
-        collection_name=COLLECTION_NAME,
-        query_vector=query_vector,
-        limit=20, # 扩大召回范围，之后进行重排序
-        score_threshold=score_threshold
-      )
-    else:
-      results = self.db_client.query_points(
-        collection_name=COLLECTION_NAME,
-        query=query_vector,
-        limit=20,
-        with_payload=True,
-        score_threshold=score_threshold
-      ).points
+    # Create prefetch for dense retrieval
+    prefetch_dense = models.Prefetch(
+      query=query_dense_vec,
+      using="text-dense",
+      limit=20,
+    )
     
-    passages = [
-      {"id": result.payload['metadata']['source'], "text": result.payload['text'], "meta": result.payload}
-      for result in results
-    ]
-    rerank_request = RerankRequest(query=query, passages=passages)
-    reranked_results = self.reranker.rerank(rerank_request)
-
-    # 从重排序后的序列中取出 TopK
-    final_results = []
-    for item in reranked_results[:top_k]:
-      final_result = SimpleNamespace()
-      final_result.payload = item['meta']
-      final_result.score = item['score']
-      final_results.append(final_result)
-
-    return final_results
+    # Create prefetch for sparse retrieval
+    prefetch_sparse = models.Prefetch(
+      query=query_sparse_vec,
+      using="text-sparse",
+      limit=20,
+    )
+    
+    # Hybrid search with RRF fusion
+    results = self.db_client.query_points(
+      collection_name=COLLECTION_NAME,
+      prefetch=[prefetch_dense, prefetch_sparse],
+      query=models.FusionQuery(fusion=models.Fusion.RRF),
+      limit=top_k,
+      with_payload=True
+    ).points
+
+    return results
 
   def format_context(self, search_results):
     context_pieces = []
     sources_summary = []
     
     for idx, hit in enumerate(search_results, 1):
-      raw_source = hit.payload['metadata']['source']
-      filename = raw_source.split('/')[-1] 
+      raw_source = hit.payload.get('source', 'unknown')
+      filename = raw_source.split('/')[-1] if '/' in raw_source else raw_source
       text = hit.payload['text']
       score = hit.score
       

qdrant_db/collection/huggingface_transformers_docs/storage.sqlite

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88a55f2d047299d73d59f44f05d0ef0bf03ca865ae5dbd5523eed72269cb0f98
-size 56549376
+oid sha256:767ba990e571262333521d2528e9d57f248e9cd63f6269907716bea20617c607
+size 62464000

qdrant_db/meta.json

@@ -1 +1 @@
-{"collections": {"huggingface_transformers_docs": {"vectors": {"size": 768, "distance": "Cosine", "hnsw_config": null, "quantization_config": null, "on_disk": null, "datatype": null, "multivector_config": null}, "shard_number": null, "sharding_method": null, "replication_factor": null, "write_consistency_factor": null, "on_disk_payload": null, "hnsw_config": null, "wal_config": null, "optimizers_config": null, "quantization_config": null, "sparse_vectors": null, "strict_mode_config": null, "metadata": null}}, "aliases": {}}
+{"collections": {"huggingface_transformers_docs": {"vectors": {"text-dense": {"size": 768, "distance": "Cosine", "hnsw_config": null, "quantization_config": null, "on_disk": null, "datatype": null, "multivector_config": null}}, "shard_number": null, "sharding_method": null, "replication_factor": null, "write_consistency_factor": null, "on_disk_payload": null, "hnsw_config": null, "wal_config": null, "optimizers_config": null, "quantization_config": null, "sparse_vectors": {"text-sparse": {"index": {"full_scan_threshold": null, "on_disk": true, "datatype": null}, "modifier": null}}, "strict_mode_config": null, "metadata": null}}, "aliases": {}}

requirements.txt

@@ -5,4 +5,4 @@ sentence-transformers
 httpx
 torch
 python-dotenv
-flashrank
+fastembed