Spaces:
Running on Zero
Running on Zero
| """Local and Hugging Face model loading for Continuous Intelligence.""" | |
| import json | |
| import math | |
| import re | |
| from contextlib import nullcontext | |
| from functools import cache | |
| from pathlib import Path, PureWindowsPath | |
| import torch | |
| from .ci_proto import Forecaster | |
| _COMMIT = re.compile(r"[0-9a-f]{40}") | |
| def _device(name: str) -> torch.device: | |
| if name == "auto": | |
| name = "cuda" if torch.cuda.is_available() else "cpu" | |
| if name not in ("cpu", "cuda"): | |
| raise ValueError("device must be 'cpu', 'cuda', or 'auto'") | |
| if name == "cuda" and not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA was requested but is not available") | |
| return torch.device(name) | |
| def _files( | |
| source: str | Path, tta_config: str | Path | None, revision: str | None | |
| ) -> tuple[Path, Path, Path | None]: | |
| path = Path(source).expanduser() | |
| if path.exists(): | |
| if revision is not None: | |
| raise ValueError("revision applies only to Hugging Face models") | |
| else: | |
| if ( | |
| isinstance(source, Path) | |
| or str(source).startswith(("/", "./", "../", "~", ".\\", "..\\")) | |
| or "\\" in str(source) | |
| or PureWindowsPath(source).drive | |
| or str(source).count("/") != 1 | |
| or path.parent.is_dir() | |
| or path.suffix in (".ckpt", ".pt", ".pth") | |
| or tta_config is not None | |
| ): | |
| raise FileNotFoundError(path) | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| info = Path(hf_hub_download(repo_id=str(source), filename="config.json", revision=revision)) | |
| commit = info.parent.name if _COMMIT.fullmatch(info.parent.name) else revision | |
| path = Path( | |
| snapshot_download( | |
| repo_id=str(source), | |
| revision=commit, | |
| allow_patterns=["model.ckpt", "model-long.ckpt", "tta.json"], | |
| ) | |
| ) | |
| info = None | |
| if path.is_dir(): | |
| if tta_config is not None: | |
| raise ValueError("a model bundle uses its own tta.json") | |
| info = path / "config.json" if (path / "config.json").is_file() else None | |
| path, tta_config = path / "model.ckpt", path / "tta.json" | |
| elif tta_config is None: | |
| raise ValueError("a standalone checkpoint requires tta_config") | |
| tta = Path(tta_config).expanduser() | |
| for file in (path, tta): | |
| if not file.is_file(): | |
| raise FileNotFoundError(file) | |
| return path, tta, info | |
| def _model_info(path: Path, forecaster: Forecaster) -> dict: | |
| """Read a bundle's ``config.json`` and check it describes the loaded weights.""" | |
| info = json.loads(path.read_text()) | |
| if not isinstance(info, dict): | |
| raise ValueError("config.json must contain a JSON object") | |
| missing = {"name", "parameters", "quantiles"} - set(info) | |
| if missing: | |
| raise ValueError(f"config.json is missing {sorted(missing)}") | |
| model = forecaster._model | |
| quantiles = forecaster.quantiles.tolist() | |
| if ( | |
| not isinstance(info["quantiles"], list) | |
| or len(info["quantiles"]) != len(quantiles) | |
| or not all( | |
| isinstance(given, (int, float)) and math.isclose(given, level, abs_tol=1e-6) | |
| for given, level in zip(info["quantiles"], quantiles, strict=True) | |
| ) | |
| ): | |
| raise ValueError("config.json quantiles do not match the checkpoint") | |
| if info["parameters"] != sum(parameter.numel() for parameter in model.parameters()): | |
| raise ValueError("config.json parameters do not match the checkpoint") | |
| return info | |
| def load_model( | |
| checkpoint: str | Path, | |
| tta_config: str | Path | None = None, | |
| device: str = "auto", | |
| *, | |
| revision: str | None = None, | |
| ) -> Forecaster: | |
| """Load a local model bundle, Hugging Face model ID, or checkpoint plus TTA file. | |
| Bundles contain ``model.ckpt`` (weights and embedded architecture) and | |
| ``tta.json``. Hub bundles also contain ``config.json``, which describes the | |
| model and is returned by ``Forecaster.info``; it is checked against the | |
| weights. Hub files are fetched together from the same revision; | |
| pass a commit SHA to pin it. Hub IDs use ``owner/model`` strings. Use a | |
| ``Path`` or ``./`` prefix to mark an ambiguous relative path as local. | |
| An optional ``model-long.ckpt`` shares the bundle's TTA and is selected | |
| when history exceeds the primary checkpoint's context length. | |
| """ | |
| placement = _device(device) | |
| path, tta_path, info_path = _files(checkpoint, tta_config, revision) | |
| payload = torch.load(path, map_location="cpu", weights_only=True, mmap=True) | |
| from .ci_proto._checkpoint import from_checkpoint | |
| forecaster = from_checkpoint(payload, tta_path, placement) | |
| placement = forecaster.device | |
| if info_path is not None: | |
| forecaster._info = _model_info(info_path, forecaster) | |
| long_path = path.with_name("model-long.ckpt") | |
| if tta_config is None and long_path.is_file() and forecaster._model.config.context_len: | |
| quantiles = forecaster.quantiles | |
| future_len = forecaster._model.config.future_len | |
| def load_long() -> Forecaster: | |
| context = torch.cuda.device(placement) if placement.type == "cuda" else nullcontext() | |
| with context: | |
| long = load_model(long_path, tta_path, device=placement.type) | |
| if not torch.equal(long.quantiles, quantiles) or ( | |
| long._model.config.future_len != future_len | |
| ): | |
| raise ValueError("long checkpoint must have the same quantiles and forecast window") | |
| return long | |
| forecaster._long_model = load_long | |
| return forecaster | |