Update everything a bit !

app.py

@@ -3,7 +3,6 @@ import spaces
 
 import gradio as gr
 import torch
-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,
@@ -13,19 +12,18 @@ from pdf2image import convert_from_path
 from PIL import Image
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-from transformers import AutoProcessor
 
-# Load model
-model_name = "vidore/colpali-v1.2"
-token = os.environ.get("HF_TOKEN")
-model = ColPali.from_pretrained(
-    "vidore/colpaligemma-3b-pt-448-base", torch_dtype=torch.bfloat16, device_map="cuda", token = token).eval()
+from colpali_engine.models import ColQwen2, ColQwen2Processor
 
-model.load_adapter(model_name)
-model = model.eval()
-processor = AutoProcessor.from_pretrained(model_name, token = token)
 
-mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
+
+model = ColQwen2.from_pretrained(
+        "manu/colqwen2-v1.0-alpha",
+        torch_dtype=torch.bfloat16,
+        device_map="cuda:0",  # or "mps" if on Apple Silicon
+    ).eval()
+processor = ColQwen2Processor.from_pretrained("manu/colqwen2-v1.0-alpha")
+
 
 
 @spaces.GPU
@@ -37,15 +35,13 @@ def search(query: str, ds, images, k):
         
     qs = []
     with torch.no_grad():
-        batch_query = process_queries(processor, [query], mock_image)
-        batch_query = {k: v.to(device) for k, v in batch_query.items()}
+        batch_query = processor.process_queries([query]).to(model.device)
         embeddings_query = model(**batch_query)
         qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
 
-    retriever_evaluator = CustomEvaluator(is_multi_vector=True)
-    scores = retriever_evaluator.evaluate(qs, ds)
+    scores = processor.score(qs, ds, device=device)
 
-    top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
+    top_k_indices = scores[0].topk(k).indices.tolist()
 
     results = []
     for idx in top_k_indices:
@@ -75,21 +71,19 @@ def convert_files(files):
 @spaces.GPU
 def index_gpu(images, ds):
     """Example script to run inference with ColPali"""
-    
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    if device != model.device:
+        model.to(device)
+        
     # run inference - docs
     dataloader = DataLoader(
         images,
         batch_size=4,
         shuffle=False,
-        collate_fn=lambda x: process_images(processor, x),
+        collate_fn=lambda x: processor.process_images(x).to(model.device),
     )
 
-    
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    if device != model.device:
-        model.to(device)
-        
-          
     for batch_doc in tqdm(dataloader):
         with torch.no_grad():
             batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
@@ -98,8 +92,6 @@ def index_gpu(images, ds):
     return f"Uploaded and converted {len(images)} pages", ds, images
 
 
-def get_example():
-    return [[["climate_youth_magazine.pdf"], "How much tropical forest is cut annually ?"]]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# ColPali: Efficient Document Retrieval with Vision Language Models 📚")
@@ -128,11 +120,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             query = gr.Textbox(placeholder="Enter your query here", label="Query")
             k = gr.Slider(minimum=1, maximum=10, step=1, label="Number of results", value=5)
 
-    # with gr.Row():
-    #    gr.Examples(
-    #        examples=get_example(),
-    #        inputs=[file, query],
-    #    )
 
     # Define the actions
     search_button = gr.Button("🔍 Search", variant="primary")