"""Forecast model adapters. Adding a model means giving a model id and naming an adapter *family*. The family list is fixed in `config.ALLOWED_ADAPTER_FAMILIES`; nothing here ever imports, downloads or executes code chosen by a user. A user-supplied model id is loaded only through an already-vetted family's loader, and `trust_remote_code` is never enabled. Every adapter reports an `inference_version` that pins the adapter's own logic alongside the model revision, so a stored slice can always be traced back to exactly what produced it. """ from __future__ import annotations import hashlib import logging import os from abc import ABC, abstractmethod from dataclasses import dataclass import numpy as np import pandas as pd from . import config log = logging.getLogger("bit.adapters") DEFAULT_QUANTILES = (0.1, 0.5, 0.9) class AdapterError(RuntimeError): pass class ModelNotAllowed(AdapterError): """The requested adapter family is not on the allow-list.""" @dataclass class Forecast: """One horizon-1 forecast per input window.""" q10: np.ndarray q50: np.ndarray q90: np.ndarray context_len: int def __post_init__(self): if not (len(self.q10) == len(self.q50) == len(self.q90)): raise AdapterError("quantile arrays have mismatched lengths") def as_frame(self, ts: pd.DatetimeIndex, inference_version: str) -> pd.DataFrame: if len(ts) != len(self.q50): raise AdapterError( f"timestamp count {len(ts)} != forecast count {len(self.q50)}" ) # Quantiles must not cross; sorting is the honest repair for tiny # numerical inversions and makes the store's validator pass. stacked = np.sort(np.vstack([self.q10, self.q50, self.q90]), axis=0) return pd.DataFrame({ "ts": ts, "q10": stacked[0], "q50": stacked[1], "q90": stacked[2], "context_len": self.context_len, "inference_version": inference_version, }) class ForecastAdapter(ABC): """Uniform interface over quantile forecasters.""" family: str = "base" adapter_version: str = "1" def __init__(self, model_id: str, revision: str | None = None, context_len: int = 512, device: str | None = None): self.model_id = model_id self.revision = revision self.context_len = context_len self.device = device or _default_device() self._model = None self._resolved_revision: str | None = None # -- interface -------------------------------------------------------- @abstractmethod def load(self, model_id: str | None = None, revision: str | None = None) -> "ForecastAdapter": """Materialise the model. Idempotent.""" @abstractmethod def predict(self, context_windows: np.ndarray) -> Forecast: """`context_windows` is (n_windows, context_len). Returns horizon-1 quantiles.""" def inference_version(self) -> str: """Identity of everything that determines the output values.""" rev = self._resolved_revision or self.revision or "unpinned" payload = f"{self.family}|{self.adapter_version}|{self.model_id}|{rev}|{self.context_len}" digest = hashlib.sha256(payload.encode()).hexdigest()[:12] return f"{config.INFERENCE_VERSION}+{self.family}.{digest}" @property def resolved_revision(self) -> str: return self._resolved_revision or self.revision or "unpinned" # -- shared helpers --------------------------------------------------- def resolve_revision(self) -> str: """Pin the model to an immutable commit sha before any inference runs.""" if self._resolved_revision: return self._resolved_revision try: from huggingface_hub import HfApi info = HfApi().model_info(self.model_id, revision=self.revision) self._resolved_revision = info.sha except Exception as e: log.warning("could not resolve revision for %s: %s", self.model_id, e) self._resolved_revision = self.revision or "unpinned" return self._resolved_revision def _default_device() -> str: try: import torch if torch.cuda.is_available(): return "cuda" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): return "mps" except Exception: pass return "cpu" def _release(device: str) -> None: """Free accelerator memory between inference chunks.""" try: import torch if device == "cuda": torch.cuda.empty_cache() elif device == "mps" and hasattr(torch, "mps"): torch.mps.empty_cache() except Exception: pass # -------------------------------------------------------------------------- # Chronos / Chronos-Bolt # -------------------------------------------------------------------------- class ChronosAdapter(ForecastAdapter): """amazon/chronos-* and amazon/chronos-bolt-* quantile forecasters.""" family = "chronos" adapter_version = "1" def load(self, model_id: str | None = None, revision: str | None = None): if model_id: self.model_id = model_id if revision: self.revision = revision if self._model is not None: return self self.resolve_revision() try: from chronos import BaseChronosPipeline except ImportError as e: raise AdapterError( "chronos-forecasting is not installed; add it to requirements.txt" ) from e import torch dtype = torch.float32 if self.device in ("cpu", "mps") else torch.bfloat16 self._model = BaseChronosPipeline.from_pretrained( self.model_id, revision=self._resolved_revision if self._resolved_revision != "unpinned" else None, device_map=self.device, torch_dtype=dtype, ) return self # Bolt emits quantiles directly and is happy with wide batches. The # original T5 Chronos *samples* forecast paths instead, so its peak memory # is batch x num_samples x context and a wide batch OOMs a GPU outright. # Chunking here rather than at the call site means every caller -- seed # script, ZeroGPU function, tests -- is protected without knowing which # variant it holds. BOLT_CHUNK = 256 T5_CHUNK = 16 T5_NUM_SAMPLES = 20 CHRONOS2_CHUNK = 128 @property def _is_bolt(self) -> bool: return "bolt" in (self.model_id or "").lower() @property def _is_chronos2(self) -> bool: """Chronos-2 (`amazon/chronos-2`). Detected from the id rather than from the loaded pipeline class so the chunk size is known before the weights are fetched. """ name = (self.model_id or "").lower() return "chronos-2" in name or "chronos2" in name @property def chunk_size(self) -> int: if self._is_chronos2: return self.CHRONOS2_CHUNK return self.BOLT_CHUNK if self._is_bolt else self.T5_CHUNK def predict(self, context_windows: np.ndarray) -> Forecast: if self._model is None: self.load() import torch ctx = np.asarray(context_windows, dtype="float32") if ctx.ndim == 1: ctx = ctx[None, :] q_levels = list(DEFAULT_QUANTILES) # Only the original T5 Chronos samples paths; Bolt and Chronos-2 emit # quantiles directly. if self._is_bolt or self._is_chronos2: extra = {} else: extra = {"num_samples": self.T5_NUM_SAMPLES} step = max(1, self.chunk_size) parts = [] for start in range(0, len(ctx), step): block = ctx[start:start + step] tensors = [torch.tensor(row) for row in block] with torch.inference_mode(): quantiles, _mean = self._model.predict_quantiles( tensors, prediction_length=1, quantile_levels=q_levels, **extra, ) parts.append(self._to_array(quantiles)) del quantiles, tensors _release(self.device) arr = np.vstack(parts) return Forecast(q10=arr[:, 0], q50=arr[:, 1], q90=arr[:, 2], context_len=ctx.shape[1]) @staticmethod def _to_array(quantiles) -> np.ndarray: """Normalise a `predict_quantiles` result to `(batch, n_quantiles)`. The two shapes differ and the difference is not cosmetic: Bolt / T5 one stacked tensor, `(batch, horizon, quantiles)` Chronos-2 a *list* of per-item tensors, each `(n_variates, horizon, quantiles)` Calling `.float()` on the list raises `AttributeError`, which is exactly what `amazon/chronos-2` did before this existed. Horizon is always 1 here, and these are univariate price series, so the leading variate axis is taken at index 0. """ if isinstance(quantiles, (list, tuple)): rows = [] for item in quantiles: array = ChronosAdapter._as_numpy(item) # (n_variates, horizon, quantiles) -> (quantiles,) while array.ndim > 1: array = array[0] rows.append(array) return np.vstack(rows) return ChronosAdapter._as_numpy(quantiles)[:, 0, :] @staticmethod def _as_numpy(value) -> np.ndarray: """Tensor -> ndarray, by duck typing rather than an isinstance check. Deliberately does not import torch: the test suite runs offline with no GPU stack installed, and a shape-normalising helper that cannot be tested without a 2GB dependency would not get tested. """ if hasattr(value, "float") and hasattr(value, "cpu"): return np.asarray(value.float().cpu().numpy()) if isinstance(value, np.ndarray): return value raise AdapterError( f"unexpected predict_quantiles result: {type(value).__name__}") # -------------------------------------------------------------------------- # TimesFM # -------------------------------------------------------------------------- class TimesFMAdapter(ForecastAdapter): """google/timesfm-* forecasters. TimesFM returns a fixed quantile grid; the 0.1/0.5/0.9 columns are selected from it rather than re-derived, so the stored numbers are the model's own. """ family = "timesfm" adapter_version = "1" _QUANTILE_GRID = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9) def load(self, model_id: str | None = None, revision: str | None = None): if model_id: self.model_id = model_id if revision: self.revision = revision if self._model is not None: return self self.resolve_revision() try: import timesfm except ImportError as e: raise AdapterError( "timesfm is not installed; add it to requirements.txt to enable " "this adapter family" ) from e backend = {"cuda": "gpu", "mps": "cpu"}.get(self.device, "cpu") self._model = timesfm.TimesFm( hparams=timesfm.TimesFmHparams( backend=backend, per_core_batch_size=32, context_len=self.context_len, horizon_len=1, ), checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=self.model_id), ) return self def predict(self, context_windows: np.ndarray) -> Forecast: if self._model is None: self.load() ctx = np.asarray(context_windows, dtype="float32") if ctx.ndim == 1: ctx = ctx[None, :] _point, quantile_out = self._model.forecast( [row for row in ctx], freq=[0] * len(ctx) ) arr = np.asarray(quantile_out)[:, 0, :] # horizon 1 grid = list(self._QUANTILE_GRID) # Column 0 of TimesFM's output is the mean; quantiles follow. offset = arr.shape[1] - len(grid) i10, i50, i90 = (grid.index(q) + offset for q in DEFAULT_QUANTILES) return Forecast(q10=arr[:, i10], q50=arr[:, i50], q90=arr[:, i90], context_len=ctx.shape[1]) # -------------------------------------------------------------------------- # Naive baselines # -------------------------------------------------------------------------- class BaselineAdapter(ForecastAdapter): """Classical forecasting baselines, exposed as first-class models. These exist so the leaderboard always carries a floor. A large pretrained forecaster that cannot beat "tomorrow looks like today" has not earned its inference cost, and burying that comparison would be the single easiest way to make this whole Space misleading. The method is selected by the model id: `baseline/naive`, `baseline/drift`, or `baseline/seasonal`. All three are deterministic and causal -- each uses only the trailing window, which ends at the bar being forecast. The interval is a Gaussian band around the point forecast, scaled by the window's own realised volatility, so coverage is comparable with a learned model's q10-q90 rather than arbitrarily wide. """ family = "baseline" adapter_version = "1" Z90 = 1.2815515655446004 # normal 90th percentile def __init__(self, model_id: str, revision: str | None = None, context_len: int = 128, device: str | None = None, season: int = 7): super().__init__(model_id, revision, context_len, device) self.season = season self.method = (model_id.split("/")[-1] or "naive").lower() if self.method not in ("naive", "drift", "seasonal"): raise AdapterError( f"unknown baseline method {self.method!r}; " "use baseline/naive, baseline/drift or baseline/seasonal" ) def load(self, model_id: str | None = None, revision: str | None = None): if model_id: self.model_id = model_id self.method = model_id.split("/")[-1].lower() # Baselines have no weights, so the "revision" is the adapter's own # version -- still pinned, still reproducible. self._resolved_revision = f"baseline-{self.adapter_version}" return self def resolve_revision(self) -> str: self._resolved_revision = f"baseline-{self.adapter_version}" return self._resolved_revision def predict(self, context_windows: np.ndarray) -> Forecast: ctx = np.asarray(context_windows, dtype="float64") if ctx.ndim == 1: ctx = ctx[None, :] last = ctx[:, -1] if self.method == "naive": point = last elif self.method == "drift": # Average per-step change across the window, added once. steps = ctx.shape[1] - 1 point = last + (ctx[:, -1] - ctx[:, 0]) / max(steps, 1) else: # seasonal k = min(self.season, ctx.shape[1]) point = ctx[:, -k] # One-step-ahead uncertainty from the window's own step volatility. sigma = np.std(np.diff(ctx, axis=1), axis=1) sigma = np.where(np.isfinite(sigma) & (sigma > 0), sigma, np.abs(last) * 1e-4) band = self.Z90 * sigma return Forecast(q10=point - band, q50=point, q90=point + band, context_len=ctx.shape[1]) # -------------------------------------------------------------------------- # Placeholder (no GPU / no model available) # -------------------------------------------------------------------------- class PlaceholderAdapter(ForecastAdapter): """Structurally identical synthetic output, permanently labelled. Exists so the UI has complete shape before every cell has real coverage. Its `inference_version` is the literal string `PLACEHOLDER`, which the store and the UI both key off to mark the data as not-real. It is deterministic: the same window always yields the same numbers. """ family = "placeholder" adapter_version = "1" def load(self, model_id: str | None = None, revision: str | None = None): if model_id: self.model_id = model_id self._resolved_revision = "PLACEHOLDER" return self def inference_version(self) -> str: return config.PLACEHOLDER_VERSION def predict(self, context_windows: np.ndarray) -> Forecast: ctx = np.asarray(context_windows, dtype="float64") if ctx.ndim == 1: ctx = ctx[None, :] last = ctx[:, -1] # Deterministic pseudo-drift seeded by the window itself, plus a band # scaled to that window's realised volatility. seedvals = np.abs(np.sum(ctx[:, -8:], axis=1) * 1e6).astype("int64") drift = np.array([ (np.random.default_rng(int(s) % (2**32)).normal(0.0, 1.0)) for s in seedvals ]) vol = np.std(np.diff(ctx, axis=1), axis=1) / np.maximum(np.abs(last), 1e-9) vol = np.clip(vol, 1e-4, 0.2) q50 = last * (1.0 + 0.25 * drift * vol) band = last * vol * 1.28 return Forecast(q10=q50 - band, q50=q50, q90=q50 + band, context_len=ctx.shape[1]) # -------------------------------------------------------------------------- # Factory # -------------------------------------------------------------------------- _FAMILIES: dict[str, type[ForecastAdapter]] = { "chronos": ChronosAdapter, "timesfm": TimesFMAdapter, "baseline": BaselineAdapter, "placeholder": PlaceholderAdapter, } def get_adapter(family: str, model_id: str, revision: str | None = None, context_len: int = 512, device: str | None = None) -> ForecastAdapter: """Build an adapter. Only allow-listed families are constructible.""" fam = (family or "").strip().lower() if fam == "placeholder": return PlaceholderAdapter(model_id, revision, context_len, device) if fam not in config.ALLOWED_ADAPTER_FAMILIES: raise ModelNotAllowed( f"adapter family {family!r} is not allowed; pick one of " f"{list(config.ALLOWED_ADAPTER_FAMILIES)}" ) return _FAMILIES[fam](model_id, revision, context_len, device) def validate_model_id(model_id: str) -> str: """Reject anything that is not a plain `owner/name` Hub id.""" mid = (model_id or "").strip() if not mid or mid.count("/") != 1: raise AdapterError(f"{model_id!r} is not a valid Hub model id (owner/name)") owner, name = mid.split("/") ok = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.") if not owner or not name or set(mid) - ok - {"/"}: raise AdapterError(f"{model_id!r} contains characters that are not allowed") if ".." in mid: raise AdapterError("model id may not contain '..'") return mid # -------------------------------------------------------------------------- # Rolling-window construction (shared by the seed script and the GPU path) # -------------------------------------------------------------------------- def build_windows(series: pd.Series, context_len: int, stride: int = 1): """Yield (timestamp, trailing window) pairs. The window for timestamp `t` ends at `t` inclusive, so the forecast stored at `t` used only data available at `t`. The engine then shifts it before any trade can act on it. """ values = series.to_numpy(dtype="float64") index = series.index n = len(values) if n <= context_len: return [], np.empty((0, context_len)) stamps, rows = [], [] for i in range(context_len, n, stride): stamps.append(index[i]) rows.append(values[i - context_len + 1: i + 1]) if not rows: return [], np.empty((0, context_len)) return pd.DatetimeIndex(stamps), np.vstack(rows)