| """ |
| Chronos-2 (amazon/chronos-2) integration. |
| |
| Chronos-2 is a 120M-parameter, encoder-only zero-shot time series |
| foundation model from Amazon Science (released 20 Oct 2025), successor |
| to the original Chronos/Chronos-Bolt family. Unlike those (and unlike |
| Moirai-MoE), Chronos-2 produces quantile forecasts DIRECTLY rather than |
| by sampling paths — confirmed by reading chronos/chronos2/pipeline.py's |
| own source (`forecast_type: ForecastType = ForecastType.QUANTILES`) in |
| a network-enabled sandbox pass, not assumed from documentation. That |
| means Chronos-2's output for identical input is deterministic by |
| construction — it does not need the seed fix moirai_model.py required. |
| |
| API confirmed two ways in that same sandbox pass: (1) the officially |
| documented usage pattern (huggingface.co/amazon/chronos-2, the |
| amazon-science/chronos-forecasting README), and (2) by installing |
| `chronos-forecasting>=2.0` with `--no-deps` (skipping torch, which |
| isn't installed here) and reading chronos/chronos2/pipeline.py's real |
| `predict_df()` signature, docstring, and implementation directly — |
| including its exact output column schema, which the README's example |
| doesn't fully spell out: |
| |
| from chronos import Chronos2Pipeline |
| pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2", device_map="cpu") |
| pred_df = pipeline.predict_df(df, id_column=..., timestamp_column=..., |
| target=..., prediction_length=H, |
| quantile_levels=[0.1, 0.25, 0.5, 0.75, 0.9]) |
| # columns: id_column, timestamp_column, "target_name", "predictions" |
| # (mean), and one column per quantile level named as str(level). |
| |
| `chronos/__init__.py` imports torch unconditionally at module level |
| (confirmed directly — unlike uni2ts's lazier import style), so even |
| constructing this class must defer the import to load()/predict(), |
| same reasoning as moirai_model.py's lazy imports but enforced one |
| level higher up. |
| |
| NETWORK NOTE: load() downloads the real checkpoint from Hugging Face on |
| first use, same as moirai_model.py. In this project's sandbox passes so |
| far, that download has never succeeded (no route to huggingface.co) -- |
| confirmed as ModelLoadError with the underlying exception attached, not |
| a fabricated result, matching moirai_model.py's own behavior in the |
| same situation. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| from forecast_types import ModelInferenceError, ModelInputError, ModelLoadError, QuantileForecast |
|
|
| CHECKPOINT = "amazon/chronos-2" |
| QUANTILE_LEVELS = [0.1, 0.25, 0.5, 0.75, 0.9] |
|
|
|
|
| class Chronos2Forecaster: |
| """Same load()/predict() contract as MoiraiMoEForecaster and |
| AutoArimaForecaster (see forecast_types.py) — no shared code with |
| either, no fallback to either.""" |
|
|
| def __init__(self, checkpoint: str = CHECKPOINT, device: str = "cpu", min_context: int = 32): |
| self.checkpoint = checkpoint |
| self.device = device |
| self.min_context = min_context |
| self.max_context = 2048 |
| self._pipeline = None |
| self._loaded = False |
|
|
| def load(self): |
| try: |
| from chronos import Chronos2Pipeline |
| except ImportError as e: |
| raise ModelLoadError( |
| f"chronos-forecasting (with torch) is not installed: {e}" |
| ) from e |
| try: |
| self._pipeline = Chronos2Pipeline.from_pretrained(self.checkpoint, device_map=self.device) |
| except Exception as e: |
| raise ModelLoadError(f"failed to load {self.checkpoint}: {e}") from e |
| self._loaded = True |
|
|
| def predict(self, context_df: pd.DataFrame, target_col: str, past_covariate_cols, |
| prediction_length: int, freq: str) -> QuantileForecast: |
| if not self._loaded: |
| raise ModelInferenceError("call load() before predict()") |
|
|
| closes = context_df[target_col] |
| if closes.isna().any() or not np.isfinite(closes.to_numpy()).all(): |
| raise ModelInputError("context contains NaN/inf values in the target column") |
|
|
| context_length = min(len(context_df), self.max_context) |
| if context_length < self.min_context: |
| raise ModelInputError( |
| f"context too short: {context_length} row(s), need at least {self.min_context}" |
| ) |
| window = context_df.iloc[-context_length:] |
|
|
| in_df = pd.DataFrame({ |
| "item_id": "series", |
| "timestamp": window.index, |
| "target": window[target_col].to_numpy(), |
| }) |
| try: |
| pred_df = self._pipeline.predict_df( |
| in_df, id_column="item_id", timestamp_column="timestamp", target="target", |
| prediction_length=prediction_length, quantile_levels=QUANTILE_LEVELS, |
| ) |
| except Exception as e: |
| raise ModelInferenceError(f"Chronos-2 predict_df failed: {e}") from e |
|
|
| try: |
| quantiles = pred_df[[str(level) for level in QUANTILE_LEVELS]].to_numpy() |
| except KeyError as e: |
| raise ModelInferenceError( |
| f"predict_df's output columns didn't match the expected quantile-level " |
| f"names {[str(l) for l in QUANTILE_LEVELS]} -- got {list(pred_df.columns)}. " |
| f"The installed chronos-forecasting version's output schema may have changed " |
| f"since this was last verified against source (spec section 75)." |
| ) from e |
|
|
| return QuantileForecast( |
| horizon=prediction_length, quantile_levels=QUANTILE_LEVELS, |
| quantiles=quantiles, samples=None, |
| context_length_used=context_length, prediction_length=prediction_length, |
| ) |
|
|
