Spaces:
Sleeping
Sleeping
| import os | |
| import base64 | |
| import tempfile | |
| from io import BytesIO | |
| from urllib.request import urlretrieve | |
| import gradio as gr | |
| from gradio_pdf import PDF | |
| import torch | |
| from pdf2image import convert_from_path | |
| from PIL import Image | |
| from torch.utils.data import DataLoader | |
| from tqdm import tqdm | |
| from colpali_engine.models import ColQwen2, ColQwen2Processor | |
| # ----------------------------- | |
| # Globals | |
| # ----------------------------- | |
| api_key = os.getenv("OPENAI_API_KEY", "") # <- use env var | |
| ds = [] # list of document embeddings (torch tensors) | |
| images = [] # list of PIL images (page-order) | |
| current_pdf_path = None # last (indexed) pdf path for preview | |
| # ----------------------------- | |
| # Model & processor | |
| # ----------------------------- | |
| device_map = "cuda:0" if torch.cuda.is_available() else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu") | |
| model = ColQwen2.from_pretrained( | |
| "vidore/colqwen2-v1.0", | |
| torch_dtype=torch.bfloat16, | |
| device_map=device_map, | |
| attn_implementation="flash_attention_2" | |
| ).eval() | |
| processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v1.0") | |
| # ----------------------------- | |
| # Utilities | |
| # ----------------------------- | |
| def encode_image_to_base64(image: Image.Image) -> str: | |
| """Encodes a PIL image to a base64 string.""" | |
| buffered = BytesIO() | |
| image.save(buffered, format="JPEG") | |
| return base64.b64encode(buffered.getvalue()).decode("utf-8") | |
| def query_gpt(query: str, retrieved_images: list[tuple[Image.Image, str]]) -> str: | |
| """Calls OpenAI's GPT model with the query and image data.""" | |
| if api_key and api_key.startswith("sk"): | |
| try: | |
| from openai import OpenAI | |
| base64_images = [encode_image_to_base64(im_caption[0]) for im_caption in retrieved_images] | |
| client = OpenAI(api_key=api_key.strip()) | |
| PROMPT = """ | |
| You are a smart assistant designed to answer questions about a PDF document. | |
| You are given relevant information in the form of PDF pages. Use them to construct a short response to the question, and cite your sources (page numbers, etc). | |
| If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents. | |
| Give detailed and extensive answers, only containing info in the pages you are given. | |
| You can answer using information contained in plots and figures if necessary. | |
| Answer in the same language as the query. | |
| Query: {query} | |
| PDF pages: | |
| """.strip() | |
| response = client.chat.completions.create( | |
| model="gpt-5-mini", | |
| messages=[ | |
| { | |
| "role": "user", | |
| "content": ( | |
| [{"type": "text", "text": PROMPT.format(query=query)}] + | |
| [{"type": "image_url", | |
| "image_url": {"url": f"data:image/jpeg;base64,{im}"}} | |
| for im in base64_images] | |
| ) | |
| } | |
| ], | |
| max_tokens=500, | |
| ) | |
| return response.choices[0].message.content | |
| except Exception: | |
| return "OpenAI API connection failure. Verify that OPENAI_API_KEY is set and valid (sk-***)." | |
| return "Set OPENAI_API_KEY in your environment to get a custom response." | |
| def _ensure_model_device(): | |
| dev = "cuda:0" if torch.cuda.is_available() else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu") | |
| if str(model.device) != dev: | |
| model.to(dev) | |
| return dev | |
| # ----------------------------- | |
| # Indexing helpers | |
| # ----------------------------- | |
| def convert_files(pdf_path: str) -> list[Image.Image]: | |
| """Convert a single PDF path into a list of PIL Images (pages).""" | |
| imgs = convert_from_path(pdf_path, thread_count=4) | |
| if len(imgs) >= 500: | |
| raise gr.Error("The number of images in the dataset should be less than 500.") | |
| return imgs | |
| def index_gpu(imgs: list[Image.Image]) -> str: | |
| """Embed a list of images (pages) with ColPali and store in globals.""" | |
| global ds, images | |
| device = _ensure_model_device() | |
| # reset previous dataset | |
| ds = [] | |
| images = imgs | |
| dataloader = DataLoader( | |
| images, | |
| batch_size=4, | |
| shuffle=False, | |
| collate_fn=lambda x: processor.process_images(x).to(model.device), | |
| ) | |
| for batch_doc in tqdm(dataloader, desc="Indexing pages"): | |
| 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"Indexed {len(images)} pages successfully." | |
| def index_from_path(pdf_path: str) -> str: | |
| """Public: index a local PDF file path.""" | |
| imgs = convert_files(pdf_path) | |
| return index_gpu(imgs) | |
| def index_from_url(url: str) -> tuple[str, str]: | |
| """ | |
| Download a PDF from URL and index it. | |
| Returns: | |
| status message, saved pdf path | |
| """ | |
| tmp_dir = tempfile.mkdtemp(prefix="colpali_") | |
| local_path = os.path.join(tmp_dir, "document.pdf") | |
| urlretrieve(url, local_path) | |
| status = index_from_path(local_path) | |
| return status, local_path | |
| # ----------------------------- | |
| # Search (MCP tool-friendly) | |
| # ----------------------------- | |
| def search(query: str, k: int): | |
| """ | |
| Search the currently indexed PDF pages for the most relevant content and | |
| generate an answer grounded ONLY in those pages. | |
| MCP tool description: | |
| - name: search | |
| - description: Retrieve top-k PDF pages relevant to a query and answer using only those pages. | |
| - input_schema: | |
| type: object | |
| properties: | |
| query: {type: string, description: "User query in natural language."} | |
| k: {type: integer, minimum: 1, maximum: 10, description: "Number of top pages to retrieve."} | |
| required: ["query"] | |
| Args: | |
| query (str): Natural-language question to search for. | |
| k (int): Number of top results to return (1–10). | |
| Returns: | |
| tuple: | |
| - results (list[tuple[PIL.Image.Image, str]]): List of (page_image, caption) pairs for a gallery. | |
| - ai_response (str): Answer grounded only in retrieved pages, with citations (page numbers). | |
| Notes: | |
| • Requires that a PDF has been indexed first. | |
| • Citations reference 1-based page numbers as shown in the gallery captions. | |
| """ | |
| global ds, images | |
| if not images or not ds: | |
| return [], "No document indexed yet. Upload a PDF or load the sample, then run Search." | |
| k = max(1, min(int(k), len(images))) | |
| device = _ensure_model_device() | |
| # Encode query | |
| qs = [] | |
| with torch.no_grad(): | |
| batch_query = processor.process_queries([query]).to(model.device) | |
| embeddings_query = model(**batch_query) | |
| qs.extend(list(torch.unbind(embeddings_query.to("cpu")))) | |
| # Score and select top-k | |
| scores = processor.score(qs, ds, device=device) | |
| top_k_indices = scores[0].topk(k).indices.tolist() | |
| # Build gallery results with 1-based page numbering | |
| results = [] | |
| for idx in top_k_indices: | |
| page_num = idx + 1 | |
| results.append((images[idx], f"Page {page_num}")) | |
| # Generate grounded response | |
| ai_response = query_gpt(query, results) | |
| return results, ai_response | |
| # ----------------------------- | |
| # Gradio UI callbacks | |
| # ----------------------------- | |
| def handle_upload(file) -> tuple[str, str | None]: | |
| """Index a user-uploaded PDF file.""" | |
| global current_pdf_path | |
| if file is None: | |
| return "Please upload a PDF.", None | |
| path = getattr(file, "name", file) | |
| status = index_from_path(path) | |
| current_pdf_path = path | |
| return status, path | |
| def handle_url(url: str) -> tuple[str, str | None]: | |
| """Index a PDF from URL (e.g., a sample).""" | |
| global current_pdf_path | |
| if not url or not url.lower().endswith(".pdf"): | |
| return "Please provide a direct PDF URL.", None | |
| status, path = index_from_url(url) | |
| current_pdf_path = path | |
| return status, path | |
| # ----------------------------- | |
| # Gradio App | |
| # ----------------------------- | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# ColPali: Efficient Document Retrieval with Vision Language Models (ColQwen2) 📚") | |
| gr.Markdown( | |
| """Demo to test ColQwen2 (ColPali) on PDF documents. | |
| ColPali is implemented from the [ColPali paper](https://arxiv.org/abs/2407.01449). | |
| This demo lets you **upload a PDF or load a sample**, then **search** for the most relevant pages and get a grounded answer. | |
| ⚠️ The model was trained on A4 portrait English PDFs; performance may drop on other formats/languages. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("## 1️⃣ Load a PDF") | |
| pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"]) | |
| index_btn = gr.Button("📥 Index Uploaded PDF", variant="secondary") | |
| url_box = gr.Textbox( | |
| label="Or index from URL", | |
| placeholder="https://example.com/file.pdf", | |
| value="https://sist.sathyabama.ac.in/sist_coursematerial/uploads/SAR1614.pdf", | |
| ) | |
| index_url_btn = gr.Button("🌐 Load Sample / From URL", variant="secondary") | |
| status_box = gr.Textbox(label="Status", interactive=False) | |
| pdf_view = PDF(label="PDF Preview") | |
| with gr.Column(scale=3): | |
| gr.Markdown("## 2️⃣ Search") | |
| query = gr.Textbox(placeholder="Enter your query here", label="Query") | |
| k_slider = gr.Slider(minimum=1, maximum=10, step=1, label="Number of results", value=5) | |
| search_button = gr.Button("🔍 Search", variant="primary") | |
| output_gallery = gr.Gallery(label="Retrieved Pages", height=600, show_label=True) | |
| output_text = gr.Textbox(label="AI Response", placeholder="Generated response based on retrieved documents") | |
| # Wiring | |
| index_btn.click(handle_upload, inputs=[pdf_input], outputs=[status_box, pdf_view]) | |
| index_url_btn.click(handle_url, inputs=[url_box], outputs=[status_box, pdf_view]) | |
| search_button.click(search, inputs=[query, k_slider], outputs=[output_gallery, output_text]) | |
| if __name__ == "__main__": | |
| # Optional: pre-load the default sample at startup. | |
| # Comment these two lines if you prefer a "cold" start. | |
| # msg, path = index_from_url("https://sist.sathyabama.ac.in/sist_coursematerial/uploads/SAR1614.pdf") | |
| # print(msg, "->", path) | |
| demo.queue(max_size=5).launch(debug=True, mcp_server=True) | |