Update colpali-main/demo/app.py

colpali-main/demo/app.py

@@ -1,5 +1,4 @@
 import os
-import sys
 
 import gradio as gr
 import torch
@@ -9,13 +8,12 @@ from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import AutoProcessor
 
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
 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, k):
+
+def search(query: str, ds, images):
     qs = []
     with torch.no_grad():
         batch_query = process_queries(processor, [query], mock_image)
@@ -26,17 +24,8 @@ def search(query: str, ds, images, k):
     # run evaluation
     retriever_evaluator = CustomEvaluator(is_multi_vector=True)
     scores = retriever_evaluator.evaluate(qs, ds)
-
-    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
-
-    #best_page = int(scores.argmax(axis=1).item())
-    #return f"The most relevant page is {best_page}", images[best_page]
+    best_page = int(scores.argmax(axis=1).item())
+    return f"The most relevant page is {best_page}", images[best_page]
 
 
 def index(file, ds):
@@ -59,6 +48,7 @@ def index(file, ds):
         ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
     return f"Uploaded and converted {len(images)} pages", ds, images
 
+
 COLORS = ["#4285f4", "#db4437", "#f4b400", "#0f9d58", "#e48ef1"]
 # Load model
 model_name = "vidore/colpali"
@@ -90,9 +80,8 @@ 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, k], outputs=[message2, output_img])
+    search_button.click(search, inputs=[query, embeds, imgs], outputs=[message2, output_img])
 
 
 if __name__ == "__main__":