Spaces:
Running on Zero
Running on Zero
| """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() | |
| 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)) | |
| 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()) | |