Update app.py

app.py

@@ -15,7 +15,7 @@ from colpali_engine.models.paligemma_colbert_architecture import ColPali
 from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
 from colpali_engine.utils.colpali_processing_utils import process_images, process_queries
 
-def search(query: str, ds, images):
+def search(query: str, ds, images, k):
     qs = []
     with torch.no_grad():
         batch_query = process_queries(processor, [query], mock_image)
@@ -26,8 +26,14 @@ def search(query: str, ds, images):
     # run evaluation
     retriever_evaluator = CustomEvaluator(is_multi_vector=True)
     scores = retriever_evaluator.evaluate(qs, ds)
-    best_page = int(scores.argmax(axis=1).item())
-    return f"The most relevant page is {best_page}", images[best_page]
+
+    top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
+
+    results = []
+    for idx in top_k_indices:
+        results.append((images[idx], f"Page {idx}"))
+
+    return results
 
 
 def index(file, ds):
@@ -83,9 +89,9 @@ with gr.Blocks() as demo:
     search_button = gr.Button("🔍 Search")
     message2 = gr.Textbox("Query not yet set")
     output_img = gr.Image()
+    k = gr.Slider(minimum=1, maximum=10, step=1, label="Number of results", value=5)
 
-    search_button.click(search, inputs=[query, embeds, imgs], outputs=[message2, output_img])
-
+    search_button.click(search, inputs=[query, embeds, imgs, k], outputs=[message2, output_img])
 
 if __name__ == "__main__":
     demo.queue(max_size=10).launch(debug=True)