Refactor retrieve function and enhance UI with Gradio components for improved query handling

app.py

@@ -367,26 +367,34 @@ return_type = List[Hit]
 
 
 ## YOUR_CODE_STARTS_HERE
-def retrieve(query: str, topk: int = 10) -> return_type:
-    ranking = bm25_retriever.retrieve(query=query, topk=3)
+def retrieve(query: str, topk: int=10) -> return_type:
+    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})
     return hits
 
-
-demo = gr.Interface(
-    fn=retrieve,
-    inputs=gr.Textbox(lines=3, placeholder="Enter your query here..."),
-    outputs="json",
-    title="CSC BM25 Retriever",
-    description="Retrieve documents based on the query using CSC BM25 Retriever",
-    examples=[
-        ["What are the differences between immunodeficiency and autoimmune diseases?"],
-        ["What are the causes of immunodeficiency?"],
-        ["What are the symptoms of immunodeficiency?"],
-    ],
-)
-## YOUR_CODE_ENDS_HERE
+with gr.Blocks(theme=gr.themes.Ocean()) as demo:
+    gr.Markdown("# BM25 Retriever")
+    gr.Markdown("Retrieve documents based on the query using BM25 Retriever")
+    query = gr.Textbox(lines=3, placeholder="Enter your query here...")
+    topk = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Top-K")
+    # Search Button
+    examples = gr.Examples(
+        examples=[
+            ["What are the differences between immunodeficiency and autoimmune diseases?"],
+            ["What are the causes of immunodeficiency?"],
+            ["What are the symptoms of immunodeficiency?"],
+        ],
+        inputs=[query],
+    )
+    search_button = gr.Button("Search", elem_id="search_button")
+    results_section = gr.JSON(elem_id="results_section")
+    search_button.click(
+        retrieve,
+        inputs=[query, topk],
+        outputs=results_section,
+    )
 demo.launch()
+## YOUR_CODE_ENDS_HERE