Spaces:
Running
Running
File size: 8,236 Bytes
e317359 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | 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"
|