from __future__ import annotations import logging import os import time from dataclasses import dataclass from typing import Any, Iterator, List, Optional import httpx import numpy as np from gluonts.dataset import Dataset as GluonDataset from gluonts.model import Forecast from gluonts.model.forecast import QuantileForecast from gluonts.model.predictor import RepresentablePredictor from tsfm_bench.eval.predictors import _ForecastConfig DEFAULT_API_URL = "https://api.tsfm.ai/v1/forecast" DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] @dataclass class TsfmApiConfig: model_id: str api_url: str = DEFAULT_API_URL api_key: str | None = None timeout: float = 120.0 freq: str = "H" min_prediction_length: int = 1 class TsfmApiPredictor(RepresentablePredictor): """Zero-shot forecaster backed by the TSFM.ai hosted API.""" def __init__( self, config: TsfmApiConfig, prediction_length: int, quantile_levels: Optional[List[float]] = None, ): super().__init__(prediction_length=prediction_length) self.config = config self.quantile_levels = quantile_levels or DEFAULT_QUANTILES self.forecast_config = _ForecastConfig.from_quantiles(self.quantile_levels) self._api_key = config.api_key or os.getenv("TSFM_API_KEY") self.leaderboard_name = config.model_id.split("/")[-1] def predict(self, dataset: GluonDataset, **kwargs) -> Iterator[Forecast]: if not self._api_key: raise RuntimeError( "TSFM_API_KEY is required for online zero-shot evaluation. " "Get a key at https://tsfm.ai/" ) headers = { "Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json", } with httpx.Client(timeout=self.config.timeout) as client: dataset_freq = getattr(dataset, "freq", None) for entry in dataset: target = np.asarray(entry["target"], dtype=np.float64) univariate = target.ndim == 1 working = target.reshape(1, -1) if univariate else target curr_freq = dataset_freq if not curr_freq: start_period = entry.get("start") if start_period and hasattr(start_period, "freqstr"): curr_freq = start_period.freqstr api_freq = _resolve_api_freq(curr_freq) forecast_rows = [] for variate_idx in range(working.shape[0]): series = working[variate_idx] series = series[np.isfinite(series)] if len(series) < 32: pad_len = 32 - len(series) series = np.pad(series, (pad_len, 0), mode='edge') payload = { "model": self.config.model_id, "inputs": [ { "item_id": str(entry.get("item_id", variate_idx)), "target": _to_api_target(series), "start": _format_start(entry.get("start")), } ], "parameters": { "prediction_length": max(self.prediction_length, self.config.min_prediction_length), "freq": api_freq, "quantiles": self.quantile_levels, }, } # Robust retry loop to handle transient 503 / 502 / 504 / 429 errors from TSFM.ai API max_retries = 5 response = None for attempt in range(1, max_retries + 1): try: response = client.post( self.config.api_url, headers=headers, json=payload, ) response.raise_for_status() break except (httpx.HTTPStatusError, httpx.RequestError) as exc: status_code = getattr(response, "status_code", "network error") if attempt == max_retries: raise exc wait_time = 2 ** attempt logger = logging.getLogger(__name__) logger.warning( f"TSFM API request failed for model {self.config.model_id} " f"(status={status_code}): {exc}. " f"Retrying in {wait_time}s... (Attempt {attempt}/{max_retries})" ) time.sleep(wait_time) forecast_rows.append( _parse_api_forecast( response.json(), self.prediction_length, self.quantile_levels, ) ) stacked = np.stack(forecast_rows, axis=0) if univariate: forecast_arrays = stacked[0] else: forecast_arrays = stacked yield QuantileForecast( forecast_arrays=forecast_arrays, forecast_keys=self.forecast_config.forecast_keys, start_date=entry["start"] + len(entry["target"]), item_id=entry["item_id"], ) def _to_api_target(series: np.ndarray) -> list[list[float]]: """TSFM.ai expects target shaped [num_timesteps][num_channels].""" return [[float(v)] for v in series.tolist()] def _format_start(start: Any) -> str: if hasattr(start, "to_timestamp"): return start.to_timestamp().isoformat() return str(start) def _parse_api_forecast( payload: dict[str, Any], prediction_length: int, quantile_levels: list[float], ) -> np.ndarray: """Parse TSFM.ai response into (num_outputs, prediction_length).""" outputs_block = payload.get("outputs") or payload.get("forecasts") or [] if not outputs_block: raise ValueError(f"Unrecognized TSFM API response: {payload}") first = outputs_block[0] rows: list[np.ndarray] = [ _flatten_forecast_values(first.get("mean"), prediction_length)] quantile_map: dict[float, np.ndarray] = {} for item in first.get("quantile_predictions", []): level = float(item.get("level", item.get("quantile", 0.5))) quantile_map[level] = _flatten_forecast_values(item.get("values"), prediction_length) for q in quantile_levels: if q in quantile_map: rows.append(quantile_map[q]) else: nearest = min(quantile_map.keys(), key=lambda k: abs(k - q), default=None) if nearest is None: rows.append(rows[0].copy()) else: rows.append(quantile_map[nearest]) return np.stack(rows, axis=0).astype(np.float64) def _flatten_forecast_values(values: Any, prediction_length: int) -> np.ndarray: """Convert nested API values like [[1.2], [1.3], ...] to 1D array.""" if values is None: raise ValueError("Missing forecast values in TSFM API response") arr = np.asarray(values, dtype=np.float64) if arr.ndim == 2 and arr.shape[1] == 1: arr = arr[:, 0] elif arr.ndim > 1: arr = arr.reshape(arr.shape[0], -1).mean(axis=1) arr = arr.reshape(-1) if arr.size < prediction_length: raise ValueError(f"Forecast length {arr.size} < expected {prediction_length}") return arr[:prediction_length] def _resolve_api_freq(freq_str: str | None) -> str: if not freq_str: return "H" fs = str(freq_str).upper() if "MIN" in fs or "T" in fs: return "T" if "H" in fs: return "H" if "D" in fs: return "D" if "W" in fs: return "W" if "M" in fs: return "M" return "H"