app.py

"""
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(
        """
        <div style="display: flex; gap: 8px; flex-wrap: wrap; margin-bottom: 15px;">
            <a href="https://arxiv.org/abs/2512.03514" target="_blank">
                <img src="https://img.shields.io/badge/arXiv-2512.03514-b31b1b.svg" alt="Paper">
            </a>
            <a href="https://github.com/adithya-s-k/colpali" target="_blank">
                <img src="https://img.shields.io/badge/GitHub-colpali-181717?logo=github" alt="GitHub">
            </a>
            <a href="https://huggingface.co/Cognitive-Lab/ColNetraEmbed" target="_blank">
                <img src="https://img.shields.io/badge/🤗%20HuggingFace-Model-yellow" alt="Model">
            </a>
        </div>
        """
    )
    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()