add reranker

app.py

@@ -11,6 +11,7 @@ from jinja2 import Environment, FileSystemLoader
 
 from backend.query_llm import generate_hf, generate_openai
 from backend.semantic_search import retrieve
+from backend.reranker import rerank_documents
 
 
 TOP_K = int(os.getenv("TOP_K", 4))
@@ -34,7 +35,7 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, api_kind, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param):
+def bot(history, api_kind, chunk_table, embedding_model, llm_model, cross_encoder, rerank_topk ):
     top_k_param = int(top_k_param)
     query = history[-1][0]
 
@@ -47,6 +48,11 @@ def bot(history, api_kind, chunk_table, embedding_model, llm_model, cross_encode
 
     #documents = retrieve(query, TOP_K)
     documents = retrieve(query, top_k_param, chunk_table, embedding_model)
+    if cross_encoder != "None" and len(documents) > 1:
+        documents = rerank_documents(query, documents, query, top_k_rerank=rerank_topk)
+        #"cross-encoder/ms-marco-MiniLM-L-6-v2"
+
+
 
 
     document_time = perf_counter() - document_start
@@ -121,7 +127,7 @@ with gr.Blocks() as demo:
             )    
     cross_encoder = gr.Radio(
                 choices=[
-                    "None"
+                    "None",
                     "BAAI/bge-reranker-large",
                     "cross-encoder/ms-marco-MiniLM-L-6-v2",
                 ],
@@ -137,20 +143,30 @@ with gr.Blocks() as demo:
                 ],
                 value="5",
                 label='top-K'
-            ) 
+            )
+    rerank_topk = gr.Radio(
+                choices=[
+                    "5",
+                    "10",
+                    "20",
+                    "50",
+                ],
+                value="5",
+                label='rerank-top-K'
+            )  
 
 
     prompt_html = gr.HTML()
     # Turn off interactivity while generating if you click
     txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [chatbot, api_kind, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param], [chatbot, prompt_html])
+            bot, [chatbot, api_kind, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param, rerank_topk], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
 
     # Turn off interactivity while generating if you hit enter
     txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [chatbot, api_kind, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param], [chatbot, prompt_html])
+            bot, [chatbot, api_kind, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param, rerank_topk], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

backend/reranker.py

@@ -0,0 +1,13 @@
+from sentence_transformers import CrossEncoder
+
+
+def rerank_documents(ce_model_name, documents, query, top_k_rerank):
+    top_k_rerank = int(top_k_rerank)
+    pairs = []
+    for doc in documents:
+        pairs.append((query, doc))
+        ce_model = CrossEncoder(ce_model_name, max_length=512)
+        scores = ce_model.predict(pairs)
+        #sorted_pairs = [(s, x[1]) for s, x in sorted(zip(scores, pairs), key=lambda p: p[0], reverse = True)]
+        reranked_docs = [x[1] for _, x in sorted(zip(scores, pairs), key=lambda p: p[0], reverse = True)]
+        return reranked_docs[:top_k_rerank]