init

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/child_passages.tsv filter=lfs diff=lfs merge=lfs -text
+data/parent_passages.tsv filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,90 @@
+import csv
+
+import gradio as gr
+import pandas as pd
+from sentence_transformers import SentenceTransformer, util
+
+bi_encoder = SentenceTransformer("phamson02/cotmae_biencoder2_170000_sbert")
+# cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+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]
+
+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}
+
+
+def f7(seq):
+    seen = set()
+    seen_add = seen.add
+    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):
+    print("Top 5 Answer by the NSE:")
+    print()
+    ans: list[str] = []
+    ##### Sematic Search #####
+    # Encode the query using the bi-encoder and find potentially relevant passages
+    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
+    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
+    hits = hits[0]  # Get the hits for the first query
+
+    ##### 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)
+
+        # 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)
+
+    top_20_hits = hits[0:20]
+    hit_child_passage_ids = [child_passage_ids[hit["corpus_id"]] for hit in top_20_hits]
+    hit_parent_passage_ids = f7(
+        [
+            "_".join(hit_child_passage_id.split("_")[:-1])
+            for hit_child_passage_id in hit_child_passage_ids
+        ]
+    )
+
+    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])
+
+    return ans[0], ans[1], ans[2], ans[3], ans[4]
+
+
+exp = [
+    "Who is steve jobs?",
+    "What is coldplay?",
+    "What is a turing test?",
+    "What is the most interesting thing about our universe?",
+    "What are the most beautiful places on earth?",
+]
+
+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")
+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")
+out4 = gr.Textbox(type="text", label="Search result 4")
+out5 = gr.Textbox(type="text", label="Search result 5")
+
+iface = gr.Interface(
+    fn=search,
+    inputs=inp,
+    outputs=[out1, out2, out3, out4, out5],
+    examples=exp,
+    article=desc,
+    title="Neural Search Engine",
+)
+iface.launch()

data/child_passages.tsv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ab5beca4e38074457dd397a7b56f30679db2a92716bd31fa11a4c303db3dea3
+size 428522185

data/parent_passages.tsv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8e35cd1a742779cbfb2ff5b7809ff92663a90ee92bd72c2d8b1d66f375e6535
+size 352871807

data/passage_embeds.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3f4f423e0c7ef2021f17b8ca0dd78b1a32c5c75d8eaca06238ac12e8661ea0e
+size 1060860322

requirements.txt

@@ -0,0 +1,4 @@
+sentence-transformers
+torch
+pandas
+gradio