src/visual_doc_rag/visual_index.py

"""Official ColModernVBERT encoding, float16 disk cache, CPU MaxSim search."""

import hashlib
import json
from importlib.metadata import version
from pathlib import Path
from time import perf_counter

import torch
from PIL import Image
from torch.nn.utils.rnn import pad_sequence

from visual_doc_rag.data import write_json
from visual_doc_rag.maxsim import maxsim
from visual_doc_rag.models import cache_model
from visual_doc_rag.pages import load_pages
from visual_doc_rag.retrieve import rank_scores


class Encoder:
    def __init__(self, data_dir: Path, model_info: dict, device: str):
        from colpali_engine.models import ColModernVBert, ColModernVBertProcessor

        self.device = torch.device(device)
        self.processor = ColModernVBertProcessor.from_pretrained(
            data_dir / "models" / model_info["adapter_path"],
            local_files_only=True,
        )
        self.model = ColModernVBert.from_pretrained(
            data_dir / "models" / model_info["base_path"],
            local_files_only=True,
            dtype=torch.float32,
            attn_implementation="sdpa",
        )
        self.model.load_adapter(
            str(data_dir / "models" / model_info["adapter_path"]),
            adapter_kwargs={"local_files_only": True},
        )
        self.model.to(self.device).eval()

    @torch.inference_mode()
    def encode(self, *, image=None, query: str | None = None) -> torch.Tensor:
        if (image is None) == (query is None):
            raise ValueError("Provide exactly one image or query.")
        inputs = (
            self.processor.process_images([image])
            if image is not None
            else self.processor.process_texts([query])
        ).to(self.device)
        embedding = self.model(**inputs)[0]
        # Remove padding, preserving every genuine text/image/prompt token.
        embedding = embedding[inputs["attention_mask"][0].bool()].detach().cpu()
        if embedding.ndim != 2 or embedding.shape[0] == 0 or not torch.isfinite(embedding).all():
            raise ValueError("Model returned an empty or nonfinite embedding.")
        return embedding


def build_visual(
    subset: str,
    data_dir: Path,
    source: dict,
    device: str,
    model_revision: str = "main",
    base_revision: str | None = None,
) -> dict:
    started = perf_counter()
    model_info = cache_model(data_dir, model_revision, base_revision)
    page_root, page_manifest = load_pages(subset, data_dir, source["revision"])
    signature = {
        "dataset_revision": source["revision"],
        "model": model_info,
        "encoder_dtype": "float32",
        "storage_dtype": "float16",
        "processor": "official defaults; process_texts without suffix; all unpadded tokens",
        "attention": "sdpa",
        "page_format": "original RGB PNG",
        "library_versions": {
            name: version(name) for name in ("colpali-engine", "transformers", "torch")
        },
    }
    fingerprint = hashlib.sha256(json.dumps(signature, sort_keys=True).encode()).hexdigest()[:16]
    root = data_dir / subset / "indexes"
    embedding_root = root / "visual" / fingerprint
    embedding_root.mkdir(parents=True, exist_ok=True)
    if device == "cuda":
        if not torch.cuda.is_available():
            raise ValueError("CUDA unavailable. Run outside the GPU sandbox or use --device cpu.")
        torch.cuda.reset_peak_memory_stats()
    loading = perf_counter()
    encoder = Encoder(data_dir, model_info, device)
    load_seconds = perf_counter() - loading
    computed = 0
    pages = []
    for index, page in enumerate(page_manifest["pages"]):
        page_id = page["corpus_id"]
        target = embedding_root / f"{page_id}.pt"
        if target.exists():
            embedding = torch.load(target, map_location="cpu", weights_only=True)
        else:
            with Image.open(page_root / page["file"]) as image:
                embedding = encoder.encode(image=image).to(torch.float16)
            temporary = target.with_suffix(".tmp")
            torch.save(embedding, temporary)
            temporary.replace(target)
            computed += 1
        if embedding.ndim != 2 or embedding.shape[1] != 128 or not torch.isfinite(embedding).all():
            raise ValueError(f"Invalid cached embedding: {target}. Remove this file and retry.")
        pages.append(
            {
                "corpus_id": page_id,
                "file": str(target.relative_to(root)),
                "tokens": embedding.shape[0],
            }
        )
        if (index + 1) % 25 == 0 or index + 1 == len(page_manifest["pages"]):
            print(f"Visual index: {index + 1}/{len(page_manifest['pages'])}", flush=True)
    result = {
        "dataset": source,
        "signature": signature,
        "pages": pages,
        "build_seconds": perf_counter() - started,
        "model_load_seconds": load_seconds,
        "new_pages_encoded": computed,
        "cached_pages_reused": len(pages) - computed,
        "index_bytes": sum((root / page["file"]).stat().st_size for page in pages),
        "peak_vram_bytes": torch.cuda.max_memory_allocated() if device == "cuda" else 0,
        "device": device,
        "device_name": torch.cuda.get_device_name() if device == "cuda" else "CPU",
    }
    write_json(root / "visual.json", result)
    return result


class VisualIndex:
    def __init__(self, root: Path, source: dict, batch_size: int = 32):
        self.manifest = json.loads((root / "visual.json").read_text(encoding="utf-8"))
        if self.manifest["dataset"]["revision"] != source["revision"]:
            raise ValueError("Visual index dataset revision differs from downloaded data.")
        self.page_ids = [page["corpus_id"] for page in self.manifest["pages"]]
        self.batches = []
        pages = self.manifest["pages"]
        for start in range(0, len(pages), batch_size):
            embeddings = [
                torch.load(root / page["file"], map_location="cpu", weights_only=True).float()
                for page in pages[start : start + batch_size]
            ]
            padded = pad_sequence(embeddings, batch_first=True)
            mask = (
                torch.arange(padded.shape[1])[None, :]
                < torch.tensor([embedding.shape[0] for embedding in embeddings])[:, None]
            )
            self.batches.append((padded, mask))

    @torch.inference_mode()
    def search(self, query: torch.Tensor) -> list[int]:
        scores = torch.cat([maxsim(query, pages, mask) for pages, mask in self.batches])
        return rank_scores(self.page_ids, scores.tolist())