Update app.py

app.py

@@ -165,7 +165,7 @@ def search(query: str, k: int = 5):
           type: object
           properties:
             query: {type: string, description: "User query in natural language."}
-            k: {type: integer, minimum: 1, maximum: 20, default: 5. description: "Number of top pages to retrieve."}
+            k: {type: integer, minimum: 1, maximum: 10, default: 5. description: "Number of top pages to retrieve."}
           required: ["query"]
 
     Args:
@@ -196,11 +196,21 @@ def search(query: str, k: int = 5):
     scores = processor.score(qs, ds, device=device)
     top_k_indices = scores[0].topk(k).indices.tolist()
 
-    print(top_k_indices)
+    # Base set & neighbor expansion
+    base = set(top_k_indices)
+    expanded = set(base)
+    for i in base:
+        expanded.add(i - 1)
+        expanded.add(i + 1)
+
+    expanded = {i for i in expanded if i >= 0 and i<=len(images)}
+
+    expanded_indices = sorted(expanded)
+    print(top_k_indices, expanded_indices)
 
     # Build gallery results with 1-based page numbering
     results = []
-    for idx in top_k_indices:
+    for idx in expanded_indices:
         page_num = idx + 1
         results.append((images[idx], f"Page {page_num}"))
 
@@ -264,7 +274,7 @@ ColPali is implemented from the [ColPali paper](https://arxiv.org/abs/2407.01449
         with gr.Column(scale=3):
             gr.Markdown("## 2️⃣ Search")
             query = gr.Textbox(placeholder="Enter your query here", label="Query")
-            k_slider = gr.Slider(minimum=1, maximum=20, step=1, label="Number of results", value=5)
+            k_slider = gr.Slider(minimum=1, maximum=10, step=1, label="Number of results", value=5)
             search_button = gr.Button("🔍 Search", variant="primary")
             output_text = gr.Textbox(label="AI Response", placeholder="Generated response based on retrieved documents")