add cross-encoder

app.py

@@ -34,7 +34,7 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, api_kind):
+def bot(history, api_kind, with_cross_encoder):
     query = history[-1][0]
 
     if not query:
@@ -44,7 +44,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, with_cross_encoder)
 
     document_time = perf_counter() - document_start
     logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
@@ -86,12 +86,13 @@ 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.Checkbox(label="Cross-encoder")
+    cross_encoder = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace")
 
     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_encoder], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

backend/semantic_search.py

@@ -2,6 +2,7 @@ import lancedb
 import os
 import gradio as gr
 from sentence_transformers import SentenceTransformer
+from sentence_transformers import CrossEncoder
 
 
 db = lancedb.connect(".lancedb")
@@ -12,13 +13,21 @@ TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
 
 retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
+cross_encoder = CrossEncoder(os.getenv("RERANK_MODEL"), max_length=512)
 
 
-def retrieve(query, k):
+def retrieve(query, k, with_cross_encoder=False):
     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]
+        if not with_cross_encoder:
+            documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
+            documents = [doc[TEXT_COLUMN] for doc in documents]
+        else:
+            documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k * 2).to_list()
+            scores = cross_encoder.predict([(query, doc[TEXT_COLUMN]) for doc in documents])
+            indexed_arr = [(elem, index) for index, elem in enumerate(scores)]
+            sorted_arr = sorted(indexed_arr, key=lambda x: x[0], reverse=True)
+            documents = [elem for elem, _ in sorted_arr[:k]]
 
         return documents