Spaces:
Running on Zero
Running on Zero
| """Cache pinned official adapter/base snapshots before loading model code.""" | |
| import json | |
| from pathlib import Path | |
| from huggingface_hub import HfApi, snapshot_download | |
| from visual_doc_rag.data import write_json | |
| MODEL_ID = "ModernVBERT/colmodernvbert" | |
| def cache_model(data_dir: Path, revision: str = "main", base_revision: str | None = None) -> dict: | |
| root = data_dir / "models" | |
| manifest_path = root / "manifest.json" | |
| if manifest_path.exists(): | |
| cached = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| if revision in {"main", cached["adapter_revision"]} and base_revision in { | |
| None, | |
| cached["base_revision"], | |
| }: | |
| return cached | |
| api = HfApi() | |
| adapter_sha = api.model_info(MODEL_ID, revision=revision).sha | |
| adapter_dir = root / "adapter" / adapter_sha | |
| snapshot_download( | |
| MODEL_ID, | |
| revision=adapter_sha, | |
| local_dir=adapter_dir, | |
| cache_dir=root / "hub", | |
| allow_patterns=["*.json", "*.safetensors", "*.jinja", "README.md", "LICENSE*"], | |
| max_workers=2, | |
| ) | |
| adapter = json.loads((adapter_dir / "adapter_config.json").read_text(encoding="utf-8")) | |
| base_id = adapter["base_model_name_or_path"] | |
| base_sha = api.model_info( | |
| base_id, revision=base_revision or adapter.get("revision") or "main" | |
| ).sha | |
| base_dir = root / "base" / base_sha | |
| snapshot_download( | |
| base_id, | |
| revision=base_sha, | |
| local_dir=base_dir, | |
| cache_dir=root / "hub", | |
| allow_patterns=["*.json", "*.safetensors", "*.jinja", "README.md", "LICENSE*"], | |
| max_workers=2, | |
| ) | |
| manifest = { | |
| "model_id": MODEL_ID, | |
| "adapter_revision": adapter_sha, | |
| "base_model_id": base_id, | |
| "base_revision": base_sha, | |
| "adapter_path": str(adapter_dir.relative_to(root)), | |
| "base_path": str(base_dir.relative_to(root)), | |
| "license": "MIT", | |
| } | |
| write_json(manifest_path, manifest) | |
| return manifest | |