Spaces:
Running
Running
import io | |
import os | |
import sys | |
from fastapi import FastAPI, File, UploadFile | |
import gradio as gr | |
import requests | |
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'))) | |
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, | |
) | |
app = FastAPI() | |
# 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)) | |
# In-memory storage | |
ds = [] | |
images = [] | |
async def index(files: List[UploadFile] = File(...)): | |
global ds, images | |
images = [] | |
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 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 {"message": f"Uploaded and converted {len(images)} pages"} | |
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")))) | |
retriever_evaluator = CustomEvaluator(is_multi_vector=True) | |
scores = retriever_evaluator.evaluate(qs, ds) | |
top_k_indices = scores.argsort(axis=1)[0][-k:][::-1] | |
results = [{"page": idx, "image": "image_placeholder"} for idx in top_k_indices] | |
return {"results": results} | |
def index_gradio(file, ds): | |
"""Upload PDFs and get embeddings.""" | |
url = "http://localhost:8082/index" | |
files = [("files", (f.name, f.file)) for f in file] | |
response = requests.post(url, files=files) | |
result = response.json() | |
return result['message'], ds, [] | |
def search_gradio(query: str, ds, images, k): | |
"""Send a search query and get results.""" | |
url = "http://localhost:8082/search" | |
payload = {'query': query, 'k': k} | |
response = requests.post(url, json=payload) | |
result = response.json() | |
results = result['results'] | |
return results | |
with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
gr.Markdown("# ColPali: Efficient Document Retrieval with Vision Language Models 📚") | |
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=[]) | |
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=1) | |
# 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_gradio, inputs=[file, embeds], outputs=[message, embeds, imgs]) | |
search_button.click(search_gradio, inputs=[query, embeds, imgs, k], outputs=[output_gallery]) | |
if __name__ == "__main__": | |
demo.queue(max_size=10).launch(debug=True) |