wip

app.py

@@ -1,18 +1,24 @@
 import gradio as gr
+import ir_datasets
+import pandas as pd
+
 from autogluon.multimodal import MultiModalPredictor
 
 
 def text_embedding(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"]   
+    return query_embedding["0"]
 
 
 def main():
@@ -20,13 +26,27 @@ def main():
         gr.Markdown("# Text Embedding for Search Queries")
         gr.Markdown("Ask an open question!")
         with gr.Row():
-            inp = gr.Textbox(show_label=False)
+            inp_single = gr.Textbox(show_label=False)
         with gr.Row():    
-            btn = gr.Button("Generate Embedding")
+            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():
-            out = gr.DataFrame(label="Embedding", show_label=True)
-        
-        btn.click(fn=text_embedding, inputs=inp, outputs=out)  
+            with gr.Column():
+                btn_fiqa = gr.Button("fiqa")
+            with gr.Column():
+                btn_faiss = gr.Button("faiss")
+        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")
+        with gr.Row():
+            out_batch = gr.File(interactive=True)
+        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)    
     
     demo.launch()   
 

automm_semantic_embedding.py

@@ -0,0 +1,43 @@
+import ir_datasets
+import pandas as pd
+from autogluon.multimodal import MultiModalPredictor
+
+
+dataset = ir_datasets.load("beir/fiqa/dev")
+
+
+
+
+dataset = ir_datasets.load("beir/fiqa/dev")
+
+docs_df = pd.DataFrame(dataset.docs_iter()).set_index("doc_id").sample(frac=0.0001)
+query_df = pd.DataFrame(dataset.queries_iter()).set_index("query_id")
+
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+
+predictor = MultiModalPredictor(
+    pipeline="feature_extraction",
+    hyperparameters={
+        "model.hf_text.checkpoint_name": model_name
+    }
+)
+
+document_embedding = predictor.extract_embedding(docs_df)
+
+query = "What happened when the dot com bubble burst?"
+query_embedding = predictor.extract_embedding([query])
+
+import numpy as np
+
+q_norm = query_embedding['0'] / np.linalg.norm(query_embedding['0'], axis=-1, keepdims=True)
+d_norm = document_embedding['text'] / np.linalg.norm(document_embedding['text'], axis=-1, keepdims=True)
+scores = d_norm.dot(q_norm[0])
+
+
+print(f'Question: {query}')
+print()
+for idx in np.argsort(-scores)[:2]:
+    print(f'Top {idx} result:')
+    print('-----------------')
+    print(docs_df['text'].iloc[idx])
+    print()

requirements.txt

@@ -1,4 +1,5 @@
 gradio
 wheel
 setuptools
+ir_datasets
 git+https://github.com/awslabs/autogluon.git@master#subdirectory=autogluon