Reranker

app.py

@@ -6,12 +6,17 @@ import logging
 from pathlib import Path
 from time import perf_counter
 
+from dotenv import load_dotenv
+print(load_dotenv())
+
+
 import gradio as gr
 from jinja2 import Environment, FileSystemLoader
 
 from backend.query_llm import generate_hf, generate_openai
 from backend.semantic_search import retrieve
 
+# load_dotenv(os.path.join(os.path.dirname(__file__), '..', '.env'))
 
 TOP_K = int(os.getenv("TOP_K", 4))
 
@@ -34,7 +39,7 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, api_kind):
+def bot(history, api_kind, rerank):
     query = history[-1][0]
 
     if not query:
@@ -44,7 +49,7 @@ def bot(history, api_kind):
     # Retrieve documents relevant to query
     document_start = perf_counter()
 
-    documents = retrieve(query, TOP_K)
+    documents = retrieve(query, TOP_K, rerank)
 
     document_time = perf_counter() - document_start
     logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
@@ -60,6 +65,8 @@ def bot(history, api_kind):
     else:
          raise gr.Error(f"API {api_kind} is not supported")
 
+    print(f"{prompt}")
+
     history[-1][1] = ""
     for character in generate_fn(prompt, history[:-1]):
         history[-1][1] = character
@@ -87,18 +94,19 @@ with gr.Blocks() as demo:
         txt_btn = gr.Button(value="Submit text", scale=1)
 
     api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace")
+    rerank = gr.Checkbox(label="Rerank", value=True)
 
     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, rerank], [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, rerank], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

backend/semantic_search.py

@@ -1,8 +1,10 @@
 import os
+import torch
 
 import gradio as gr
 import lancedb
 from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
 db = lancedb.connect(".lancedb")
 
@@ -12,15 +14,38 @@ TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
 
 retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
+reranker_model = os.getenv("RERANKER_MODEL", None)
 
+if reranker_model:
+    reranker = AutoModelForSequenceClassification.from_pretrained(reranker_model)
+    tokenizer = AutoTokenizer.from_pretrained(reranker_model)
+    reranker_pipeline = pipeline("text-classification", model=reranker, tokenizer=tokenizer)
 
-def retrieve(query, k):
+
+def retrieve(query, k, rerank=True):
     query_vec = retriever.encode(query)
     try:
-        documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
-        documents = [doc[TEXT_COLUMN] for doc in documents]
+        num_retrieve = k * (5 if rerank else 1)
+        documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(num_retrieve).to_list()
+        docs = [doc[TEXT_COLUMN] for doc in documents]
+
+        if not rerank:
+            return docs
+
+        assert reranker_model, "Reranker model is not provided"
 
-        return documents
+        reranked_documents = []
+        for i in range(0, len(docs), BATCH_SIZE):
+            batch_texts = docs[i:i+BATCH_SIZE]
+            inputs = tokenizer([query]*len(batch_texts), batch_texts, return_tensors="pt", padding=True, truncation=True)
+            with torch.no_grad():
+                outputs = reranker(**inputs)
+                logits = outputs.logits.squeeze().tolist()
+                reranked_documents.extend(zip(batch_texts, logits))
+
+        reranked_documents.sort(key=lambda x: x[1], reverse=True)
+        return [doc[0] for doc in reranked_documents[:k]]
 
     except Exception as e:
         raise gr.Error(str(e))
+