Update app.py

app.py

@@ -1,24 +1,26 @@
+import logging
 import gradio as gr
-import os
-import query_index
 import datasets
 import sentence_transformers
 
-def query(text, k=5):
-    model = sentence_transformers.SentenceTransformer(
-    "dangvantuan/sentence-camembert-large", device="cpu")
+logging.disable(logging.CRITICAL)
 
-    dataset = datasets.load_dataset("json", data_files=["./data/dataset.json"], split="train")
-    dataset.load_faiss_index("embeddings", "index.faiss")
+model = sentence_transformers.SentenceTransformer(
+    "dangvantuan/sentence-camembert-large", device="cuda"
+)
+
+dataset = datasets.load_dataset("json", data_files=["./data/dataset.json"], split="train")
+dataset.load_faiss_index("embeddings", "index.faiss")
 
-    query_embedding = model.encode(text)
+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"],
@@ -28,16 +30,25 @@ def query(text, k=5):
     ):
         start = start
         end = end
-        print(f"title: {title}")
-        print(f"transcript: [{str(start)+' ====> '+str(end)}] {text}")
-        print(f"link: {url}")
-        print("*" * 10)
+        result = {
+            "title": title,
+            "transcript": f"[{str(start)+' ====> '+str(end)}] {text}",
+            "link": url
+        }
+        results.append(result)
+    return results
 
 iface = gr.Interface(
-  fn=query, 
-  inputs='text',
-  outputs='text',
-  examples=[["Qu'est ce qui t'a fait le plus progresser?"]]
+    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
+    ]
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()