Update app.py

app.py

@@ -583,24 +583,28 @@ return_type = List[Hit]
 
 ## YOUR_CODE_STARTS_HERE
 
-def search(query) -> List[Hit]:
-    return_type: List[Hit] = []
-    bm_25_retriever = BM25Retriever(index_dir="output/bm25_index")
-    ranking = bm_25_retriever.retrieve(query)
-    for rank in ranking:
-      hit = {
-          "cid": rank,
-          "score": ranking[rank],
-          "text": bm_25_retriever.index.doc_texts[bm_25_retriever.index.cid2docid[rank]]
-      }
-      return_type.append(hit)
-
-    return return_type
+def search(query: str) -> List[Hit]:
+    bm25_index = BM25Index.build_from_documents(
+        documents=iter(sciq.corpus),
+        ndocs=12160,
+        show_progress_bar=True
+    )
+    bm25_index.save("output/bm25_index")
+    bm25_retriever = BM25Retriever(index_dir="output/bm25_index")
+    ranking = bm25_retriever.retrieve(query=query)
+    hits = []
+    for cid, score in ranking.items():
+        doc = next((doc for doc in sciq.corpus if doc.collection_id == cid), None)
+        if doc:
+            hits.append({"cid": cid, "score": score, "text": doc.text})
+    return hits
 
 demo = gr.Interface(
     fn=search,
-    inputs=["text"],
-    outputs=gr.Textbox()
+    inputs=gr.Textbox(lines=2, placeholder="Enter your query here..."),
+    outputs=gr.JSON(label="Search Results"),
+    title="SciQ Search Engine",
+    description="Enter a query to search the SciQ dataset using BM25.",
 )
 
 ## YOUR_CODE_ENDS_HERE