Update app.py

app.py

@@ -3,6 +3,13 @@ import sys
 
 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,
+    process_queries,
+)
+import spaces
 from pdf2image import convert_from_path
 from PIL import Image
 from torch.utils.data import DataLoader
@@ -11,11 +18,23 @@ from transformers import AutoProcessor
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), './colpali-main')))
 
-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
+# Load model
+model_name = "vidore/colpali"
+token = os.environ.get("HF_TOKEN")
+model = ColPali.from_pretrained(
+    "google/paligemma-3b-mix-448", torch_dtype=torch.bfloat16, device_map="cpu", token = token).eval()
 
+model.load_adapter(model_name)
+processor = AutoProcessor.from_pretrained(model_name, token = token)
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+if device != model.device:
+    model.to(device)
+mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
+
+
+@spaces.GPU
 def search(query: str, ds, images, k):
+        
     qs = []
     with torch.no_grad():
         batch_query = process_queries(processor, [query], mock_image)
@@ -23,7 +42,6 @@ def search(query: str, ds, images, k):
         embeddings_query = model(**batch_query)
         qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
 
-    # run evaluation
     retriever_evaluator = CustomEvaluator(is_multi_vector=True)
     scores = retriever_evaluator.evaluate(qs, ds)
 
@@ -36,19 +54,24 @@ def search(query: str, ds, images, k):
     return results
 
 
-def index(file, ds):
+@spaces.GPU
+def index(files, ds):
     """Example script to run inference with ColPali"""
     images = []
-    for f in file:
+    for f in files:
         images.extend(convert_from_path(f))
 
+    if len(images) >= 150:
+        raise gr.Error("The number of images in the dataset should be less than 150.")
+
     # run inference - docs
     dataloader = DataLoader(
         images,
         batch_size=4,
         shuffle=False,
         collate_fn=lambda x: process_images(processor, x),
-    )
+    )   
+          
     for batch_doc in tqdm(dataloader):
         with torch.no_grad():
             batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
@@ -56,42 +79,41 @@ 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"
-token = os.environ.get("HF_TOKEN")
-model = ColPali.from_pretrained(
-    "google/paligemma-3b-mix-448", torch_dtype=torch.bfloat16, device_map="cpu", token=token
-).eval()
-model.load_adapter(model_name)
-processor = AutoProcessor.from_pretrained(model_name, token=token)
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-if device != model.device:
-    model.to(device)
-mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
 
-with gr.Blocks() as demo:
-    gr.Markdown("# ColPali: Efficient Document Retrieval with Vision Language Models 📚🔍")
+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 📚")
     gr.Markdown("## 1️⃣ Upload PDFs")
-    file = gr.File(file_types=["pdf"], file_count="multiple")
 
-    gr.Markdown("## 2️⃣ Convert the PDFs and upload")
-    convert_button = gr.Button("🔄 Convert and upload")
-    message = gr.Textbox("Files not yet uploaded")
-    embeds = gr.State(value=[])
-    imgs = gr.State(value=[])
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("## 1️⃣ Upload PDFs")
+            file = gr.File(file_types=["pdf"], file_count="multiple", label="Upload PDFs")
 
-    # Define the actions
-    convert_button.click(index, inputs=[file, embeds], outputs=[message, embeds, imgs])
+            convert_button = gr.Button("🔄 Index documents")
+            message = gr.Textbox("Files not yet uploaded", label="Status")
+            embeds = gr.State(value=[])
+            imgs = gr.State(value=[])
 
-    gr.Markdown("## 3️⃣ Search")
-    query = gr.Textbox(placeholder="Enter your query here")
-    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)
+        with gr.Column(scale=3):
+            gr.Markdown("## 2️⃣ Search")
+            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)
 
-    search_button.click(search, inputs=[query, embeds, imgs, k], outputs=[message2, output_img])
+    # with gr.Row():
+    #    gr.Examples(
+    #        examples=get_example(),
+    #        inputs=[file, query],
+    #    )
+
+    # Define the actions
+    search_button = gr.Button("🔍 Search", variant="primary")
+    output_gallery = gr.Gallery(label="Retrieved Documents", height=600, show_label=True)
+
+    convert_button.click(index, inputs=[file, embeds], outputs=[message, embeds, imgs])
+    search_button.click(search, inputs=[query, embeds, imgs, k], outputs=[output_gallery])
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch(debug=True)
+    demo.queue(max_size=10).launch(debug=True)