added rerank component

app.py

@@ -4,6 +4,8 @@ import gradio as gr
 from openai import OpenAI
 from qdrant_client import QdrantClient
 from sentence_transformers import SentenceTransformer
+from flashrank import Ranker, RerankRequest
+from types import SimpleNamespace
 
 API_KEY = os.environ.get('DEEPSEEK_API_KEY')
 BASE_URL = "https://api.deepseek.com"
@@ -40,6 +42,8 @@ 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()
     
@@ -47,18 +51,34 @@ class HFRAG:
       results = self.db_client.search(
         collection_name=COLLECTION_NAME,
         query_vector=query_vector,
-        limit=top_k,
+        limit=20, # 扩大召回范围，之后进行重排序
         score_threshold=score_threshold
       )
     else:
       results = self.db_client.query_points(
         collection_name=COLLECTION_NAME,
         query=query_vector,
-        limit=top_k,
+        limit=20,
         with_payload=True,
         score_threshold=score_threshold
       ).points
-    return results
+    
+    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
 
   def format_context(self, search_results):
     context_pieces = []

requirements.txt

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