Update app.py

app.py

@@ -1,22 +1,25 @@
+import io
 import os
 import sys
 
+from fastapi import FastAPI, File, UploadFile
+from typing import List
+import torch
+from pdf2image import convert_from_path
+from PIL import Image
+from torch.utils.data import DataLoader
+from transformers import AutoProcessor
+
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), './colpali-main')))
 
-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
-from tqdm import tqdm
-from transformers import AutoProcessor
+
+app = FastAPI()
 
 # Load model
 model_name = "vidore/colpali"
@@ -31,88 +34,52 @@ if device != model.device:
     model.to(device)
 mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
 
+# In-memory storage
+ds = []
+images = []
 
-@spaces.GPU
-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()}
-        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)
-
-    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
-
-
-@spaces.GPU
-def index(files, ds):
-    """Example script to run inference with ColPali"""
+@app.post("/index")
+async def index(files: List[UploadFile] = File(...)):
+    global ds, images
     images = []
-    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
+    ds = []
+    for file in files:
+        content = await file.read()
+        pdf_image_list = convert_from_path(io.BytesIO(content))
+        images.extend(pdf_image_list)
+    
     dataloader = DataLoader(
         images,
         batch_size=4,
         shuffle=False,
         collate_fn=lambda x: process_images(processor, x),
-    )   
-          
-    for batch_doc in tqdm(dataloader):
+    )
+    for batch_doc in dataloader:
         with torch.no_grad():
             batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
             embeddings_doc = model(**batch_doc)
         ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
-    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 📚")
+    
+    return {"message": f"Uploaded and converted {len(images)} pages"}
 
-    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")
-
-            convert_button = gr.Button("🔄 Index documents")
-            message = gr.Textbox("Files not yet uploaded", label="Status")
-            embeds = gr.State(value=[])
-            imgs = gr.State(value=[])
+@app.post("/search")
+async def search(query: str, k: int):
+    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()}
+        embeddings_query = model(**batch_query)
+        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
 
-        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)
+    retriever_evaluator = CustomEvaluator(is_multi_vector=True)
+    scores = retriever_evaluator.evaluate(qs, ds)
 
-    # with gr.Row():
-    #    gr.Examples(
-    #        examples=get_example(),
-    #        inputs=[file, query],
-    #    )
+    top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
 
-    # Define the actions
-    search_button = gr.Button("🔍 Search", variant="primary")
-    output_gallery = gr.Gallery(label="Retrieved Documents", height=600, show_label=True)
+    results = [{"page": idx, "image": "image_placeholder"} for idx in top_k_indices]
 
-    convert_button.click(index, inputs=[file, embeds], outputs=[message, embeds, imgs])
-    search_button.click(search, inputs=[query, embeds, imgs, k], outputs=[output_gallery])
+    return {"results": results}
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch(debug=True)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8082, reload=True)