Update app.py

app.py

@@ -1,25 +1,24 @@
+import argparse
 import logging
-import gradio as gr
+
 import datasets
-import sentence_transformers
+import gradio as gr
 
 logging.disable(logging.CRITICAL)
 
-model = sentence_transformers.SentenceTransformer(
-    "dangvantuan/sentence-camembert-large", device="cuda"
-)
+model_name = "dangvantuan/sentence-camembert-large"
+model = gr.load(model_name)
 
-dataset = datasets.load_dataset("json", data_files=["./dataset.json"], split="train")
+dataset = datasets.load_dataset("json", data_files=["./data/dataset.json"], split="train")
 dataset.load_faiss_index("embeddings", "index.faiss")
 
-def search(query: str, k: int):
+def search(query, k):
     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"],
@@ -32,23 +31,27 @@ def search(query: str, k: int):
         end = end
         result = {
             "title": title,
-            "transcript": f"[{str(start)+' ====> '+str(end)}] {text}",
-            "link": url
+            "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
-    ]
+    search,
+    inputs=[
+        gr.inputs.Textbox(label="Query"),
+        gr.inputs.Number(label="K", default=3, min_value=1, max_value=10),
+    ],
+    outputs=[
+        gr.outputs.Textbox(label="Title"),
+        gr.outputs.Textbox(label="Transcript"),
+        gr.outputs.Textbox(label="Link"),
+    ],
+    title="Camembert and Faiss-powered Search Engine",
+    description="Search through a dataset using Camembert and Faiss",
+    theme="light",
+    layout="vertical",
 )
 
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()