| """Stage B: cache teacher quantile forecasts (TimesFM-2.5-200M + Chronos-2) for distillation. |
| |
| Run in the teacher venv (~/tvenv, system torch): |
| ~/tvenv/bin/python -m modular_ts.teacher_gen --source synth --n 200000 --out teach/synth |
| ~/tvenv/bin/python -m modular_ts.teacher_gen --source gep --n 100000 --out teach/gep --gep_selection gep_selection.json |
| |
| Each shard teach/<src>/shard_XXXX.npz holds |
| ctx (n, C) float32 context, NaN = unobserved (left-padded for short contexts) |
| fut (n, F) float32 ground-truth continuation (synthetic only; NaN for GEP = teacher-only input) |
| q_tfm, q_cr2 (n, H, 9) float32 teacher quantiles 0.1..0.9 for the first H future steps |
| Teacher weights: google/timesfm-2.5-200m-pytorch (Apache-2.0), amazon/chronos-2 (Apache-2.0). |
| GiftEvalPretrain inputs use the same GIFT-Eval/fev/TIME leakage filter as run1 (realdata.filter_datasets). |
| NEVER uses Timestral weights. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import time |
|
|
| import numpy as np |
| import torch |
|
|
| Q9 = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] |
|
|
|
|
| def interp_nans(x: np.ndarray) -> np.ndarray: |
| m = np.isfinite(x) |
| if m.all(): |
| return x |
| idx = np.arange(len(x)) |
| return np.interp(idx, idx[m], x[m]).astype(np.float32) |
|
|
|
|
| class Teachers: |
| def __init__(self, C, H, dev="cuda", tfm_bs=256, cr_bs=512): |
| import timesfm |
| from chronos import Chronos2Pipeline |
| torch.set_float32_matmul_precision("high") |
| self.H = H |
| self.cr = Chronos2Pipeline.from_pretrained("amazon/chronos-2", device_map=dev) |
| self.cr_bs = cr_bs |
| self.tfm = timesfm.TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch") |
| self.tfm.compile(timesfm.ForecastConfig(max_context=C, max_horizon=H, normalize_inputs=True, per_core_batch_size=tfm_bs, |
| use_continuous_quantile_head=True, force_flip_invariance=True, |
| infer_is_positive=True, fix_quantile_crossing=True)) |
|
|
| def __call__(self, ctx: np.ndarray): |
| """ctx (n, C) with NaN -> q_tfm, q_cr2 (n, H, 9).""" |
| seqs = [] |
| for row in ctx: |
| f = np.isfinite(row) |
| first = int(np.argmax(f)) if f.any() else len(row) - 1 |
| seqs.append(row[first:]) |
| with torch.no_grad(): |
| qs, _ = self.cr.predict_quantiles([torch.from_numpy(s) for s in seqs], prediction_length=self.H, |
| quantile_levels=Q9, batch_size=self.cr_bs) |
| q_cr2 = torch.stack([q[0] for q in qs]).float().cpu().numpy() |
| _, qt = self.tfm.forecast(horizon=self.H, inputs=[interp_nans(s) for s in seqs]) |
| q_tfm = np.asarray(qt)[:, :self.H, 1:10].astype(np.float32) |
| return q_tfm, q_cr2 |
|
|
|
|
| def synth_windows(gen, n, C, F, rng): |
| vals, roles = gen.group_batch(n, 1) |
| v = vals[:, 0].float().cpu().numpy() |
| T = C + F |
| |
| clen = np.where(rng.random(n) < 0.5, C, np.exp(rng.uniform(math.log(48), math.log(C), n))).astype(int) |
| obs = np.ones((n, T), bool) |
| obs[np.arange(T)[None] < (C - clen)[:, None]] = False |
| miss = rng.random(n) < 0.15 |
| pm = rng.uniform(0, 0.1, n)[:, None] |
| obs &= ~(miss[:, None] & (rng.random((n, T)) < pm)) |
| v = np.where(obs, v, np.nan).astype(np.float32) |
| return v[:, :C], v[:, C:] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--source", choices=["synth", "gep"], required=True) |
| ap.add_argument("--n", type=int, default=100000) |
| ap.add_argument("--shard", type=int, default=8192) |
| ap.add_argument("--C", type=int, default=2048) |
| ap.add_argument("--F", type=int, default=1024) |
| ap.add_argument("--H", type=int, default=512) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--seed", type=int, default=123) |
| ap.add_argument("--gep_selection", default="gep_selection.json") |
| ap.add_argument("--max_minutes", type=float, default=0) |
| args = ap.parse_args() |
| os.makedirs(args.out, exist_ok=True) |
| dev = torch.device("cuda") |
| rng = np.random.default_rng(args.seed) |
| T = args.C + args.F |
| if args.source == "synth": |
| from modular_ts.synth import SynthGen |
| gen = SynthGen(T, dev, seed=args.seed) |
| else: |
| from modular_ts.realdata import FastGiftEvalPretrain |
| gep = FastGiftEvalPretrain(T, args.shard * 2 * (T // 32), dev, json.load(open(args.gep_selection)), seed=args.seed, |
| mv_frac=0.0, short_ctx_p=0.3) |
| json.dump(gep.manifest(), open(os.path.join(args.out, "gep_manifest.json"), "w"), indent=1) |
| tea = Teachers(args.C, args.H) |
| t0, done, k = time.time(), 0, 0 |
| while done < args.n: |
| if args.max_minutes and (time.time() - t0) / 60 > args.max_minutes: |
| print("time limit", flush=True) |
| break |
| p = os.path.join(args.out, f"shard_{k:04d}.npz") |
| if os.path.exists(p): |
| done += args.shard; k += 1 |
| continue |
| if args.source == "synth": |
| gen = SynthGen(T, dev, seed=args.seed * 1000 + k) |
| rng = np.random.default_rng(args.seed * 1000 + k) |
| ctx, fut = synth_windows(gen, args.shard, args.C, args.F, rng) |
| else: |
| gep.g.manual_seed(args.seed * 1000 + k); gep.rng.seed(args.seed * 1000 + k) |
| b = gep.next_batch() |
| v = torch.where(b["observed"][:, 0], b["values"][:, 0], torch.full_like(b["values"][:, 0], float("nan"))).cpu().numpy() |
| ok = np.isfinite(v[:, :args.C]).sum(1) >= 48 |
| v = v[ok][:args.shard] |
| ctx, fut = v[:, :args.C], np.full((len(v), args.F), np.nan, np.float32) |
| ts = time.time() |
| q_tfm, q_cr2 = tea(ctx) |
| good = np.isfinite(q_tfm).all((1, 2)) & np.isfinite(q_cr2).all((1, 2)) |
| np.savez(p, ctx=ctx[good], fut=fut[good], q_tfm=q_tfm[good].astype(np.float32), q_cr2=q_cr2[good].astype(np.float32)) |
| done += int(good.sum()); k += 1 |
| print(json.dumps({"shard": k, "n": int(good.sum()), "dropped": int((~good).sum()), "teach_s": round(time.time() - ts, 1), |
| "done": done, "min": round((time.time() - t0) / 60, 2)}), flush=True) |
| print("TEACH_DONE", done, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
