"""GIFT-Eval harness for Modular-TS (official gift-eval Dataset + gluonts evaluate_model). Adapted from gift-eval/notebooks/timestral.ipynb (Apache-2.0). Our weights load into the Apache-2.0 `timestral.TimestralForecaster` (block-AR decoding, advancing refine frame, clampsym TTA). Metrics, settings and the leaderboard aggregate match the official board: evaluate_model(axis=None, mask_invalid_label=True, allow_nan_forecast=False), 9 quantiles 0.1..0.9 for CRPS (mean_weighted_sum_quantile_loss), then the geometric mean over configs of (model / Seasonal_Naive) for MASE and CRPS. python -m modular_ts.eval_gift --export runs/smoke/export --mode quick --out runs/smoke/gift_quick python -m modular_ts.eval_gift --export ... --mode full # all 97 configs -> all_results.csv Requires: pip install --no-deps git+https://github.com/SalesforceAIResearch/gift-eval pip install "gluonts~=0.15.1" python-dotenv toolz orjson datasets pyarrow pandas Data: Salesforce/GiftEval is downloaded (only the needed folders) to $GIFT_EVAL. """ from __future__ import annotations import argparse import json import logging import math import os import time import warnings from pathlib import Path import numpy as np import pandas as pd import torch import modular_ts.model # noqa: F401 (puts the vendored timestral package on sys.path) ROOT = Path(__file__).resolve().parent.parent LB_DIR = ROOT / "gift-eval" / "results" PROPS = ROOT / "gift-eval" / "notebooks" / "dataset_properties.json" SHORT_DATASETS = "m4_yearly m4_quarterly m4_monthly m4_weekly m4_daily m4_hourly electricity/15T electricity/H electricity/D electricity/W solar/10T solar/H solar/D solar/W hospital covid_deaths us_births/D us_births/M us_births/W saugeenday/D saugeenday/M saugeenday/W temperature_rain_with_missing kdd_cup_2018_with_missing/H kdd_cup_2018_with_missing/D car_parts_with_missing restaurant hierarchical_sales/D hierarchical_sales/W LOOP_SEATTLE/5T LOOP_SEATTLE/H LOOP_SEATTLE/D SZ_TAXI/15T SZ_TAXI/H M_DENSE/H M_DENSE/D ett1/15T ett1/H ett1/D ett1/W ett2/15T ett2/H ett2/D ett2/W jena_weather/10T jena_weather/H jena_weather/D bitbrains_fast_storage/5T bitbrains_fast_storage/H bitbrains_rnd/5T bitbrains_rnd/H bizitobs_application bizitobs_service bizitobs_l2c/5T bizitobs_l2c/H" MED_LONG_DATASETS = "electricity/15T electricity/H solar/10T solar/H kdd_cup_2018_with_missing/H LOOP_SEATTLE/5T LOOP_SEATTLE/H SZ_TAXI/15T M_DENSE/H ett1/15T ett1/H ett2/15T ett2/H jena_weather/10T jena_weather/H bitbrains_fast_storage/5T bitbrains_rnd/5T bizitobs_application bizitobs_service bizitobs_l2c/5T bizitobs_l2c/H" PRETTY = {"saugeenday": "saugeen", "temperature_rain_with_missing": "temperature_rain", "kdd_cup_2018_with_missing": "kdd_cup_2018", "car_parts_with_missing": "car_parts"} # Quick subset: small, domain-diverse configs (12 short + 2 long). Used for smoke checks only. QUICK = [("m4_hourly", "short"), ("m4_weekly", "short"), ("us_births/D", "short"), ("saugeenday/D", "short"), ("hospital", "short"), ("covid_deaths", "short"), ("ett1/H", "short"), ("jena_weather/H", "short"), ("bizitobs_service", "short"), ("car_parts_with_missing", "short"), ("electricity/D", "short"), ("LOOP_SEATTLE/D", "short"), ("ett2/H", "long"), ("bizitobs_l2c/H", "long")] QUANTILE_LEVELS = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] METRIC_COLUMNS = ["MSE[mean]", "MSE[0.5]", "MAE[0.5]", "MASE[0.5]", "MAPE[0.5]", "sMAPE[0.5]", "MSIS", "RMSE[mean]", "NRMSE[mean]", "ND[0.5]", "mean_weighted_sum_quantile_loss"] COMPARE = ["Seasonal_Naive", "TimeSSM-2.5M", "Timestral-5M", "TTM-R3-PT", "TinyCast"] def all_configs(properties): short, ml = SHORT_DATASETS.split(), MED_LONG_DATASETS.split() out = [] for name in dict.fromkeys(short + ml): for term in ["short", "medium", "long"]: if term != "short" and name not in ml: continue out.append((name, term, config_key(name, term, properties))) assert len(out) == 97 return out def config_key(name, term, properties): if "/" in name: k, f = name.split("/") k = PRETTY.get(k.lower(), k.lower()) else: k = PRETTY.get(name.lower(), name.lower()) f = properties[k]["frequency"] return f"{k}/{f}/{term}" def ensure_data(names, root): from huggingface_hub import snapshot_download tops = sorted({n.split("/")[0] for n in names}) need = [t for t in tops if not (Path(root) / t).exists()] if need: snapshot_download("Salesforce/GiftEval", repo_type="dataset", local_dir=root, allow_patterns=[f"{t}/**" for t in need]) class Predictor: def __init__(self, fc, prediction_length, use_past_covariates, max_rows=1024, max_tokens=2**24): self.fc, self.h, self.cov = fc, prediction_length, use_past_covariates self.max_rows, self.max_tokens = max_rows, max_tokens lv = list(fc.quantile_levels) self.qi = [next(i for i, l in enumerate(lv) if abs(l - q) < 1e-9) for q in QUANTILE_LEVELS] self.keys = [str(q) for q in QUANTILE_LEVELS] + ["mean"] def predict(self, test_input): from gluonts.model.forecast import QuantileForecast entries = list(test_input) mc = self.fc.max_context tg, cv = [], [] for e in entries: t = np.asarray(e["target"], dtype=np.float32) tg.append((t[None] if t.ndim == 1 else t)[:, -mc:]) c = e.get("past_feat_dynamic_real") if self.cov else None if c is not None: c = np.asarray(c, dtype=np.float32) c = (c[None] if c.ndim == 1 else c)[:, -mc:] cv.append(c) nrows = tg[0].shape[0] + (cv[0].shape[0] if cv[0] is not None else 0) outs = [None] * len(entries) order = sorted(range(len(entries)), key=lambda i: tg[i].shape[-1]) s = 0 while s < len(order): e_ = s + 1 while e_ < len(order): L = tg[order[e_]].shape[-1] + self.h if (e_ + 1 - s) * nrows > self.max_rows or (e_ + 1 - s) * nrows * L > self.max_tokens: break e_ += 1 self._batch(order[s:e_], tg, cv, outs) s = e_ res = [] for e, q in zip(entries, outs): arr = np.concatenate([q, q[..., QUANTILE_LEVELS.index(0.5), None]], -1).transpose(2, 1, 0) if np.asarray(e["target"]).ndim == 1: arr = arr[..., 0] res.append(QuantileForecast(forecast_arrays=arr, forecast_keys=self.keys, start_date=e["start"] + np.asarray(e["target"]).shape[-1], item_id=e.get("item_id"))) return res def _batch(self, idx, tg, cv, outs): try: L = max(tg[i].shape[-1] for i in idx) ctx = np.stack([_lpad(tg[i], L) for i in idx]) past = np.stack([_lpad(cv[i], L) for i in idx]) if cv[idx[0]] is not None else None f = self.fc.forecast(ctx, self.h, past_covariates=past) q = f.quantiles[..., self.qi].float().cpu().numpy() except torch.cuda.OutOfMemoryError: torch.cuda.empty_cache() if len(idx) == 1: raise hlf = len(idx) // 2 self._batch(idx[:hlf], tg, cv, outs) self._batch(idx[hlf:], tg, cv, outs) return for j, i in enumerate(idx): outs[i] = q[j] def _lpad(x, L): p = L - x.shape[-1] return x if p == 0 else np.concatenate([np.full((x.shape[0], p), np.nan, np.float32), x], -1) def evaluate_config(fc, name, term): from gift_eval.data import Dataset from gluonts.ev.metrics import MAE, MAPE, MASE, MSE, MSIS, ND, NRMSE, RMSE, SMAPE, MeanWeightedSumQuantileLoss from gluonts.model import evaluate_model from gluonts.time_feature import get_seasonality metrics = [MSE(forecast_type="mean"), MSE(forecast_type=0.5), MAE(), MASE(), MAPE(), SMAPE(), MSIS(), RMSE(), NRMSE(), ND(), MeanWeightedSumQuantileLoss(quantile_levels=QUANTILE_LEVELS)] ds = Dataset(name=name, term=term, to_univariate=False) pred = Predictor(fc, ds.prediction_length, bool(ds.past_feat_dynamic_real_dim), max_rows=int(os.environ.get("MAX_ROWS", "1024"))) res = evaluate_model(pred, test_data=ds.test_data, metrics=metrics, batch_size=512, axis=None, mask_invalid_label=True, allow_nan_forecast=False, seasonality=get_seasonality(ds.freq)) return {k: float(res[k].iloc[0]) for k in res.columns} def gmean(x): return float(math.exp(np.log(np.asarray(x, dtype=float)).mean())) def aggregate(per_cfg: dict[str, dict]) -> dict: keys = sorted(per_cfg) sn = pd.read_csv(LB_DIR / "Seasonal_Naive" / "all_results.csv").set_index("dataset") out = {} m = [per_cfg[k]["MASE[0.5]"] / sn.loc[k, "eval_metrics/MASE[0.5]"] for k in keys] c = [per_cfg[k]["mean_weighted_sum_quantile_loss"] / sn.loc[k, "eval_metrics/mean_weighted_sum_quantile_loss"] for k in keys] out["ours"] = {"MASE": gmean(m), "CRPS": gmean(c)} for name in COMPARE: p = LB_DIR / name / "all_results.csv" if not p.exists(): continue df = pd.read_csv(p).set_index("dataset") if not all(k in df.index for k in keys): continue out[name] = { "MASE": gmean([df.loc[k, "eval_metrics/MASE[0.5]"] / sn.loc[k, "eval_metrics/MASE[0.5]"] for k in keys]), "CRPS": gmean([df.loc[k, "eval_metrics/mean_weighted_sum_quantile_loss"] / sn.loc[k, "eval_metrics/mean_weighted_sum_quantile_loss"] for k in keys]), } out["n_configs"] = len(keys) return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--export", required=True) ap.add_argument("--mode", default="quick", choices=["quick", "full"]) ap.add_argument("--configs", default="", help="override: 'name:term,name:term'") ap.add_argument("--out", required=True) ap.add_argument("--gift_eval_dir", default=os.environ.get("GIFT_EVAL", str(ROOT / "data" / "gift_eval"))) ap.add_argument("--max_context", type=int, default=0, help="0 = model's trained context_length") ap.add_argument("--model_name", default="Modular-TS-2.5M") ap.add_argument("--tta", default="", help="v3 test-time ensemble spec, e.g. 'c8192,c2048,d4' (modular_ts.tta)") ap.add_argument("--tta_alpha", type=float, default=0.5) ap.add_argument("--tta_backtest", action="store_true") ap.add_argument("--tta_tau", type=float, default=0.5) ap.add_argument("--tta_bt_windows", type=int, default=1) ap.add_argument("--tta_prior", type=float, default=0.0) ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") args = ap.parse_args() logging.getLogger("gluonts.model.forecast").setLevel(logging.ERROR) warnings.filterwarnings("ignore") os.environ["GIFT_EVAL"] = args.gift_eval_dir if os.environ.get("EVAL_MEM_FRAC") and torch.cuda.is_available(): # share the GPU with a training job torch.cuda.set_per_process_memory_fraction(float(os.environ["EVAL_MEM_FRAC"])) props = json.load(open(PROPS)) if args.configs: cfgs = [tuple(x.split(":")) for x in args.configs.split(",")] elif args.mode == "quick": cfgs = QUICK else: cfgs = [(n, t) for n, t, _ in all_configs(props)] ensure_data([n for n, _ in cfgs], args.gift_eval_dir) from timestral import TimestralForecaster fc = TimestralForecaster.from_pretrained(args.export, device=args.device, max_context=args.max_context or None) if args.tta: from modular_ts.tta import EnsembleForecaster fc = EnsembleForecaster(fc, args.tta, alpha=args.tta_alpha, backtest=args.tta_backtest, tau=args.tta_tau, bt_windows=args.tta_bt_windows, prior=args.tta_prior) out = Path(args.out) (out / "per_config").mkdir(parents=True, exist_ok=True) per = {} for name, term in cfgs: key = config_key(name, term, props) p = out / "per_config" / (key.replace("/", "__") + ".json") if p.exists(): per[key] = json.loads(p.read_text())["metrics"] continue t = time.time() try: met = evaluate_config(fc, name, term) except Exception as e: # keep going; a failed config is reported, not silently dropped print(f"{key}: Error {e!r}"[:400], flush=True) continue p.write_text(json.dumps({"dataset": key, "metrics": met, "seconds": time.time() - t}, indent=1)) per[key] = met print(f"{key}: MASE={met['MASE[0.5]']:.4f} CRPS={met['mean_weighted_sum_quantile_loss']:.4f} ({time.time()-t:.0f}s)", flush=True) agg = aggregate(per) json.dump(agg, open(out / "summary.json", "w"), indent=1) print("SUMMARY", json.dumps(agg), flush=True) if args.mode == "full": rows = [] for k, met in sorted(per.items()): ds_key = k.split("/")[0] rows.append({"dataset": k, "model": args.model_name, **{f"eval_metrics/{c}": met[c] for c in METRIC_COLUMNS}, "domain": props[ds_key]["domain"], "num_variates": props[ds_key]["num_variates"]}) pd.DataFrame(rows).to_csv(out / "all_results.csv", index=False) if __name__ == "__main__": main()