Multi-modal RAG

import base64
from openai import OpenAI
from sentence_transformers import SentenceTransformer

client = OpenAI()
embedder = SentenceTransformer("BAAI/bge-large-en-v1.5")

def caption_image(image_path: str, context: str = "") -> str:
    """Use GPT-4o to generate a retrieval-optimized caption."""
    with open(image_path, "rb") as f:
        image_b64 = base64.b64encode(f.read()).decode()

    prompt = (
        "You are generating a detailed description of this image for "
        "a search index. Include ALL visible text, numbers, labels, "
        "legends, axis titles, and visual patterns. If it's a table, "
        "reproduce the full table structure. If it's a chart, describe "
        "trends, values, and comparisons."
    )
    if context:
        prompt += f"\nSurrounding document context: {context}"

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {
                    "url": f"data:image/png;base64,{image_b64}"
                }}
            ]
        }],
        max_tokens=1000
    )
    return response.choices[0].message.content

def build_multimodal_index(documents: list):
    """Build a unified text index from mixed content."""
    entries = []

    for doc in documents:
        if doc["type"] == "text":
            entries.append({
                "content": doc["text"],
                "source": doc["source"],
                "modality": "text",
            })
        elif doc["type"] == "image":
            caption = caption_image(
                doc["path"],
                context=doc.get("surrounding_text", "")
            )
            entries.append({
                "content": caption,
                "source": doc["path"],
                "modality": "image",
                "caption": caption,  # store for display
            })
        elif doc["type"] == "table":
            # Tables get BOTH Markdown and NL description
            caption = caption_image(doc["screenshot_path"])
            entries.append({
                "content": caption,
                "source": doc["source"],
                "modality": "table",
                "markdown": doc.get("markdown", ""),
            })

    texts = [e["content"] for e in entries]
    embeddings = embedder.encode(texts, normalize_embeddings=True)
    return entries, embeddings  # Store in your vector DB

from colpali_engine.models import ColPali, ColPaliProcessor
from pdf2image import convert_from_path
import torch

# Load ColPali (requires ~6GB GPU memory)
model = ColPali.from_pretrained(
    "vidore/colpali-v1.2",
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)
processor = ColPaliProcessor.from_pretrained("vidore/colpali-v1.2")

def index_pdf(pdf_path: str) -> list[torch.Tensor]:
    """Convert each page to an image and embed it."""
    pages = convert_from_path(pdf_path, dpi=144)
    page_embeddings = []

    for page_image in pages:
        inputs = processor.process_images([page_image])
        inputs = {k: v.to("cuda") for k, v in inputs.items()}
        with torch.no_grad():
            embs = model(**inputs)  # (1, n_patches, 128)
        page_embeddings.append(embs.squeeze(0).cpu())

    return page_embeddings  # List of (n_patches, 128) tensors

def search(query: str, page_embeddings: list, top_k: int = 5):
    """Retrieve most relevant pages using MaxSim scoring."""
    inputs = processor.process_queries([query])
    inputs = {k: v.to("cuda") for k, v in inputs.items()}
    with torch.no_grad():
        q_embs = model(**inputs).squeeze(0)  # (n_tokens, 128)

    scores = []
    for page_idx, p_embs in enumerate(page_embeddings):
        # MaxSim: for each query token, max similarity across patches
        sim_matrix = q_embs @ p_embs.to("cuda").T  # (n_tokens, n_patches)
        max_sims = sim_matrix.max(dim=1).values     # (n_tokens,)
        score = max_sims.sum().item()
        scores.append((page_idx, score))

    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]

# Usage
embeddings = index_pdf("quarterly_report.pdf")
results = search("Q3 revenue breakdown by segment", embeddings)
for page_idx, score in results:
    print(f"Page {page_idx + 1}: score={score:.2f}")