Enhance BM25 index creation and retrieval functionality; save index to output directory

app.py

@@ -357,22 +357,29 @@ bm25_index = BM25Index.build_from_documents(
 
 
 class Hit(TypedDict):
-    cid: str
-    score: float
-    text: str
-
+  cid: str
+  score: float
+  text: str
 
 demo: Optional[gr.Interface] = None  # Assign your gradio demo to this variable
 return_type = List[Hit]
-
+bm25_index = BM25Index.build_from_documents(
+    documents=iter(sciq.corpus),
+    ndocs=12160,
+    show_progress_bar=True,
+    k1=0.9,
+    b=0.4,
+)
+bm25_index.save("output/bm25_index")
+bm25_retriever = BM25Retriever(index_dir="output/bm25_index")
 
 ## YOUR_CODE_STARTS_HERE
 def retrieve(query: str, topk: int = 10) -> return_type:
-    ranking = bm25_retriever.retrieve(query=query, topk=3)
+    ranking = bm25_retriever.retrieve(query=query, topk=topk)
     hits = []
     for cid, score in ranking.items():
         text = bm25_retriever.index.doc_texts[bm25_retriever.index.cid2docid[cid]]
-        hits.append({"cid": cid, "score": score, "text": text})
+        hits.append(Hit(cid=cid, score=score, text=text))
     return hits
 
 
@@ -388,5 +395,5 @@ demo = gr.Interface(
         ["What are the symptoms of immunodeficiency?"],
     ],
 )
-## YOUR_CODE_ENDS_HERE
 demo.launch()
+## YOUR_CODE_ENDS_HERE