"""Model wrappers for TimesFM and Chronos-2 with graceful degradation. Both wrappers expose a uniform interface:: model.available -> bool model.label -> str (e.g. "TimesFM 2.5", "Chronos-2") model.forecast(series, horizon, covariates=None) -> np.ndarray If the heavy ML dependencies or the model weights are unavailable (not installed, offline, download failed, OOM, ...) the wrapper keeps ``available = False`` and ``forecast`` returns a damped naive forecast. That guarantees ``main.py`` always starts and ``/forecast`` always responds, which is exactly what Phase 5 asks for. When the real weights are present the wrappers use them. The API surfaces of these libraries changed across releases, so each loader tries the current (2.x) API first and falls back to the legacy API. Every inference call is wrapped so a runtime surprise degrades to the naive path instead of 500-ing the request. """ from __future__ import annotations import logging from typing import Optional, Sequence import numpy as np log = logging.getLogger("forecast.models") def naive_forecast(series: Sequence[float], horizon: int) -> np.ndarray: """Damped-drift fallback forecast. Continues the last observed step, damping it each period so the line does not run away over the horizon. Falls back to persistence (repeat the last value) when there is too little history. """ arr = np.asarray(series, dtype=float) arr = arr[np.isfinite(arr)] if arr.size == 0: return np.zeros(int(horizon), dtype=float) last = float(arr[-1]) step = float(arr[-1] - arr[-2]) if arr.size >= 2 else 0.0 damp = 0.6 out = [] cur = last for _ in range(int(horizon)): cur = cur + step step *= damp out.append(cur) return np.asarray(out, dtype=float) class TimesFMModel: """Google TimesFM wrapper (tries the 2.5 API, then the legacy 2.0 API).""" def __init__(self) -> None: self.available = False self.label = "TimesFM 2.0" self._impl = None self._mode: Optional[str] = None self._load() def _load(self) -> None: try: import timesfm # type: ignore except Exception as exc: # pragma: no cover - import guard log.warning("timesfm not importable (%s) — using naive fallback", exc) return # --- Preferred: TimesFM 2.5 (current PyPI API) --------------------- try: import torch # type: ignore torch.set_float32_matmul_precision("high") model = timesfm.TimesFM_2p5_200M_torch.from_pretrained( "google/timesfm-2.5-200m-pytorch" ) model.compile( timesfm.ForecastConfig( max_context=1024, max_horizon=256, normalize_inputs=True, use_continuous_quantile_head=True, force_flip_invariance=True, infer_is_positive=False, # probabilities can sit near 0 fix_quantile_crossing=True, ) ) self._impl = model self._mode = "2.5" self.label = "TimesFM 2.5" self.available = True log.info("Loaded TimesFM 2.5 (google/timesfm-2.5-200m-pytorch)") return except Exception as exc: log.warning("TimesFM 2.5 load failed (%s) — trying legacy 2.0 API", exc) # --- Fallback: TimesFM 2.0 (legacy hparams/checkpoint API) --------- try: tfm = timesfm.TimesFm( hparams=timesfm.TimesFmHparams( backend="cpu", per_core_batch_size=32, horizon_len=128, context_len=2048, ), checkpoint=timesfm.TimesFmCheckpoint( huggingface_repo_id="google/timesfm-2.0-500m-pytorch" ), ) self._impl = tfm self._mode = "2.0" self.label = "TimesFM 2.0" self.available = True log.info("Loaded TimesFM 2.0 (google/timesfm-2.0-500m-pytorch)") except Exception as exc: log.warning("TimesFM 2.0 load failed (%s) — using naive fallback", exc) def forecast( self, series: Sequence[float], horizon: int, covariates: Optional[Sequence[Sequence[float]]] = None, ) -> np.ndarray: if not self.available or self._impl is None: return naive_forecast(series, horizon) arr = np.asarray(series, dtype=float) try: if self._mode == "2.5": point, _quantiles = self._impl.forecast(horizon=int(horizon), inputs=[arr]) return np.asarray(point[0], dtype=float)[:horizon] # legacy 2.0 point, _ = self._impl.forecast([arr], freq=[0]) return np.asarray(point[0], dtype=float)[:horizon] except Exception as exc: log.warning("TimesFM forecast failed (%s) — naive fallback", exc) return naive_forecast(series, horizon) class ChronosModel: """Amazon Chronos-2 wrapper (falls back to a Chronos-Bolt pipeline).""" def __init__(self) -> None: self.available = False self.label = "Chronos-2" self._impl = None self._mode: Optional[str] = None self._load() def _load(self) -> None: # --- Preferred: Chronos-2 ----------------------------------------- try: from chronos import Chronos2Pipeline # type: ignore self._impl = Chronos2Pipeline.from_pretrained( "amazon/chronos-2", device_map="cpu" ) self._mode = "chronos2" self.label = "Chronos-2" self.available = True log.info("Loaded Chronos-2 (amazon/chronos-2)") return except Exception as exc: log.warning("Chronos-2 load failed (%s) — trying Chronos-Bolt", exc) # --- Fallback: Chronos-Bolt (older base pipeline) ----------------- try: from chronos import BaseChronosPipeline # type: ignore self._impl = BaseChronosPipeline.from_pretrained( "amazon/chronos-bolt-base", device_map="cpu" ) self._mode = "bolt" self.label = "Chronos-Bolt" self.available = True log.info("Loaded Chronos-Bolt (amazon/chronos-bolt-base)") except Exception as exc: log.warning("Chronos-Bolt load failed (%s) — using naive fallback", exc) def forecast( self, series: Sequence[float], horizon: int, covariates: Optional[Sequence[Sequence[float]]] = None, ) -> np.ndarray: if not self.available or self._impl is None: return naive_forecast(series, horizon) arr = np.asarray(series, dtype=float) try: if self._mode == "chronos2": return self._forecast_chronos2(arr, horizon, covariates) return self._forecast_bolt(arr, horizon) except Exception as exc: log.warning("Chronos forecast failed (%s) — naive fallback", exc) return naive_forecast(series, horizon) def _forecast_chronos2( self, arr: np.ndarray, horizon: int, covariates: Optional[Sequence[Sequence[float]]], ) -> np.ndarray: import pandas as pd ts = pd.date_range("2000-01-01", periods=arr.size, freq="D") frame = {"id": ["series"] * arr.size, "timestamp": ts, "target": arr} # Chronos-2 natively supports past covariates; attach any we were given. if covariates: for idx, cov in enumerate(covariates): cov_arr = np.asarray(cov, dtype=float) if cov_arr.size == arr.size: frame[f"cov_{idx}"] = cov_arr context_df = pd.DataFrame(frame) pred_df = self._impl.predict_df( context_df, prediction_length=int(horizon), quantile_levels=[0.1, 0.5, 0.9], id_column="id", timestamp_column="timestamp", target="target", ) return self._extract_median(pred_df, horizon) @staticmethod def _extract_median(pred_df, horizon: int) -> np.ndarray: import pandas as pd # noqa: F401 # predict_df output column naming varies; prefer the 0.5 quantile/mean. for candidate in ("0.5", 0.5, "median", "mean", "predictions", "target"): if candidate in pred_df.columns: return np.asarray(pred_df[candidate].to_numpy(), dtype=float)[:horizon] # Last resort: first purely-numeric column that is not an id/time column. for col in pred_df.columns: if col in ("id", "timestamp", "item_id"): continue try: return np.asarray(pred_df[col].to_numpy(), dtype=float)[:horizon] except Exception: continue raise ValueError("could not locate a forecast column in predict_df output") def _forecast_bolt(self, arr: np.ndarray, horizon: int) -> np.ndarray: import torch # type: ignore context = torch.tensor(arr, dtype=torch.float32) # BaseChronosPipeline exposes predict_quantiles(...) -> (quantiles, mean) quantiles, mean = self._impl.predict_quantiles( context=context, prediction_length=int(horizon), quantile_levels=[0.1, 0.5, 0.9], ) median = quantiles[0, :, 1] # the 0.5 quantile return np.asarray(median, dtype=float)[:horizon]