""" Gradio Demo for Document Retrieval - Hugging Face Spaces with ZeroGPU This script creates a Gradio interface for testing both BiGemma3 and ColGemma3 models with PDF document upload, automatic conversion to images, and query-based retrieval. Features: - PDF upload with automatic conversion to images - Model selection: NetraEmbed (BiGemma3), ColNetraEmbed (ColGemma3), or Both - Query input with top-k selection (default: 5) - Similarity score display - Side-by-side comparison when both models are selected - ZeroGPU integration for efficient GPU usage """ import io import gc import math from typing import List, Optional, Tuple import gradio as gr import torch import spaces from pdf2image import convert_from_path from PIL import Image import matplotlib.pyplot as plt import numpy as np import seaborn as sns from einops import rearrange # Import from colpali_engine from colpali_engine.models import BiGemma3, BiGemmaProcessor3, ColGemma3, ColGemmaProcessor3 from colpali_engine.interpretability import get_similarity_maps_from_embeddings from colpali_engine.interpretability.similarity_map_utils import normalize_similarity_map # Configuration MAX_BATCH_SIZE = 32 # Maximum pages to process at once device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Device: {device}") if torch.cuda.is_available(): print(f"GPU: {torch.cuda.get_device_name(0)}") # Global state for models and indexed documents class DocumentIndex: def __init__(self): self.images: List[Image.Image] = [] self.bigemma_embeddings = None self.colgemma_embeddings = None self.bigemma_model = None self.bigemma_processor = None self.colgemma_model = None self.colgemma_processor = None doc_index = DocumentIndex() # Helper functions def pdf_to_images(pdf_path: str) -> List[Image.Image]: """Convert PDF to list of PIL Images with error handling.""" try: print(f"Converting PDF to images: {pdf_path}") images = convert_from_path(pdf_path, dpi=200) print(f"Converted {len(images)} pages") return images except Exception as e: print(f"❌ PDF conversion error: {str(e)}") raise gr.Error(f"Failed to convert PDF: {str(e)}") @spaces.GPU def load_bigemma_model(): """Load BiGemma3 model and processor.""" if doc_index.bigemma_model is None: print("Loading BiGemma3 (NetraEmbed)...") try: doc_index.bigemma_processor = BiGemmaProcessor3.from_pretrained( "Cognitive-Lab/NetraEmbed", use_fast=True, ) doc_index.bigemma_model = BiGemma3.from_pretrained( "Cognitive-Lab/NetraEmbed", torch_dtype=torch.bfloat16, device_map=device, ) doc_index.bigemma_model.eval() print("✓ BiGemma3 loaded successfully") except Exception as e: print(f"❌ Failed to load BiGemma3: {str(e)}") raise gr.Error(f"Failed to load BiGemma3: {str(e)}") return "✅ BiGemma3 loaded" @spaces.GPU def load_colgemma_model(): """Load ColGemma3 model and processor.""" if doc_index.colgemma_model is None: print("Loading ColGemma3 (ColNetraEmbed)...") try: doc_index.colgemma_model = ColGemma3.from_pretrained( "Cognitive-Lab/ColNetraEmbed", dtype=torch.bfloat16, device_map=device, ) doc_index.colgemma_model.eval() doc_index.colgemma_processor = ColGemmaProcessor3.from_pretrained( "Cognitive-Lab/ColNetraEmbed", use_fast=True, ) print("✓ ColGemma3 loaded successfully") except Exception as e: print(f"❌ Failed to load ColGemma3: {str(e)}") raise gr.Error(f"Failed to load ColGemma3: {str(e)}") return "✅ ColGemma3 loaded" def unload_models(): """Unload models and free GPU memory.""" try: if doc_index.bigemma_model is not None: del doc_index.bigemma_model del doc_index.bigemma_processor doc_index.bigemma_model = None doc_index.bigemma_processor = None if doc_index.colgemma_model is not None: del doc_index.colgemma_model del doc_index.colgemma_processor doc_index.colgemma_model = None doc_index.colgemma_processor = None # Clear embeddings and images doc_index.bigemma_embeddings = None doc_index.colgemma_embeddings = None doc_index.images = [] # Force garbage collection gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() return "✅ Models unloaded and GPU memory cleared" except Exception as e: return f"❌ Error unloading models: {str(e)}" @spaces.GPU def index_bigemma_images(images: List[Image.Image]) -> torch.Tensor: """Index images with BiGemma3 model.""" model, processor = doc_index.bigemma_model, doc_index.bigemma_processor batch_images = processor.process_images(images).to(device) embeddings = model(**batch_images, embedding_dim=768) return embeddings @spaces.GPU def index_colgemma_images(images: List[Image.Image]) -> torch.Tensor: """Index images with ColGemma3 model.""" model, processor = doc_index.colgemma_model, doc_index.colgemma_processor batch_images = processor.process_images(images).to(device) embeddings = model(**batch_images) return embeddings def index_document(pdf_file, model_choice: str): """Upload and index a PDF document.""" if pdf_file is None: return "⚠️ Please upload a PDF document first." try: status = [] # Convert PDF to images status.append("⏳ Converting PDF to images...") doc_index.images = pdf_to_images(pdf_file.name) num_pages = len(doc_index.images) status.append(f"✓ Converted PDF to {num_pages} images") if num_pages > MAX_BATCH_SIZE: status.append(f"⚠️ Large PDF ({num_pages} pages). Processing in batches...") # Index with BiGemma3 if model_choice in ["NetraEmbed (BiGemma3)", "Both"]: if doc_index.bigemma_model is None: status.append("⏳ Loading BiGemma3 model...") load_bigemma_model() status.append("✓ BiGemma3 loaded") else: status.append("✓ Using cached BiGemma3 model") status.append("⏳ Encoding images with BiGemma3...") doc_index.bigemma_embeddings = index_bigemma_images(doc_index.images) status.append(f"✓ Indexed with BiGemma3 (shape: {doc_index.bigemma_embeddings.shape})") # Index with ColGemma3 if model_choice in ["ColNetraEmbed (ColGemma3)", "Both"]: if doc_index.colgemma_model is None: status.append("⏳ Loading ColGemma3 model...") load_colgemma_model() status.append("✓ ColGemma3 loaded") else: status.append("✓ Using cached ColGemma3 model") status.append("⏳ Encoding images with ColGemma3...") doc_index.colgemma_embeddings = index_colgemma_images(doc_index.images) status.append(f"✓ Indexed with ColGemma3 (shape: {doc_index.colgemma_embeddings.shape})") return "\n".join(status) + "\n\n✅ Document ready for querying!" except Exception as e: import traceback error_details = traceback.format_exc() print(f"Indexing error: {error_details}") return f"❌ Error indexing document: {str(e)}" @spaces.GPU def generate_colgemma_heatmap( image: Image.Image, query: str, query_embedding: torch.Tensor, image_embedding: torch.Tensor, model, processor, ) -> Image.Image: """Generate heatmap overlay for ColGemma3 results.""" try: # Re-process the single image to get the proper batch_images dict for image mask batch_images = processor.process_images([image]).to(device) # Create image mask manually if "input_ids" in batch_images and hasattr(model.config, "image_token_id"): image_token_id = model.config.image_token_id image_mask = batch_images["input_ids"] == image_token_id else: image_mask = torch.ones( image_embedding.shape[0], image_embedding.shape[1], dtype=torch.bool, device=device ) # Calculate n_patches from actual number of image tokens num_image_tokens = image_mask.sum().item() n_side = int(math.sqrt(num_image_tokens)) if n_side * n_side == num_image_tokens: n_patches = (n_side, n_side) else: n_patches = (16, 16) # Generate similarity maps similarity_maps_list = get_similarity_maps_from_embeddings( image_embeddings=image_embedding, query_embeddings=query_embedding, n_patches=n_patches, image_mask=image_mask, ) similarity_map = similarity_maps_list[0] # Aggregate across all query tokens if similarity_map.dtype == torch.bfloat16: similarity_map = similarity_map.float() aggregated_map = torch.mean(similarity_map, dim=0) # Convert the image to an array img_array = np.array(image.convert("RGBA")) # Normalize the similarity map similarity_map_array = normalize_similarity_map(aggregated_map).to(torch.float32).cpu().numpy() similarity_map_array = rearrange(similarity_map_array, "h w -> w h") # Create PIL image from similarity map similarity_map_image = Image.fromarray((similarity_map_array * 255).astype("uint8")).resize( image.size, Image.Resampling.BICUBIC ) # Create matplotlib figure _, ax = plt.subplots(figsize=(10, 10)) ax.imshow(img_array) ax.imshow( similarity_map_image, cmap=sns.color_palette("mako", as_cmap=True), alpha=0.5, ) ax.set_axis_off() plt.tight_layout() # Convert to PIL Image buffer = io.BytesIO() plt.savefig(buffer, format="png", dpi=150, bbox_inches="tight", pad_inches=0) buffer.seek(0) heatmap_image = Image.open(buffer).copy() plt.close() return heatmap_image except Exception as e: print(f"❌ Heatmap generation error: {str(e)}") return image @spaces.GPU def query_bigemma(query: str, top_k: int) -> Tuple[str, List]: """Query indexed documents with BiGemma3.""" model, processor = doc_index.bigemma_model, doc_index.bigemma_processor # Encode query batch_query = processor.process_texts([query]).to(device) query_embedding = model(**batch_query, embedding_dim=768) # Compute scores scores = processor.score(qs=query_embedding, ps=doc_index.bigemma_embeddings) # Get top-k results top_k_actual = min(top_k, len(doc_index.images)) top_indices = scores[0].argsort(descending=True)[:top_k_actual] # Format results results_text = "### BiGemma3 (NetraEmbed) Results\n\n" gallery_images = [] for rank, idx in enumerate(top_indices): score = scores[0, idx].item() results_text += f"**Rank {rank + 1}:** Page {idx.item() + 1} - Score: {score:.4f}\n" gallery_images.append( (doc_index.images[idx.item()], f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.4f})") ) return results_text, gallery_images @spaces.GPU def query_colgemma(query: str, top_k: int, show_heatmap: bool = False) -> Tuple[str, List]: """Query indexed documents with ColGemma3.""" model, processor = doc_index.colgemma_model, doc_index.colgemma_processor # Encode query batch_query = processor.process_queries([query]).to(device) query_embedding = model(**batch_query) # Compute scores scores = processor.score_multi_vector(qs=query_embedding, ps=doc_index.colgemma_embeddings) # Get top-k results top_k_actual = min(top_k, len(doc_index.images)) top_indices = scores[0].argsort(descending=True)[:top_k_actual] # Format results results_text = "### ColGemma3 (ColNetraEmbed) Results\n\n" gallery_images = [] for rank, idx in enumerate(top_indices): score = scores[0, idx].item() results_text += f"**Rank {rank + 1}:** Page {idx.item() + 1} - Score: {score:.2f}\n" # Generate heatmap if requested if show_heatmap: heatmap_image = generate_colgemma_heatmap( image=doc_index.images[idx.item()], query=query, query_embedding=query_embedding, image_embedding=doc_index.colgemma_embeddings[idx.item()].unsqueeze(0), model=model, processor=processor, ) gallery_images.append( (heatmap_image, f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.2f})") ) else: gallery_images.append( (doc_index.images[idx.item()], f"Rank {rank + 1} - Page {idx.item() + 1} (Score: {score:.2f})") ) return results_text, gallery_images def query_documents( query: str, model_choice: str, top_k: int, show_heatmap: bool = False ) -> Tuple[Optional[str], Optional[str], Optional[List], Optional[List]]: """Query the indexed documents.""" if not doc_index.images: return "⚠️ Please upload and index a document first.", None, None, None if not query.strip(): return "⚠️ Please enter a query.", None, None, None try: results_bi = None results_col = None gallery_images_bi = [] gallery_images_col = [] # Query with BiGemma3 if model_choice in ["NetraEmbed (BiGemma3)", "Both"]: if doc_index.bigemma_embeddings is None: return "⚠️ Please index the document with BiGemma3 first.", None, None, None results_bi, gallery_images_bi = query_bigemma(query, top_k) # Query with ColGemma3 if model_choice in ["ColNetraEmbed (ColGemma3)", "Both"]: if doc_index.colgemma_embeddings is None: return "⚠️ Please index the document with ColGemma3 first.", None, None, None results_col, gallery_images_col = query_colgemma(query, top_k, show_heatmap) # Return results based on model choice if model_choice == "NetraEmbed (BiGemma3)": return results_bi, None, gallery_images_bi, None elif model_choice == "ColNetraEmbed (ColGemma3)": return results_col, None, None, gallery_images_col else: # Both return results_bi, results_col, gallery_images_bi, gallery_images_col except Exception as e: import traceback error_details = traceback.format_exc() print(f"Query error: {error_details}") return f"❌ Error during query: {str(e)}", None, None, None # Create Gradio interface with gr.Blocks(title="NetraEmbed Demo") as demo: # Header section gr.Markdown("# NetraEmbed") gr.HTML( """
Paper GitHub Model
""" ) gr.Markdown( """ **🚀 Universal Multilingual Multimodal Document Retrieval** Upload a PDF document, select your model(s), and query using semantic search. **Available Models:** - **NetraEmbed (BiGemma3)**: Single-vector embedding - Fast retrieval with cosine similarity - **ColNetraEmbed (ColGemma3)**: Multi-vector embedding - High-quality retrieval with MaxSim scoring and heatmaps """ ) with gr.Row(): # Column 1: Model Selection with gr.Column(scale=1): gr.Markdown("### 🤖 Model Selection") model_select = gr.Radio( choices=["NetraEmbed (BiGemma3)", "ColNetraEmbed (ColGemma3)", "Both"], value="Both", label="Select Model(s)", ) # Column 2: Document Upload with gr.Column(scale=1): gr.Markdown("### 📄 Upload & Index") pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"]) index_btn = gr.Button("📥 Index Document", variant="primary") index_status = gr.Textbox(label="Status", lines=6, interactive=False) # Column 3: Query with gr.Column(scale=1): gr.Markdown("### 🔎 Query Document") query_input = gr.Textbox( label="Enter Query", placeholder="e.g., financial report, organizational structure...", lines=2, ) with gr.Row(): top_k_slider = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Top K", scale=2) heatmap_checkbox = gr.Checkbox(label="Heatmaps", value=False, scale=1) query_btn = gr.Button("🔍 Search", variant="primary") gr.Markdown("---") # Results section gr.Markdown("### 📊 Results") with gr.Row(): with gr.Column(scale=1): bigemma_results = gr.Markdown(value="*BiGemma3 results will appear here...*") bigemma_gallery = gr.Gallery( label="BiGemma3 - Top Retrieved Pages", columns=2, height="auto", ) with gr.Column(scale=1): colgemma_results = gr.Markdown(value="*ColGemma3 results will appear here...*") colgemma_gallery = gr.Gallery( label="ColGemma3 - Top Retrieved Pages", columns=2, height="auto", ) # Event handlers index_btn.click( fn=index_document, inputs=[pdf_upload, model_select], outputs=[index_status], ) query_btn.click( fn=query_documents, inputs=[query_input, model_select, top_k_slider, heatmap_checkbox], outputs=[bigemma_results, colgemma_results, bigemma_gallery, colgemma_gallery], ) # Launch the app demo.launch()