more options

app.py

@@ -1,35 +1,34 @@
 import gradio as gr
 from pyserini.search.lucene import LuceneSearcher
 import os
+import json
 
-# Download the index if it doesn't exist
-if not os.path.exists('msmarco-passage'):
-    os.system('python -c "from pyserini.search import LuceneSearcher; LuceneSearcher.from_prebuilt_index(\'msmarco-passage\')"')
+def initialize_searcher(index_name):
+    if not os.path.exists(index_name):
+        os.system(f'python -c "from pyserini.search import LuceneSearcher; LuceneSearcher.from_prebuilt_index(\'{index_name}\')"')
+    searcher = LuceneSearcher.from_prebuilt_index(index_name)
+    searcher.set_bm25(k1=0.9, b=0.4)
+    return searcher
 
-searcher = LuceneSearcher.from_prebuilt_index('msmarco-passage')
-searcher.set_bm25(k1=0.9, b=0.4)
-
-def search_pyserini(query):
+def search_pyserini(query, top_k, index_name):
     try:
-        hits = searcher.search(query, k=10)
+        searcher = initialize_searcher(index_name)
+        hits = searcher.search(query, k=top_k)
         results = []
         for i, hit in enumerate(hits):
             doc = searcher.doc(hit.docid)
-            content = doc.raw()
+            doc_dict = json.loads(doc.raw())
             results.append({
                 "rank": i + 1,
                 "doc_id": hit.docid,
                 "score": hit.score,
-                "content": content
+                "content": doc_dict['contents']
             })
-        return results
+        return format_results(results)
     except Exception as e:
-        return [{"error": str(e)}]
+        return f"<div class='error'>An error occurred: {str(e)}</div>"
 
 def format_results(results):
-    if isinstance(results, list) and len(results) > 0 and "error" in results[0]:
-        return f"<div class='error'>An error occurred: {results[0]['error']}</div>"
-    
     html = "<div class='results-container'>"
     for result in results:
         html += f"""
@@ -48,14 +47,14 @@ css = """
 }
 .results-container {
     display: flex;
-    flex-wrap: wrap;
+    flex-direction: column;
     gap: 20px;
 }
 .result-item {
     border: 1px solid #ddd;
     border-radius: 8px;
     padding: 15px;
-    width: calc(50% - 10px);
+    width: 100%;
     box-shadow: 0 2px 4px rgba(0,0,0,0.1);
 }
 .result-item h3 {
@@ -81,6 +80,23 @@ with gr.Blocks(css=css) as iface:
     gr.Markdown("# Pyserini Search Interface")
     gr.Markdown("Enter a query to search using Pyserini with BM25 scoring (k1=0.9, b=0.4).")
     
+    with gr.Row():
+        index_input = gr.Textbox(
+            value="msmarco-passage",
+            lines=1,
+            label="Prebuilt Index Name",
+            placeholder="Enter the name of the prebuilt index"
+        )
+    
+    with gr.Row():
+        top_k_slider = gr.Slider(
+            minimum=1,
+            maximum=100,
+            value=10,
+            step=1,
+            label="Number of top results to return"
+        )
+    
     with gr.Row():
         query_input = gr.Textbox(
             lines=1, 
@@ -96,31 +112,8 @@ with gr.Blocks(css=css) as iface:
     
     search_button.click(
         fn=search_pyserini,
-        inputs=query_input,
-        outputs=output,
-        _js="(results) => format_results(results)"  # Client-side formatting
+        inputs=[query_input, top_k_slider, index_input],
+        outputs=output
     )
 
-# Add the JavaScript function to format results
-iface.load(js="""
-function format_results(results) {
-    if (Array.isArray(results) && results.length > 0 && results[0].hasOwnProperty('error')) {
-        return `<div class='error'>An error occurred: ${results[0].error}</div>`;
-    }
-    
-    let html = "<div class='results-container'>";
-    for (let result of results) {
-        html += `
-        <div class='result-item'>
-            <h3>Rank ${result.rank} (Score: ${result.score.toFixed(4)})</h3>
-            <p class='doc-id'>Doc ID: ${result.doc_id}</p>
-            <p class='content'>${result.content}</p>
-        </div>
-        `;
-    }
-    html += "</div>";
-    return html;
-}
-""")
-
 iface.launch()