wip

app.py

@@ -1,29 +1,59 @@
 import gradio as gr
 import ir_datasets
 import pandas as pd
+import numpy as np
 
 from autogluon.multimodal import MultiModalPredictor
 
 
-def text_embedding(query: str):
+query_embedding = None
+document_embedding = None
+docs_df = None
+
+def text_embedding_batch():
+    model_name = "sentence-transformers/all-MiniLM-L6-v2"
+    dataset = ir_datasets.load("beir/fiqa/dev")
+    docs_df = pd.DataFrame(dataset.docs_iter()).set_index("doc_id").sample(frac=0.0001)
+    predictor = MultiModalPredictor(
+        pipeline="feature_extraction",
+        hyperparameters={
+            "model.hf_text.checkpoint_name": model_name
+        }
+    )
+    embedding = predictor.extract_embedding(docs_df)
+    query_embedding = embedding["text"]
+    return query_embedding
+
+
+def text_embedding_single(query: str):
     model_name = "sentence-transformers/all-MiniLM-L6-v2"
-    # dataset = ir_datasets.load("beir/fiqa/dev")
-    # docs_df = pd.DataFrame(dataset.docs_iter()).set_index("doc_id").sample(frac=0.001)
     predictor = MultiModalPredictor(
         pipeline="feature_extraction",
         hyperparameters={
             "model.hf_text.checkpoint_name": model_name
         }
     )
-    # query_embedding = predictor.extract_embedding(docs_df)
-    # return query_embedding["text"]   
-    query_embedding = predictor.extract_embedding([query])
-    return query_embedding["0"]
+    embedding = predictor.extract_embedding([query])
+    document_embedding = embedding["0"]
+    return document_embedding
+
+
+def rank_document():
+    q_norm = query_embedding / np.linalg.norm(query_embedding, axis=-1, keepdims=True)
+    print(q_norm)
+    d_norm = document_embedding / np.linalg.norm(document_embedding, axis=-1, keepdims=True)
+    scores = d_norm.dot(q_norm[0])
+    print(scores)
 
+    result = []
+    for idx in np.argsort(-scores)[:2]:
+        result.append(docs_df['text'].iloc[idx])
+    return result
+    
 
 def main():
     with gr.Blocks(title="OpenSearch Demo") as demo:
-        gr.Markdown("# Text Embedding for Search Queries")
+        gr.Markdown("# Semantic Search with Autogluon")
         gr.Markdown("Ask an open question!")
         with gr.Row():
             inp_single = gr.Textbox(show_label=False)
@@ -31,23 +61,25 @@ def main():
             btn_single = gr.Button("Generate Embedding")
         with gr.Row():
             out_single = gr.DataFrame(label="Embedding", show_label=True)
-        gr.Markdown("You can select one of the sample datasets for batch inference")
-        with gr.Row():
-            with gr.Column():
-                btn_fiqa = gr.Button("fiqa")
-            with gr.Column():
-                btn_faiss = gr.Button("faiss")
+        gr.Markdown("You can select one of the sample datasets for document embedding")
         with gr.Row():
-            out_batch = gr.DataFrame(label="Embedding", show_label=True)
-        gr.Markdown("You can also try out our batch inference by uploading a file")
+            btn_fiqa = gr.Button("fiqa")
         with gr.Row():
-            out_batch = gr.File(interactive=True)
+            out_batch = gr.DataFrame(label="Sample Embeddings", show_label=True, row_count=5)
+        gr.Markdown("Now rank the documents and pick the top 3 most relevant from the dataset")
         with gr.Row():    
-            btn_file = gr.Button("Generate Embedding")
-            
-        btn_single.click(fn=text_embedding, inputs=inp_single, outputs=out_single)
-        btn_file.click(fn=text_embedding, inputs=inp_single, outputs=out_single)    
-    
+            btn_rank = gr.Button("Rank documents")
+        with gr.Row():
+            out_rank = gr.DataFrame(label="Top ranked documents", show_label=True, row_count=5)
+        # with gr.Row():
+        #     out_batch = gr.File(interactive=True)
+        # with gr.Row():    
+        #     btn_file = gr.Button("Generate Embedding")
+        
+        btn_single.click(fn=text_embedding_single, inputs=inp_single, outputs=out_single)
+        btn_fiqa.click(fn=text_embedding_batch, inputs=None, outputs=out_batch)
+        btn_rank.click(fn=rank_document, inputs=None, outputs=out_rank)
+        # btn_file.click(fn=text_embedding_batch, inputs=inp_single, outputs=out_single)    
     demo.launch()