update

app.py

@@ -1,21 +1,27 @@
 import csv
+from typing import Any
 
 import gradio as gr
 import pandas as pd
 from sentence_transformers import SentenceTransformer, util
 from underthesea import word_tokenize
+from retriever_trainer import PretrainedColBERT
+
 
 bi_encoder = SentenceTransformer("phamson02/cotmae_biencoder2_170000_sbert")
-# cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+colbert = PretrainedColBERT(
+    pretrained_model_name="phamson02/colbert2.1_290000",
+)
 corpus_embeddings = pd.read_pickle("data/passage_embeds.pkl")
 
 with open("data/child_passages.tsv", "r") as f:
     tsv_reader = csv.reader(f, delimiter="\t")
     child_passage_ids = [row[0] for row in tsv_reader]
+    child_passages = [row[1] for row in tsv_reader]
 
 with open("data/parent_passages.tsv", "r") as f:
     tsv_reader = csv.reader(f, delimiter="\t")
-    parent_passages = {row[0]: row[1] for row in tsv_reader}
+    parent_passages_map = {row[0]: row[1] for row in tsv_reader}
 
 
 def f7(seq):
@@ -24,7 +30,7 @@ def f7(seq):
     return [x for x in seq if not (x in seen or seen_add(x))]
 
 
-def search(query: str, top_k: int = 100, reranking: bool = False):
+def search(query: str, reranking: bool = False, top_k: int = 100):
     query = word_tokenize(query, format="text")
 
     print("Top 5 Answer by the NSE:")
@@ -36,19 +42,22 @@ def search(query: str, top_k: int = 100, reranking: bool = False):
     hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
     hits = hits[0]  # Get the hits for the first query
 
+    top_k_child_passages = [child_passages[hit["corpus_id"]] for hit in hits]
+    top_k_child_passage_ids = [child_passage_ids[hit["corpus_id"]] for hit in hits]
+
     ##### Re-Ranking #####
     # Now, score all retrieved passages with the cross_encoder
     if reranking:
-        cross_inp = [[query, corpus[hit["corpus_id"]]] for hit in hits]
-        cross_scores = cross_encoder.predict(cross_inp)
+        colbert_scores: list[dict[str, Any]] = colbert.rerank(
+            query=query, documents=top_k_child_passages, top_k=100
+        )
 
-        # Sort results by the cross-encoder scores
-        for idx in range(len(cross_scores)):
-            hits[idx]["cross-score"] = cross_scores[idx]
-
-        hits = sorted(hits, key=lambda x: x["cross-score"], reverse=True)
+        # Reorder child passage ids based on the reranking
+        top_k_child_passage_ids = [
+            top_k_child_passage_ids[score["corpus_id"]] for score in colbert_scores
+        ]
 
-    top_20_hits = hits[0:20]
+    top_20_hits = top_k_child_passage_ids[0:20]
     hit_child_passage_ids = [child_passage_ids[hit["corpus_id"]] for hit in top_20_hits]
     hit_parent_passage_ids = f7(
         [
@@ -60,7 +69,7 @@ def search(query: str, top_k: int = 100, reranking: bool = False):
     assert len(hit_parent_passage_ids) >= 5, "Not enough unique parent passages found"
 
     for hit in hit_parent_passage_ids[:5]:
-        ans.append(parent_passages[hit])
+        ans.append(parent_passages_map[hit])
 
     return ans[0], ans[1], ans[2], ans[3], ans[4]
 
@@ -76,6 +85,8 @@ exp = [
 desc = "This is a semantic search engine powered by SentenceTransformers (Nils_Reimers) with a retrieval and reranking system on Wikipedia corous. This will return the top 5 results. So Quest on with Transformers."
 
 inp = gr.Textbox(lines=1, placeholder=None, label="search you query here")
+reranking_checkbox = gr.Checkbox(label="Enable reranking")
+
 out1 = gr.Textbox(type="text", label="Search result 1")
 out2 = gr.Textbox(type="text", label="Search result 2")
 out3 = gr.Textbox(type="text", label="Search result 3")
@@ -84,7 +95,7 @@ out5 = gr.Textbox(type="text", label="Search result 5")
 
 iface = gr.Interface(
     fn=search,
-    inputs=inp,
+    inputs=[inp, reranking_checkbox],
     outputs=[out1, out2, out3, out4, out5],
     examples=exp,
     article=desc,

requirements.txt

@@ -2,4 +2,5 @@ sentence-transformers
 torch
 pandas
 gradio
-underthesea
+underthesea
+retriever-trainer[colbert] @ git+https://[email protected]/phamson02/retriever_trainer.git@rerank