Update query_index.py

query_index.py

@@ -1,18 +1,10 @@
-import argparse
 import logging
-
+import gradio as gr
 import datasets
 import sentence_transformers
 
-import utils
-
 logging.disable(logging.CRITICAL)
 
-parser = argparse.ArgumentParser()
-parser.add_argument("--query", type=str, required=True)
-parser.add_argument("--k", type=int, default=5)
-args = parser.parse_args()
-
 model = sentence_transformers.SentenceTransformer(
     "dangvantuan/sentence-camembert-large", device="cuda"
 )
@@ -20,24 +12,43 @@ model = sentence_transformers.SentenceTransformer(
 dataset = datasets.load_dataset("json", data_files=["./data/dataset.json"], split="train")
 dataset.load_faiss_index("embeddings", "index.faiss")
 
-query_embedding = model.encode(args.query)
-_, retrieved_examples = dataset.get_nearest_examples(
-    "embeddings",
-    query_embedding,
-    k=args.k,
+def search(query: str, k: int):
+    query_embedding = model.encode(query)
+    _, retrieved_examples = dataset.get_nearest_examples(
+        "embeddings",
+        query_embedding,
+        k=k,
+    )
+
+    results = []
+    for text, start, end, title, url in zip(
+        retrieved_examples["text"],
+        retrieved_examples["start"],
+        retrieved_examples["end"],
+        retrieved_examples["title"],
+        retrieved_examples["url"],
+    ):
+        start = start
+        end = end
+        result = {
+            "title": title,
+            "transcript": f"[{str(start)+' ====> '+str(end)}] {text}",
+            "link": url
+        }
+        results.append(result)
+    return results
+
+iface = gr.Interface(
+    fn=search,
+    inputs=["text", "number"],
+    outputs=gr.outputs.JSON(),
+    title="Search Dataset",
+    description="Search a dataset using Camembert and Faiss.",
+    example=[
+        "Enter a query to search for.",
+        5
+    ]
 )
 
-
-for text, start, end, title, url in zip(
-    retrieved_examples["text"],
-    retrieved_examples["start"],
-    retrieved_examples["end"],
-    retrieved_examples["title"],
-    retrieved_examples["url"],
-):
-    start = start
-    end = end
-    print(f"title: {title}")
-    print(f"transcript: [{str(start)+' ====> '+str(end)}] {text}")
-    print(f"link: {url}")
-    print("*" * 10)
+if __name__ == "__main__":
+    iface.launch()