late commit

app.py

@@ -11,9 +11,13 @@ from jinja2 import Environment, FileSystemLoader
 
 from backend.query_llm import generate_hf, generate_openai
 from backend.semantic_search import retrieve
+from backend.cross_encoder import rerank_with_cross_encoder
+
+
 
 
 TOP_K = int(os.getenv("TOP_K", 4))
+TOP_K_RERANK = int(os.getenv("TOP_K_RERANK", 40))
 
 proj_dir = Path(__file__).parent
 # Setting up the logging
@@ -34,7 +38,7 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, api_kind):
+def bot(history, api_kind, cross_enc):
     query = history[-1][0]
 
     if not query:
@@ -42,12 +46,32 @@ def bot(history, api_kind):
 
     logger.info('Retrieving documents...')
     # Retrieve documents relevant to query
-    document_start = perf_counter()
+    documents = []
+    if not cross_enc:
+        document_start = perf_counter()
+
+        documents = retrieve(query, TOP_K)
+
+        document_time = perf_counter() - document_start
+        logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
+
+    else:
+        document_start = perf_counter()
+
+        documents = retrieve(query, TOP_K_RERANK)
+
+        document_time = perf_counter() - document_start
+        logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
+
+        logger.info('Reranking documents')
+        document_start = perf_counter()
+
+        documents = rerank_with_cross_encoder(cross_enc, documents, query)
+
+        document_time = perf_counter() - document_start
 
-    documents = retrieve(query, TOP_K)
+        logger.info(f'Finished Reranking documents in {round(document_time, 2)} seconds...')
 
-    document_time = perf_counter() - document_start
-    logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
 
     # Create Prompt
     prompt = template.render(documents=documents, query=query)
@@ -86,19 +110,20 @@ with gr.Blocks() as demo:
                 )
         txt_btn = gr.Button(value="Submit text", scale=1)
 
-    api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace")
+    api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace", label="LLM")
+    cross_enc = gr.Radio(choices=["None", "cross-encoder/ms-marco-MiniLM-L-6-v2", "BAAI/bge-reranker-large"], value=None, label="Cross Encoder")
 
     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], [chatbot, prompt_html])
+            bot, [chatbot, api_kind, cross_enc], [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], [chatbot, prompt_html])
+            bot, [chatbot, api_kind, cross_enc], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

backend/cross_encoder.py

@@ -0,0 +1,29 @@
+import os
+
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+cross_encoder = None
+cross_enc_tokenizer = None
+
+TOP_K_RERANK = os.getenv("TOP_K_RERANK", 40)
+
+
+@torch.no_grad()
+def rerank_with_cross_encoder(cross_enc_name, documents, query):
+    if cross_enc_name is None or len(documents) <= 1:
+        return documents
+
+    global cross_encoder, cross_enc_tokenizer
+    if cross_encoder is None or cross_encoder.name_or_path != cross_enc_name:
+        cross_encoder = AutoModelForSequenceClassification.from_pretrained(cross_enc_name)
+        cross_encoder.eval()
+        cross_enc_tokenizer = AutoTokenizer.from_pretrained(cross_enc_name)
+
+    features = cross_enc_tokenizer(
+        [query] * len(documents), documents, padding=True, truncation=True, return_tensors="pt"
+    )
+    scores = cross_encoder(**features).logits.squeeze()
+    ranks = torch.argsort(scores, descending=True)
+    documents = [documents[i] for i in ranks[:TOP_K_RERANK]]
+    return documents