| """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 |
|
|
| 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 = [("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(): |
| 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: |
| 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() |
|
|
