File size: 8,987 Bytes
b428368 | 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """Auto-generated by cascade Toto2Trainer. Loads the trained checkpoint and
decodes the full horizon in one forward pass via contiguous patch masking
(CPM) — no autoregressive sampling. Exposes:
forecast(history, horizon, num_samples) -> (1, num_samples, horizon)
the cascade validator contract — sample paths drawn once from the
decoded quantiles (seeded per window for validator consensus).
forecast_quantiles(history, horizon) -> (1, horizon, num_q)
forecast_quantiles_batch(histories, horizon) -> (B, horizon, num_q)
the quantile head directly — what benchmark CRPS consumes; batched
across series so eval sweeps amortize the forward passes.
"""
from __future__ import annotations
import hashlib
import importlib.util
import json
import sys
from pathlib import Path
import numpy as np
import torch
# Single-pass CPM decoding is stable to ~768 steps (Toto 2.0 tech report);
# longer horizons block-decode: commit the median per block, then continue.
STABLE_DECODE_STEPS = 768
def _load_model_module(d: Path):
spec = importlib.util.spec_from_file_location("cascade_ckpt_model", d / "model.py")
mod = importlib.util.module_from_spec(spec)
# Register before exec: model.py defines an @dataclass, and the dataclass
# machinery does sys.modules.get(cls.__module__).__dict__ during class
# creation — which is None (AttributeError) unless the module is registered.
sys.modules[spec.name] = mod
spec.loader.exec_module(mod)
return mod
class Wrapper:
def __init__(self, checkpoint_dir, device: str = "cpu"):
d = Path(checkpoint_dir)
self.device = device
cfg_obj = json.loads((d / "config.json").read_text())
self.m = _load_model_module(d)
self.cfg = self.m.Toto2Config(**cfg_obj["toto2"])
self.quantile_levels = [float(v) for v in cfg_obj["quantile_levels"]]
self.levels = torch.tensor(self.quantile_levels, dtype=torch.float32, device=device)
self.model = self.m.Toto2Model(self.cfg).to(device).eval()
from safetensors.torch import load_file
state = load_file(str(d / "weights.safetensors"))
self.model.load_state_dict(state)
# ── CPM decoding ──────────────────────────────────────────────────────────
def _prep(self, histories):
"""Left-pad (with the first value) or truncate each 1-D history to the
context window. Returns the real-space context ``(B, window_len)`` in
float64 — standardization happens per decode block, from full
precision, so large-level series keep their fluctuations."""
ps = self.cfg.patch_size
n_ctx = max(2, self.cfg.context_length // ps)
window_len = n_ctx * ps
rows = []
for h in histories:
h = np.asarray(h, dtype=np.float64).reshape(-1)
if h.shape[0] < window_len:
pad = np.full(window_len - h.shape[0], h[0] if h.size else 0.0)
h = np.concatenate([pad, h])
else:
h = h[-window_len:]
rows.append(h)
return torch.as_tensor(np.stack(rows), dtype=torch.float64, device=self.device)
@torch.no_grad()
def _decode_block_z(self, z, block: int):
"""One CPM forward pass: append ``block`` masked patches to the
normalized context ``(B, L)`` and read their z-space quantiles
``(B, block*patch_size, num_q)``."""
ps = self.cfg.patch_size
# keep as much context as the positional table allows
ctx_p = min(z.shape[1] // ps, self.cfg.max_patches - block)
ctx = z[:, -ctx_p * ps :].view(z.shape[0], ctx_p, ps)
filler = torch.zeros(z.shape[0], block, ps, dtype=ctx.dtype, device=self.device)
mask = torch.zeros(z.shape[0], ctx_p + block, dtype=ctx.dtype, device=self.device)
mask[:, ctx_p:] = 1.0
pred = self.model(torch.cat([ctx, filler], dim=1), mask=mask)
# position i predicts patch i+1 → the horizon patches come from
# positions ctx_p-1 .. ctx_p+block-2.
q = pred[:, ctx_p - 1 : ctx_p + block - 1] # (B, block, ps, nq)
q, _ = torch.sort(q, dim=-1) # prevent quantile crossing
return q.reshape(z.shape[0], block * ps, -1)
@torch.no_grad()
def _decode_quantiles(self, x, horizon: int):
"""Block-decode real-space quantiles ``(B, horizon, num_q)`` from the
real-space context ``x`` ``(B, L)``.
Each block re-runs the causal scaler over history + committed medians
and unscales with the resulting end-of-context anchor. Committed
patches are *observed* context for later blocks, and in training the
causal stats advance through every observed patch — so the anchor must
advance with them; reusing the pre-horizon anchor would feed blocks ≥ 2
a scale/location regime the model never sees in training. Clamp bounds
are fixed from the original context (min/max ± 1e4x anchor scale, per
the report) so committed medians can't widen them.
"""
ps = self.cfg.patch_size
stable = max(1, min(STABLE_DECODE_STEPS // ps, self.cfg.max_patches - 2))
remaining = -(-int(horizon) // ps)
lo = hi = None
out = []
while remaining > 0:
block = min(remaining, stable)
z, loc_t, scale_t = self.m.causal_standardize(x)
loc = loc_t[:, -1:].double().unsqueeze(-1) # (B, 1, 1)
scale = scale_t[:, -1:].double().unsqueeze(-1)
if lo is None:
lo = x.min(dim=-1, keepdim=True).values.unsqueeze(-1) - 1e4 * scale
hi = x.max(dim=-1, keepdim=True).values.unsqueeze(-1) + 1e4 * scale
qz = self._decode_block_z(z.to(torch.float32), block)
q = torch.sinh(qz.double()) * scale + loc # (B, block*ps, nq)
q = torch.clamp(q, min=lo, max=hi)
out.append(q)
remaining -= block
if remaining > 0:
x = torch.cat([x, q[..., q.shape[-1] // 2]], dim=1)
return torch.cat(out, dim=1)[:, : int(horizon)]
# ── quantile head (benchmark path) ────────────────────────────────────────
@torch.no_grad()
def forecast_quantiles_batch(self, histories, horizon: int) -> np.ndarray:
"""Decode ``len(histories)`` series in one batch → real-space quantiles
``(B, horizon, num_q)`` at ``self.quantile_levels``. arcsinh + affine
are monotone increasing, so quantiles map pointwise."""
q = self._decode_quantiles(self._prep(list(histories)), horizon)
return q.detach().cpu().numpy().astype(np.float64)
def forecast_quantiles(self, history, horizon: int) -> np.ndarray:
return self.forecast_quantiles_batch([history], horizon)
# ── validator contract (sample paths) ─────────────────────────────────────
@torch.no_grad()
def forecast(self, history, horizon: int, num_samples: int) -> np.ndarray:
hist = np.asarray(history, dtype=np.float64).reshape(-1)
# Deterministic per-window sampling: seed from the (raw history, horizon,
# num_samples) so every validator computes identical scores and king vs
# challenger share the uniform draws (paired Monte-Carlo).
seed_src = hist.tobytes() + int(horizon).to_bytes(8, "big") + int(num_samples).to_bytes(8, "big")
seed = int.from_bytes(hashlib.sha256(seed_src).digest()[:8], "big") & ((1 << 63) - 1)
generator = torch.Generator(device=self.device)
generator.manual_seed(seed)
q = self._decode_quantiles(self._prep([hist]), horizon)[0] # (h, nq) real-space
# One draw per step per path via the piecewise-linear inverse CDF of the
# decoded quantiles (already clamped and monotone in the level).
# Quantiles decode once; samples never feed back.
nq = q.shape[-1]
levels = self.levels
u = torch.rand(int(num_samples), int(horizon), device=self.device, generator=generator)
idx = torch.searchsorted(levels, u.clamp(levels[0].item(), levels[-1].item()))
idx = idx.clamp(1, nq - 1)
i_lo = idx - 1
i_hi = idx
qe = q.unsqueeze(0).expand(u.shape[0], -1, -1) # (ns, h, nq)
vl = torch.gather(qe, -1, i_lo.unsqueeze(-1)).squeeze(-1)
vh = torch.gather(qe, -1, i_hi.unsqueeze(-1)).squeeze(-1)
ql = levels[i_lo].double(); qh = levels[i_hi].double()
frac = ((u.double() - ql) / (qh - ql).clamp_min(1e-8)).clamp(0, 1)
out = vl + frac * (vh - vl) # (ns, h)
return out.detach().cpu().numpy().reshape(1, int(num_samples), int(horizon))
|