Change prompt; Change Rerank model

backend/semantic_search.py

@@ -2,7 +2,8 @@ import lancedb
 import os
 import gradio as gr
 from sentence_transformers import SentenceTransformer
-from FlagEmbedding import FlagReranker
+from sentence_transformers import CrossEncoder
+# from FlagEmbedding import FlagReranker
 
 
 db = lancedb.connect(".lancedb")
@@ -13,7 +14,8 @@ TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
 
 retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
-reranker = FlagReranker(os.getenv("RERANKER_MODEL", 'BAAI/bge-reranker-large'), use_fp16=True)
+# reranker = FlagReranker(os.getenv("RERANKER_MODEL", 'BAAI/bge-reranker-large'), use_fp16=True)
+reranker = CrossEncoder(os.getenv("RERANKER_MODEL", 'cross-encoder/ms-marco-MiniLM-L-6-v2'), max_length=512)
 
 def retrieve(query, k):
     query_vec = retriever.encode(query)
@@ -29,7 +31,7 @@ def retrieve(query, k):
 def rerank(documents, query, k):
     try:
         query_pairs = [[query, doc] for doc in documents]
-        scores = reranker.compute_score(query_pairs)
+        scores = reranker.predict(query_pairs)
         scored_documents = list(zip(documents, scores))
         scored_documents.sort(key=lambda x: x[1], reverse=True)
         top_k_documents = [doc for doc, _ in scored_documents[:k]]

templates/template.j2

@@ -1,4 +1,4 @@
-Instructions: Use the following unique documents in the Context section to answer the Query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+Instructions: Use the following unique sorted by relevance documents in the Context section to answer the Query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
 Context:
 {% for doc in documents %}
 ---