tsagentkit 1.0.2__py3-none-any.whl

tsagentkit/models/adapters/chronos.py

@@ -0,0 +1,387 @@
+"""Amazon Chronos2 TSFM adapter.
+
+Adapter for Amazon's Chronos2 time series forecasting models.
+Chronos2 is a family of pretrained models based on T5 architecture
+that supports zero-shot forecasting.
+
+Reference: https://github.com/amazon-science/chronos-forecasting
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+from tsagentkit.utils import quantile_col_name
+
+from .base import TSFMAdapter
+
+if TYPE_CHECKING:
+    from tsagentkit.contracts import ForecastResult, ModelArtifact
+    from tsagentkit.series import TSDataset
+
+
+class ChronosAdapter(TSFMAdapter):
+    """Adapter for Amazon Chronos2 time series models.
+
+    Chronos2 is a family of pretrained time series forecasting models
+    based on T5 architecture. It supports zero-shot forecasting on
+    unseen time series data.
+
+    Available model sizes:
+        - small: AutoGluon Chronos-2-Small (fast, lightweight)
+        - base: Amazon Chronos-2 (best accuracy)
+
+    Example:
+        >>> config = AdapterConfig(model_name="chronos", model_size="small")
+        >>> adapter = ChronosAdapter(config)
+        >>> adapter.load_model()
+        >>> result = adapter.predict(dataset, horizon=30)
+
+    Reference:
+        https://github.com/amazon-science/chronos-forecasting
+    """
+
+    # HuggingFace model IDs for Chronos2
+    MODEL_SIZES = {
+        "small": "autogluon/chronos-2-small",
+        "base": "amazon/chronos-2",
+    }
+
+    def load_model(self) -> None:
+        """Load Chronos model from HuggingFace.
+
+        Downloads and caches the model if not already present.
+
+        Raises:
+            ImportError: If chronos-forecasting is not installed
+            RuntimeError: If model loading fails
+        """
+        try:
+            from chronos import Chronos2Pipeline
+        except ImportError as e:
+            raise ImportError(
+                "chronos-forecasting>=2.0.0 is required for ChronosAdapter. "
+                "Install with: pip install 'chronos-forecasting>=2.0.0'"
+            ) from e
+
+        model_id = self.MODEL_SIZES.get(
+            self.config.model_size,
+            self.MODEL_SIZES["small"]
+        )
+
+        try:
+            self._model = Chronos2Pipeline.from_pretrained(
+                model_id,
+                device_map=self._device,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load Chronos model: {e}") from e
+
+    def fit(
+        self,
+        dataset: TSDataset,
+        prediction_length: int,
+        quantiles: list[float] | None = None,
+    ) -> ModelArtifact:
+        """Prepare Chronos for prediction.
+
+        Chronos is a zero-shot model and doesn't require training.
+
+        Args:
+            dataset: Dataset to validate
+            prediction_length: Forecast horizon
+            quantiles: Optional quantile levels
+
+        Returns:
+            ModelArtifact with model reference
+        """
+        from tsagentkit.contracts import ModelArtifact
+
+        if not self.is_loaded:
+            self.load_model()
+
+        self._validate_dataset(dataset)
+
+        return ModelArtifact(
+            model=self._model,
+            model_name=f"chronos-{self.config.model_size}",
+            config={
+                "model_size": self.config.model_size,
+                "device": self._device,
+                "prediction_length": prediction_length,
+                "quantiles": quantiles,
+            },
+        )
+
+    def predict(
+        self,
+        dataset: TSDataset,
+        horizon: int,
+        quantiles: list[float] | None = None,
+    ) -> ForecastResult:
+        """Generate forecasts using Chronos.
+
+        Supports covariate-informed forecasting when dataset contains
+        past_covariates (past_x) and/or future_covariates (future_x).
+
+        Args:
+            dataset: Historical data for context, optionally with covariates
+            horizon: Number of steps to forecast
+            quantiles: Quantile levels for probabilistic forecasts
+
+        Returns:
+            ForecastResult with predictions and provenance
+        """
+        if not self.is_loaded:
+            self.load_model()
+
+        context_df, future_df = self._to_chronos_df(dataset, horizon)
+
+        # Use predict_df for pandas-friendly API
+        # quantile_levels must not be None for Chronos 2.0
+        quantile_levels = quantiles if quantiles is not None else [0.1, 0.5, 0.9]
+        pred_df = self._model.predict_df(
+            context_df,
+            future_df=future_df if future_df is not None else None,
+            id_column="item_id",
+            timestamp_column="timestamp",
+            target="target",
+            prediction_length=horizon,
+            quantile_levels=quantile_levels,
+        )
+
+        return self._to_forecast_result(pred_df, dataset, horizon, quantiles)
+
+    def _to_chronos_df(
+        self, dataset: TSDataset, horizon: int
+    ) -> tuple[pd.DataFrame, pd.DataFrame | None]:
+        """Convert TSDataset to Chronos DataFrame format with covariates.
+
+        Args:
+            dataset: Input dataset with optional covariates
+            horizon: Forecast horizon for generating future timestamps
+
+        Returns:
+            Tuple of (context_df, future_df):
+            - context_df: Historical data with target and covariates
+            - future_df: Future covariates only (no target), or None if no future covariates
+        """
+        # Start with base columns
+        df = dataset.df[["unique_id", "ds", "y"]].copy()
+        df = df.sort_values(["unique_id", "ds"]).reset_index(drop=True)
+
+        # Handle missing values in target
+        if df["y"].isna().any():
+            df["y"] = df.groupby("unique_id")["y"].transform(
+                self._handle_missing_values
+            )
+
+        # Merge past covariates if available
+        if dataset.past_x is not None:
+            df = self._merge_covariates(df, dataset.past_x, "past")
+
+        # Limit context length if specified
+        if self.config.max_context_length:
+            df = df.groupby("unique_id", as_index=False).tail(
+                self.config.max_context_length
+            )
+
+        # Rename columns for Chronos format
+        context_df = df.rename(
+            columns={"unique_id": "item_id", "ds": "timestamp", "y": "target"}
+        )
+
+        # Prepare future_df if future covariates are available
+        future_df = None
+        if dataset.future_x is not None:
+            future_df = self._prepare_future_df(dataset, horizon)
+
+        return context_df, future_df
+
+    def _merge_covariates(
+        self, df: pd.DataFrame, covariates: pd.DataFrame, cov_type: str
+    ) -> pd.DataFrame:
+        """Merge covariates into main DataFrame.
+
+        Args:
+            df: Main DataFrame with unique_id, ds, y
+            covariates: Covariate DataFrame
+            cov_type: Type of covariate ("past" or "future")
+
+        Returns:
+            Merged DataFrame
+        """
+        # Ensure covariates have proper index columns
+        cov_df = covariates.copy()
+
+        # Reset index if needed to get unique_id and ds as columns
+        if isinstance(cov_df.index, pd.MultiIndex):
+            cov_df = cov_df.reset_index()
+
+        # Ensure required columns exist
+        if "unique_id" not in cov_df.columns and "id" in cov_df.columns:
+            cov_df = cov_df.rename(columns={"id": "unique_id"})
+        if "ds" not in cov_df.columns and "timestamp" in cov_df.columns:
+            cov_df = cov_df.rename(columns={"timestamp": "ds"})
+
+        # Merge on unique_id and ds
+        merge_cols = ["unique_id", "ds"]
+        available_cols = [c for c in merge_cols if c in cov_df.columns]
+
+        if len(available_cols) == 2:
+            # Handle missing values in covariates before merging
+            cov_cols = [c for c in cov_df.columns if c not in merge_cols]
+            for col in cov_cols:
+                if cov_df[col].isna().any():
+                    cov_df[col] = self._handle_missing_values(cov_df[col])
+
+            df = df.merge(cov_df, on=merge_cols, how="left")
+
+        return df
+
+    def _prepare_future_df(
+        self, dataset: TSDataset, horizon: int
+    ) -> pd.DataFrame | None:
+        """Prepare future covariates DataFrame for Chronos.
+
+        Args:
+            dataset: Input dataset with future covariates
+            horizon: Forecast horizon
+
+        Returns:
+            Future covariates DataFrame without target column
+        """
+        if dataset.future_x is None:
+            return None
+
+        future_df = dataset.future_x.copy()
+
+        # Reset index if needed
+        if isinstance(future_df.index, pd.MultiIndex):
+            future_df = future_df.reset_index()
+
+        # Rename columns to Chronos format
+        if "unique_id" not in future_df.columns and "id" in future_df.columns:
+            future_df = future_df.rename(columns={"id": "unique_id"})
+        if "ds" not in future_df.columns and "timestamp" in future_df.columns:
+            future_df = future_df.rename(columns={"timestamp": "ds"})
+
+        # Rename to Chronos expected column names
+        future_df = future_df.rename(
+            columns={"unique_id": "item_id", "ds": "timestamp"}
+        )
+
+        # Ensure timestamp column exists - generate if needed
+        if "timestamp" not in future_df.columns and dataset.future_index is not None:
+            future_df = future_df.reset_index()
+            if "ds" in future_df.columns:
+                future_df = future_df.rename(columns={"ds": "timestamp"})
+
+        # Handle missing values in future covariates
+        for col in future_df.columns:
+            if col not in ["item_id", "timestamp"] and future_df[col].isna().any():
+                future_df[col] = self._handle_missing_values(future_df[col])
+
+        return future_df
+
+    def _handle_missing_values(
+        self, values: pd.Series | np.ndarray
+    ) -> pd.Series | np.ndarray:
+        """Fill missing values using linear interpolation.
+
+        Args:
+            values: Series or array that may contain NaNs
+
+        Returns:
+            Values with NaNs filled
+        """
+        is_array = isinstance(values, np.ndarray)
+        s = pd.Series(values).astype(float)
+        s = s.interpolate(method="linear", limit_direction="both")
+        if s.isna().any():
+            fill_val = 0.0 if pd.isna(s.mean()) else s.mean()
+            s = s.fillna(fill_val)
+        return s.values if is_array else s
+
+    def _to_forecast_result(
+        self,
+        pred_df: pd.DataFrame,
+        dataset: TSDataset,
+        horizon: int,
+        quantiles: list[float] | None,
+    ) -> ForecastResult:
+        """Convert Chronos predictions to ForecastResult.
+
+        Args:
+            pred_df: DataFrame from Chronos predict_df
+            dataset: Original dataset
+            horizon: Forecast horizon
+            quantiles: Quantile levels
+
+        Returns:
+            ForecastResult with predictions
+        """
+        from tsagentkit.contracts import ForecastResult
+
+        # Map column names
+        result_df = pred_df.rename(
+            columns={
+                "item_id": "unique_id",
+                "timestamp": "ds",
+                "predictions": "yhat",
+            }
+        )
+
+        # Handle quantile columns (Chronos returns them as strings like "0.1", "0.5")
+        if quantiles:
+            quantile_cols = {}
+            for col in result_df.columns:
+                if col not in ["unique_id", "ds", "yhat"]:
+                    try:
+                        q_val = float(col)
+                        if 0 < q_val < 1:
+                            quantile_cols[q_val] = col
+                    except (TypeError, ValueError):
+                        continue
+
+            for q in quantiles:
+                if quantile_cols:
+                    nearest = min(quantile_cols, key=lambda v: abs(v - q))
+                    result_df[quantile_col_name(q)] = result_df[quantile_cols[nearest]]
+
+        # Select and order columns
+        keep_cols = ["unique_id", "ds", "yhat"]
+        for q in quantiles or []:
+            col = quantile_col_name(q)
+            if col in result_df.columns:
+                keep_cols.append(col)
+
+        result_df = result_df[keep_cols].copy()
+        result_df["model"] = f"chronos-{self.config.model_size}"
+
+        provenance = self._create_provenance(dataset, horizon, quantiles)
+
+        return ForecastResult(
+            df=result_df,
+            provenance=provenance,
+            model_name=f"chronos-{self.config.model_size}",
+            horizon=horizon,
+        )
+
+    def get_model_signature(self) -> str:
+        """Return model signature for provenance."""
+        return f"chronos-{self.config.model_size}-{self._device}"
+
+    @classmethod
+    def _check_dependencies(cls) -> None:
+        """Check if Chronos dependencies are installed."""
+        try:
+            import chronos  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "chronos-forecasting is required. "
+                "Install with: pip install chronos-forecasting"
+            ) from e

tsagentkit/models/adapters/moirai.py

@@ -0,0 +1,256 @@
+"""Salesforce Moirai 2.0 TSFM adapter.
+
+Adapter for Salesforce's Moirai 2.0 universal time series forecasting model.
+Moirai 2.0 is a transformer-based model with improved architecture.
+
+Reference: https://github.com/SalesforceAIResearch/uni2ts
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+from tsagentkit.time import normalize_pandas_freq
+from tsagentkit.utils import quantile_col_name
+
+from .base import TSFMAdapter
+
+if TYPE_CHECKING:
+    from tsagentkit.contracts import ForecastResult, ModelArtifact
+    from tsagentkit.series import TSDataset
+
+
+class MoiraiAdapter(TSFMAdapter):
+    """Adapter for Salesforce Moirai 2.0 foundation model.
+
+    Moirai 2.0 is a universal time series forecasting transformer with
+    improved architecture over Moirai 1.x.
+
+    Available model sizes:
+        - small: 384d model, 6 layers (recommended for fast inference)
+
+    Example:
+        >>> config = AdapterConfig(model_name="moirai", model_size="small")
+        >>> adapter = MoiraiAdapter(config)
+        >>> adapter.load_model()
+        >>> result = adapter.predict(dataset, horizon=30)
+
+    Reference:
+        https://github.com/SalesforceAIResearch/uni2ts
+    """
+
+    # HuggingFace model ID for Moirai 2.0 (only small available currently)
+    MODEL_ID = "Salesforce/moirai-2.0-R-small"
+
+    # Default context length for Moirai 2.0
+    DEFAULT_CONTEXT_LENGTH = 512
+
+    def load_model(self) -> None:
+        """Load Moirai 2.0 model from HuggingFace.
+
+        Downloads and caches the model if not already present.
+
+        Raises:
+            ImportError: If uni2ts is not installed
+            RuntimeError: If model loading fails
+        """
+        try:
+            from uni2ts.model.moirai2 import Moirai2Module
+        except ImportError as e:
+            raise ImportError(
+                "uni2ts>=2.0.0 is required for MoiraiAdapter. "
+                "Install with: pip install 'uni2ts @ git+https://github.com/SalesforceAIResearch/uni2ts.git'"
+            ) from e
+
+        self._module = Moirai2Module.from_pretrained(self.MODEL_ID)
+
+    def fit(
+        self,
+        dataset: TSDataset,
+        prediction_length: int,
+        quantiles: list[float] | None = None,
+    ) -> ModelArtifact:
+        """Prepare Moirai for prediction.
+
+        Moirai is a zero-shot model and doesn't require training.
+
+        Args:
+            dataset: Dataset to validate
+            prediction_length: Forecast horizon
+            quantiles: Optional quantile levels
+
+        Returns:
+            ModelArtifact with model reference
+        """
+        from tsagentkit.contracts import ModelArtifact
+
+        if not self.is_loaded:
+            self.load_model()
+
+        self._validate_dataset(dataset)
+
+        return ModelArtifact(
+            model=self._module,
+            model_name="moirai-2.0",
+            config={
+                "model_size": "small",
+                "device": self._device,
+                "prediction_length": prediction_length,
+                "quantiles": quantiles,
+            },
+        )
+
+    def predict(
+        self,
+        dataset: TSDataset,
+        horizon: int,
+        quantiles: list[float] | None = None,
+    ) -> ForecastResult:
+        """Generate forecasts using Moirai 2.0.
+
+        Args:
+            dataset: Historical data for context
+            horizon: Number of steps to forecast
+            quantiles: Quantile levels for probabilistic forecasts
+
+        Returns:
+            ForecastResult with predictions and provenance
+        """
+        if not self.is_loaded:
+            self.load_model()
+
+        from gluonts.dataset.common import ListDataset
+        from uni2ts.model.moirai2 import Moirai2Forecast
+
+        from tsagentkit.contracts import ForecastResult
+
+        freq = normalize_pandas_freq(dataset.freq)
+        context_length = self._get_context_length(dataset, horizon)
+
+        model = Moirai2Forecast(
+            module=self._module,
+            prediction_length=horizon,
+            context_length=context_length,
+            target_dim=1,
+            feat_dynamic_real_dim=0,
+            past_feat_dynamic_real_dim=0,
+        )
+        predictor = model.create_predictor(
+            batch_size=self.config.prediction_batch_size or 32
+        )
+
+        entries = []
+        meta = []
+        for uid in dataset.series_ids:
+            series_df = dataset.get_series(uid).sort_values("ds")
+            values = series_df["y"].values.astype(np.float32)
+            if np.any(np.isnan(values)):
+                values = self._handle_missing_values(values)
+            entries.append(
+                {
+                    "item_id": uid,
+                    "start": series_df["ds"].iloc[0],
+                    "target": values,
+                }
+            )
+            meta.append({"uid": uid, "last_date": series_df["ds"].max()})
+
+        gluonts_ds = ListDataset(entries, freq=freq)
+        forecast_it = predictor.predict(gluonts_ds)
+
+        offset = pd.tseries.frequencies.to_offset(freq)
+        result_rows = []
+        for meta_item, forecast in zip(meta, forecast_it, strict=False):
+            uid = meta_item["uid"]
+            last_date = meta_item["last_date"]
+            future_dates = pd.date_range(
+                start=last_date + offset,
+                periods=horizon,
+                freq=freq,
+            )
+
+            point_forecast = (
+                forecast.quantile(0.5)
+                if quantiles and 0.5 in quantiles
+                else forecast.mean
+            )
+            point_forecast = np.asarray(point_forecast).flatten()
+
+            quantile_arrays: dict[float, np.ndarray] = {}
+            if quantiles:
+                for q in quantiles:
+                    try:
+                        quantile_arrays[q] = np.asarray(forecast.quantile(q)).flatten()
+                    except Exception:
+                        quantile_arrays[q] = point_forecast
+
+            for h in range(horizon):
+                row = {
+                    "unique_id": uid,
+                    "ds": future_dates[h],
+                    "yhat": float(point_forecast[h]),
+                }
+                for q in quantiles or []:
+                    row[quantile_col_name(q)] = float(quantile_arrays[q][h])
+                result_rows.append(row)
+
+        result_df = pd.DataFrame(result_rows)
+        result_df["model"] = "moirai-2.0"
+        provenance = self._create_provenance(dataset, horizon, quantiles)
+
+        return ForecastResult(
+            df=result_df,
+            provenance=provenance,
+            model_name="moirai-2.0",
+            horizon=horizon,
+        )
+
+    def _get_context_length(self, dataset: TSDataset, horizon: int) -> int:
+        """Get appropriate context length for prediction.
+
+        Args:
+            dataset: Input dataset
+            horizon: Forecast horizon
+
+        Returns:
+            Context length (capped at model max)
+        """
+        max_series_len = int(
+            dataset.df.groupby("unique_id").size().max()
+        ) if not dataset.df.empty else 0
+
+        context_length = self.config.max_context_length or max_series_len
+        context_length = max(context_length, horizon)
+        context_length = min(context_length, self.DEFAULT_CONTEXT_LENGTH)
+        return max(1, int(context_length))
+
+    def _handle_missing_values(self, values: np.ndarray) -> np.ndarray:
+        """Handle missing values in series.
+
+        Args:
+            values: Array that may contain NaNs
+
+        Returns:
+            Array with NaNs filled
+        """
+        s = pd.Series(values)
+        s = s.interpolate(method="linear", limit_direction="both")
+        return s.fillna(s.mean()).values.astype(np.float32)
+
+    def get_model_signature(self) -> str:
+        """Return model signature for provenance."""
+        return f"moirai-2.0-{self._device}"
+
+    @classmethod
+    def _check_dependencies(cls) -> None:
+        """Check if Moirai dependencies are installed."""
+        try:
+            from uni2ts.model.moirai2 import Moirai2Module  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "uni2ts>=2.0.0 is required. "
+                "Install with: pip install 'uni2ts @ git+https://github.com/SalesforceAIResearch/uni2ts.git'"
+            ) from e