Spaces:
Running on Zero
Running on Zero
| """Reconstruct ci-proto weights, calibration and inference backend.""" | |
| import json | |
| import math | |
| from pathlib import Path | |
| import torch | |
| from ._forecast import Forecaster | |
| from ._model import Model | |
| def _load_tta(path: str | Path) -> dict[str, float]: | |
| booleans = { | |
| "postprocess", | |
| "tta_diff", | |
| "tta_sign_flip", | |
| "tta_running_totals", | |
| "target_encode_categoricals", | |
| "clip_nonnegative", | |
| "integer_round", | |
| "bfloat16", | |
| "calendar_covariates", | |
| } | |
| constants = { | |
| "trend_threshold", | |
| "trend_min_r2", | |
| "trend_window", | |
| "diff_band_exponent", | |
| "diff_band_scale", | |
| "raw_band_exponent", | |
| "raw_band_scale", | |
| "differencing_max_horizon", | |
| "static_groups", | |
| "joint_neighbours", | |
| "chunk_joint", | |
| } | |
| required = booleans | constants | |
| raw = json.loads(Path(path).read_text()) | |
| if not isinstance(raw, dict): | |
| raise ValueError("tta config must contain a JSON object") | |
| missing = required - raw.keys() | |
| unknown = raw.keys() - required | |
| if missing or unknown: | |
| raise ValueError( | |
| f"invalid tta fields: missing={sorted(missing)}, unknown={sorted(unknown)}" | |
| ) | |
| result = {} | |
| for key in booleans: | |
| if type(raw[key]) is not bool: | |
| raise ValueError(f"{key} must be a boolean") | |
| result[key] = raw[key] | |
| for key in constants: | |
| value = raw[key] | |
| if key in ("differencing_max_horizon", "static_groups", "joint_neighbours", "chunk_joint"): | |
| valid = ( | |
| type(value) is int | |
| and value >= (1 if key == "differencing_max_horizon" else 0) | |
| and not (key in ("joint_neighbours", "chunk_joint") and value == 1) | |
| ) | |
| else: | |
| if type(value) not in (int, float) or not math.isfinite(value): | |
| raise ValueError(f"{key} must be a finite number") | |
| valid = value >= 0 if key.endswith("exponent") else value > 0 | |
| if key in ("trend_min_r2", "trend_window"): | |
| valid = 0 < value <= 1 | |
| value = float(value) | |
| if not valid: | |
| raise ValueError(f"invalid {key}: {value}") | |
| result[key] = value | |
| return result | |
| def from_checkpoint(payload: object, tta_config: str | Path, device: torch.device) -> Forecaster: | |
| """Build ci-proto from a loaded checkpoint and its calibration.""" | |
| tta = _load_tta(tta_config) | |
| metadata = payload.get("hyper_parameters") if isinstance(payload, dict) else None | |
| if not isinstance(metadata, dict) or not isinstance(metadata.get("model"), dict): | |
| raise ValueError("checkpoint must contain a hyper_parameters.model mapping") | |
| try: | |
| model = Model(metadata["model"]) | |
| except (KeyError, AttributeError, TypeError, ZeroDivisionError, OverflowError) as error: | |
| raise ValueError(f"invalid checkpoint model configuration: {error}") from error | |
| state = payload.get("state_dict") | |
| if not isinstance(state, dict) or not state or any(not isinstance(key, str) for key in state): | |
| raise ValueError("checkpoint must contain a nonempty state_dict with string keys") | |
| if all(key.startswith("model.") for key in state): | |
| state = {key.removeprefix("model."): value for key, value in state.items()} | |
| model.load_state_dict(state, strict=True) | |
| levels = model.quantiles | |
| if ( | |
| len(levels) % 2 == 0 | |
| or not torch.isclose(levels[len(levels) // 2], torch.tensor(0.5), atol=1e-6, rtol=0) | |
| or not torch.allclose( | |
| levels, | |
| 1 - levels.flip(0), | |
| atol=1e-6, | |
| rtol=0, | |
| ) | |
| ): | |
| raise ValueError("ci-proto TTA requires symmetric quantiles centred on the median") | |
| model = model.to(device).eval() | |
| from ._cuda import configure_backend | |
| configure_backend(model) | |
| return Forecaster(model, tta) | |