#!/usr/bin/env python3 """Run live TS-Bench + TSFM.ai evaluation inside the HF Space (or locally).""" from __future__ import annotations import csv import json import logging import os import sys from datetime import datetime, timezone from pathlib import Path from typing import Any import yaml SPACE_ROOT = Path(__file__).resolve().parents[1] if str(SPACE_ROOT) not in sys.path: sys.path.insert(0, str(SPACE_ROOT)) os.environ.setdefault("TS_BENCH_ROOT", str(SPACE_ROOT / "vendor" / "ts_bench")) from tsfm_bench.data.registry import load_data_source, load_dataset_properties from tsfm_bench.data.ts_bench import TsBenchDataSource from tsfm_bench.eval.api_predictor import DEFAULT_QUANTILES, TsfmApiConfig, TsfmApiPredictor from tsfm_bench.eval.online_eval import run_online_eval_for_dataset from src.eval_schedule import with_next_eval_fields from src.benchmark_config import BENCHMARK_INTERVAL_SECONDS logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger(__name__) DATA_CONFIG = SPACE_ROOT / "configs" / "datasets" / "ts_bench.yaml" MODEL_CONFIG = SPACE_ROOT / "configs" / "models" / "online_tsfm.yaml" OUTPUT_ROOT = Path(os.getenv("TSFM_RESULTS_PATH", str(SPACE_ROOT / "results"))) EVAL_STATE_PATH = OUTPUT_ROOT / "eval_state.json" def model_display_name(model_spec: dict[str, Any]) -> str: return model_spec.get("display_name") or model_spec["model_id"] def model_output_slug(model_spec: dict[str, Any]) -> str: import re name = model_spec.get("display_name") or model_spec["model_id"] return re.sub(r"[^a-zA-Z0-9]+", "_", name).strip("_").lower() def load_model_specs(path: Path) -> list[dict[str, Any]]: return yaml.safe_load(path.read_text()).get("models", []) def write_model_results( model_spec: dict[str, Any], rows: list[list[Any]], output_root: Path, meta: dict[str, Any], ) -> None: model_name = model_display_name(model_spec) out_dir = output_root / model_output_slug(model_spec) out_dir.mkdir(parents=True, exist_ok=True) csv_path = out_dir / "all_results.csv" with csv_path.open("w", newline="") as handle: writer = csv.writer(handle) writer.writerow( [ "dataset", "model", "eval_metrics/MSE[mean]", "eval_metrics/MSE[0.5]", "eval_metrics/MAE[0.5]", "eval_metrics/MASE[0.5]", "eval_metrics/MAPE[0.5]", "eval_metrics/sMAPE[0.5]", "eval_metrics/MSIS", "eval_metrics/RMSE[mean]", "eval_metrics/NRMSE[mean]", "eval_metrics/ND[0.5]", "eval_metrics/mean_weighted_sum_quantile_loss", "domain", "num_variates", ] ) writer.writerows(rows) config = { "model": model_name, "model_type": model_spec.get("model_type", "zero-shot"), "model_dtype": "float32", "model_link": model_spec.get( "model_link", f"https://tsfm.ai/models/{model_spec['model_id']}" ), "code_link": "https://github.com/zhouziyu02/TSFM_Bench/blob/main/space/scripts/run_space_eval.py", "org": model_spec.get("org", "TSFM.ai"), "testdata_leakage": "No", "replication_code_available": "Yes", "api_model_id": model_spec["model_id"], } (out_dir / "config.json").write_text(json.dumps(config, indent=4) + "\n") (out_dir / "online_meta.json").write_text(json.dumps(meta, indent=4) + "\n") def write_dataset_properties(properties: dict[str, dict[str, Any]], output_root: Path) -> None: rows = [ { "dataset": key, "domain": values["domain"], "frequency": values["frequency"], "num_variates": values["num_variates"], } for key, values in sorted(properties.items()) ] csv_path = output_root / "dataset_properties.csv" with csv_path.open("w", newline="") as handle: writer = csv.DictWriter( handle, fieldnames=["dataset", "domain", "frequency", "num_variates"], ) writer.writeheader() writer.writerows(rows) def write_eval_state(payload: dict[str, Any]) -> None: OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) EVAL_STATE_PATH.write_text(json.dumps(payload, indent=4) + "\n") def write_run_metadata(output_root: Path, payload: dict[str, Any]) -> None: (output_root / "online_status.json").write_text(json.dumps(payload, indent=4) + "\n") def make_predictor(model_spec: dict[str, Any], pred_len: int) -> TsfmApiPredictor: predictor = TsfmApiPredictor( config=TsfmApiConfig(model_id=model_spec["model_id"]), prediction_length=pred_len, quantile_levels=DEFAULT_QUANTILES, ) predictor.leaderboard_name = model_display_name(model_spec) return predictor def run_evaluation() -> int: if not os.getenv("TSFM_API_KEY"): write_eval_state( { "status": "disabled", "message": "TSFM_API_KEY not configured — set it in HF Space secrets.", "updated_at": datetime.now(timezone.utc).isoformat(), } ) logger.error("TSFM_API_KEY is missing; skipping evaluation") return 1 started = datetime.now(timezone.utc).isoformat() write_eval_state( { "status": "running", "started_at": started, "message": "Collecting TS-Bench data and calling TSFM.ai API…", "updated_at": started, } ) try: source = load_data_source(DATA_CONFIG) if isinstance(source, TsBenchDataSource) and source._settings.auto_refresh: source.refresh() datasets = source.list_datasets() if not datasets: raise RuntimeError("No TS-Bench tasks available after data collection") properties = load_dataset_properties(DATA_CONFIG) model_specs = load_model_specs(MODEL_CONFIG) OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) write_dataset_properties(properties, OUTPUT_ROOT) failed_models: list[str] = [] all_model_meta: dict[str, Any] = {} for model_spec in model_specs: model_name = model_display_name(model_spec) rows: list[list[Any]] = [] dataset_meta: list[dict[str, Any]] = [] try: for ds_name in datasets: pred_len = source.get_prediction_length(ds_name) predictor = make_predictor(model_spec, pred_len) try: logger.info("Evaluating %s on %s", model_name, ds_name) result = run_online_eval_for_dataset(source, ds_name, predictor) except Exception: logger.exception("Skipping %s on %s", model_name, ds_name) continue rows.append( [ result.dataset, model_name, result.metrics["MSE[mean]"], result.metrics["MSE[0.5]"], result.metrics["MAE[0.5]"], result.metrics["MASE[0.5]"], result.metrics["MAPE[0.5]"], result.metrics["sMAPE[0.5]"], result.metrics["MSIS"], result.metrics["RMSE[mean]"], result.metrics["NRMSE[mean]"], result.metrics["ND[0.5]"], result.metrics["mean_weighted_sum_quantile_loss"], result.domain, result.num_variates, ] ) dataset_meta.append( { "dataset": ds_name, "data_fetched_at": result.data_fetched_at, "context_length": result.context_length, "prediction_length": result.prediction_length, } ) if not rows: raise RuntimeError(f"No successful datasets for {model_name}") model_meta = { "model": model_name, "api_model_id": model_spec["model_id"], "evaluated_at": datetime.now(timezone.utc).isoformat(), "datasets": dataset_meta, } write_model_results(model_spec, rows, OUTPUT_ROOT, model_meta) all_model_meta[model_name] = model_meta except Exception: logger.exception("Failed evaluating %s", model_name) failed_models.append(model_name) finished = datetime.now(timezone.utc).isoformat() status = "ok" if not failed_models else "partial" write_run_metadata( OUTPUT_ROOT, { "status": status, "started_at": started, "finished_at": finished, "data_source": "ts_bench", "data_config": str(DATA_CONFIG), "ts_bench_root": os.environ.get("TS_BENCH_ROOT", ""), "models": all_model_meta, "failed_models": failed_models, }, ) write_eval_state( with_next_eval_fields( { "status": status, "started_at": started, "finished_at": finished, "failed_models": failed_models, "task_count": len(datasets), "model_count": len(model_specs) - len(failed_models), "message": "Evaluation complete", "updated_at": finished, }, finished, BENCHMARK_INTERVAL_SECONDS, ) ) return 0 if status == "ok" else 2 except Exception as exc: finished = datetime.now(timezone.utc).isoformat() logger.exception("Evaluation run failed") write_eval_state( with_next_eval_fields( { "status": "error", "started_at": started, "finished_at": finished, "message": str(exc), "updated_at": finished, }, finished, BENCHMARK_INTERVAL_SECONDS, ) ) return 1 if __name__ == "__main__": raise SystemExit(run_evaluation())