File size: 6,449 Bytes
51bb0b7 | 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 | """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
# context length: 50% full C, else log-uniform in [48, C]; 15% of series get iid missing values
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) # per-shard seed -> resumable without duplicates
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()
|