modular_ts/eval_gift.py

"""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()