Spaces:
Running on Zero
Running on Zero
File size: 3,915 Bytes
ba6c090 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """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)
|