Step 3

backend/semantic_search.py

@@ -1,7 +1,7 @@
-import lancedb
+import lancedb # type: ignore
 import os
-import gradio as gr
-from sentence_transformers import SentenceTransformer
+import gradio as gr # type: ignore
+from sentence_transformers import SentenceTransformer, CrossEncoder # type: ignore
 
 
 db = lancedb.connect(".lancedb")
@@ -10,17 +10,25 @@ TABLE = db.open_table(os.getenv("TABLE_NAME"))
 VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
 TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
+RERANKER = os.getenv("RERANKER", "cross-encoder/ms-marco-MiniLM-L-6-v2")
 
 retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
+reranker = CrossEncoder(RERANKER)
 
 
-def retrieve(query, k):
+def retrieve(query, k, rerank_factor=3):
  query_vec = retriever.encode(query)
  try:
- documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
+ documents = (
+ TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN)
+ .limit(k * rerank_factor)
+ .to_list()
+ )
  documents = [doc[TEXT_COLUMN] for doc in documents]
-
- return documents
+ scores = reranker.predict([(query, doc) for doc in documents])
+ best_scores_and_documents = sorted(zip(scores, documents), reverse=True)[:k]
+ best_documents = [doc[1] for doc in best_scores_and_documents]
+ return best_documents
 
  except Exception as e:
  raise gr.Error(str(e))