tsagentkit 1.0.2__py3-none-any.whl

tsagentkit/models/__init__.py

@@ -0,0 +1,182 @@
+"""Models module for tsagentkit.
+
+Provides model fitting, prediction, and TSFM (Time-Series Foundation Model)
+adapters for various forecasting backends.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from datetime import UTC
+from typing import TYPE_CHECKING, Any
+
+from tsagentkit.contracts import ForecastResult, ModelArtifact, Provenance
+
+# Import adapters submodules
+from tsagentkit.models import adapters
+from tsagentkit.models.baselines import fit_baseline, is_baseline_model, predict_baseline
+from tsagentkit.models.sktime import SktimeModelBundle, fit_sktime, predict_sktime
+from tsagentkit.utils import normalize_quantile_columns
+
+if TYPE_CHECKING:
+    from tsagentkit.contracts import TaskSpec
+    from tsagentkit.router import PlanSpec
+    from tsagentkit.series import TSDataset
+
+
+def _is_tsfm_model(model_name: str) -> bool:
+    return model_name.lower().startswith("tsfm-")
+
+
+def _is_sktime_model(model_name: str) -> bool:
+    return model_name.lower().startswith("sktime-")
+
+
+def _build_adapter_config(model_name: str, config: dict[str, Any]) -> adapters.AdapterConfig:
+    adapter_name = model_name.split("tsfm-", 1)[-1]
+    return adapters.AdapterConfig(
+        model_name=adapter_name,
+        model_size=config.get("model_size", "base"),
+        device=config.get("device"),
+        cache_dir=config.get("cache_dir"),
+        batch_size=config.get("batch_size", 32),
+        prediction_batch_size=config.get("prediction_batch_size", 100),
+        quantile_method=config.get("quantile_method", "sample"),
+        num_samples=config.get("num_samples", 100),
+        max_context_length=config.get("max_context_length"),
+    )
+
+
+def _fit_model_name(
+    model_name: str,
+    dataset: TSDataset,
+    plan: PlanSpec,
+    covariates: Any | None = None,
+) -> ModelArtifact:
+    """Fit a model by name with baseline or TSFM dispatch."""
+    config: dict[str, Any] = {
+        "horizon": dataset.task_spec.horizon,
+        "season_length": dataset.task_spec.season_length or 1,
+        "quantiles": plan.quantiles,
+    }
+
+    if _is_tsfm_model(model_name):
+        adapter_name = model_name.split("tsfm-", 1)[-1]
+        adapter_config = _build_adapter_config(model_name, {})
+        adapter = adapters.AdapterRegistry.create(adapter_name, adapter_config)
+        adapter.fit(
+            dataset=dataset,
+            prediction_length=dataset.task_spec.horizon,
+            quantiles=plan.quantiles,
+        )
+        return ModelArtifact(
+            model=adapter,
+            model_name=model_name,
+            config=config,
+            metadata={"adapter": adapter_name},
+        )
+
+    if is_baseline_model(model_name):
+        return fit_baseline(model_name, dataset, config)
+
+    if _is_sktime_model(model_name):
+        return fit_sktime(model_name, dataset, plan, covariates=covariates)
+
+    raise ValueError(f"Unknown model name: {model_name}")
+
+
+def fit(
+    dataset: TSDataset,
+    plan: PlanSpec,
+    on_fallback: Callable[[str, str, Exception], None] | None = None,
+    covariates: Any | None = None,
+) -> ModelArtifact:
+    """Fit a model using the plan's fallback ladder."""
+    from tsagentkit.router import execute_with_fallback
+
+    def _fit(model_name: str, ds: TSDataset) -> ModelArtifact:
+        return _fit_model_name(model_name, ds, plan, covariates=covariates)
+
+    artifact, _ = execute_with_fallback(
+        fit_func=_fit,
+        dataset=dataset,
+        plan=plan,
+        on_fallback=on_fallback,
+    )
+    return artifact
+
+
+def _basic_provenance(
+    dataset: TSDataset,
+    spec: TaskSpec,
+    artifact: ModelArtifact,
+) -> Provenance:
+    from datetime import datetime
+
+    from tsagentkit.utils import compute_data_signature
+
+    return Provenance(
+        run_id=f"model_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}",
+        timestamp=datetime.now(UTC).isoformat(),
+        data_signature=compute_data_signature(dataset.df),
+        task_signature=spec.model_hash(),
+        plan_signature=artifact.signature,
+        model_signature=artifact.signature,
+        metadata={"provenance_incomplete": True},
+    )
+
+
+def predict(
+    dataset: TSDataset,
+    artifact: ModelArtifact,
+    spec: TaskSpec,
+    covariates: Any | None = None,
+) -> ForecastResult:
+    """Generate predictions for baseline or TSFM models."""
+    if isinstance(artifact.model, adapters.TSFMAdapter):
+        result = artifact.model.predict(
+            dataset=dataset,
+            horizon=spec.horizon,
+            quantiles=artifact.config.get("quantiles"),
+        )
+        df = normalize_quantile_columns(result.df)
+        if "model" not in df.columns:
+            df = df.copy()
+            df["model"] = artifact.model_name
+        return ForecastResult(
+            df=df,
+            provenance=result.provenance,
+            model_name=artifact.model_name,
+            horizon=spec.horizon,
+        )
+
+    if is_baseline_model(artifact.model_name):
+        forecast_df = predict_baseline(
+            model_artifact=artifact,
+            dataset=dataset,
+            horizon=spec.horizon,
+            quantiles=artifact.config.get("quantiles"),
+        )
+        if "model" not in forecast_df.columns:
+            forecast_df = forecast_df.copy()
+            forecast_df["model"] = artifact.model_name
+        provenance = _basic_provenance(dataset, spec, artifact)
+        return ForecastResult(
+            df=normalize_quantile_columns(forecast_df),
+            provenance=provenance,
+            model_name=artifact.model_name,
+            horizon=spec.horizon,
+        )
+
+    if isinstance(artifact.model, SktimeModelBundle):
+        return predict_sktime(
+            dataset=dataset,
+            artifact=artifact,
+            spec=spec,
+            covariates=covariates,
+        )
+
+    raise ValueError(f"Unknown model type for prediction: {artifact.model_name}")
+
+
+__all__ = ["fit", "predict", "adapters"]

tsagentkit/models/adapters/__init__.py

@@ -0,0 +1,83 @@
+"""TSFM (Time-Series Foundation Model) adapters for tsagentkit.
+
+This module provides unified adapters for popular time-series foundation models:
+- Amazon Chronos
+- Salesforce Moirai
+- Google TimesFM
+
+Example:
+    >>> from tsagentkit.models.adapters import AdapterConfig, AdapterRegistry
+    >>> from tsagentkit.models.adapters import ChronosAdapter
+    >>>
+    >>> # Register adapters
+    >>> AdapterRegistry.register("chronos", ChronosAdapter)
+    >>>
+    >>> # Create and use adapter
+    >>> config = AdapterConfig(model_name="chronos", model_size="base")
+    >>> adapter = AdapterRegistry.create("chronos", config)
+    >>> adapter.load_model()
+    >>> result = adapter.predict(dataset, horizon=30)
+"""
+
+from __future__ import annotations
+
+# Base classes
+from .base import AdapterConfig, TSFMAdapter
+from .registry import AdapterRegistry
+
+__all__ = [
+    # Base classes
+    "AdapterConfig",
+    "TSFMAdapter",
+    "AdapterRegistry",
+]
+
+# Lazily import specific adapters as they become available
+# This prevents import errors when dependencies are not installed
+
+
+def _try_import_chronos():
+    """Try to import ChronosAdapter if dependencies are available."""
+    try:
+        from .chronos import ChronosAdapter
+
+        return ChronosAdapter
+    except ImportError:
+        return None
+
+
+def _try_import_moirai():
+    """Try to import MoiraiAdapter if dependencies are available."""
+    try:
+        from .moirai import MoiraiAdapter
+
+        return MoiraiAdapter
+    except ImportError:
+        return None
+
+
+def _try_import_timesfm():
+    """Try to import TimesFMAdapter if dependencies are available."""
+    try:
+        from .timesfm import TimesFMAdapter
+
+        return TimesFMAdapter
+    except ImportError:
+        return None
+
+
+# Auto-register available adapters
+_chronos = _try_import_chronos()
+if _chronos:
+    AdapterRegistry.register("chronos", _chronos)
+    __all__.append("ChronosAdapter")
+
+_moirai = _try_import_moirai()
+if _moirai:
+    AdapterRegistry.register("moirai", _moirai)
+    __all__.append("MoiraiAdapter")
+
+_timesfm = _try_import_timesfm()
+if _timesfm:
+    AdapterRegistry.register("timesfm", _timesfm)
+    __all__.append("TimesFMAdapter")

tsagentkit/models/adapters/base.py

@@ -0,0 +1,321 @@
+"""Base adapter class for Time-Series Foundation Models.
+
+Provides a unified interface for integrating external TSFMs like Chronos,
+Moirai, and TimesFM with the tsagentkit pipeline.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import UTC
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+
+    from tsagentkit.contracts import ForecastResult, ModelArtifact, Provenance
+    from tsagentkit.series import TSDataset
+
+
+@dataclass(frozen=True)
+class AdapterConfig:
+    """Configuration for TSFM adapter.
+
+    Attributes:
+        model_name: Name of the model/adapter
+        model_size: Model size variant (small, base, large)
+        device: Compute device (cuda, mps, cpu, or None for auto)
+        cache_dir: Directory for model caching
+        batch_size: Batch size for training/fitting
+        prediction_batch_size: Batch size for prediction
+        quantile_method: Method for quantile prediction (sample, direct)
+        num_samples: Number of samples for probabilistic forecasting
+        max_context_length: Maximum context length the model accepts
+    """
+
+    model_name: str
+    model_size: Literal["tiny", "small", "base", "large"] = "base"
+    device: str | None = None  # Auto-detect if None
+    cache_dir: str | None = None
+    batch_size: int = 32
+    prediction_batch_size: int = 100
+    quantile_method: Literal["sample", "direct"] = "sample"
+    num_samples: int = 100
+    max_context_length: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        valid_sizes = {"tiny", "small", "base", "large"}
+        if self.model_size not in valid_sizes:
+            raise ValueError(
+                f"Invalid model_size '{self.model_size}'. "
+                f"Must be one of: {valid_sizes}"
+            )
+
+        valid_methods = {"sample", "direct"}
+        if self.quantile_method not in valid_methods:
+            raise ValueError(
+                f"Invalid quantile_method '{self.quantile_method}'. "
+                f"Must be one of: {valid_methods}"
+            )
+
+
+class TSFMAdapter(ABC):
+    """Abstract base class for Time-Series Foundation Model adapters.
+
+    Provides unified interface for different TSFMs while handling
+    model-specific quirks internally. Subclasses implement specific
+    adapters for models like Chronos, Moirai, and TimesFM.
+
+    Example:
+        >>> config = AdapterConfig(model_name="chronos", model_size="base")
+        >>> adapter = ChronosAdapter(config)
+        >>> adapter.load_model()
+        >>> result = adapter.predict(dataset, horizon=30)
+
+    Attributes:
+        config: Adapter configuration
+        _model: The underlying model instance (lazy loaded)
+        _device: Resolved compute device
+    """
+
+    def __init__(self, config: AdapterConfig):
+        """Initialize adapter with configuration.
+
+        Args:
+            config: Adapter configuration
+        """
+        self.config = config
+        self._model: Any | None = None
+        self._device = self._resolve_device()
+
+    def _resolve_device(self) -> str:
+        """Resolve compute device (cuda/mps/cpu).
+
+        Returns the best available device based on hardware support.
+        Priority: cuda > mps > cpu
+
+        Returns:
+            Device string for PyTorch
+        """
+        if self.config.device:
+            return self.config.device
+
+        # Try to import torch for device detection
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                return "cuda"
+            elif torch.backends.mps.is_available():
+                return "mps"
+        except ImportError:
+            pass
+
+        return "cpu"
+
+    @abstractmethod
+    def load_model(self) -> None:
+        """Load the foundation model with caching.
+
+        Downloads and caches the model if not already present.
+        Should be called before fit() or predict().
+
+        Raises:
+            ImportError: If required dependencies are not installed
+            RuntimeError: If model loading fails
+        """
+        pass
+
+    @abstractmethod
+    def fit(
+        self,
+        dataset: TSDataset,
+        prediction_length: int,
+        quantiles: list[float] | None = None,
+    ) -> ModelArtifact:
+        """Prepare model for prediction on the dataset.
+
+        Note: Most TSFMs are zero-shot and don't require traditional fitting.
+        This method validates compatibility and may perform preprocessing.
+
+        Args:
+            dataset: Training dataset
+            prediction_length: Forecast horizon
+            quantiles: Optional quantile levels for probabilistic forecasts
+
+        Returns:
+            ModelArtifact with model reference and configuration
+
+        Raises:
+            ValueError: If dataset is incompatible with model
+            RuntimeError: If preparation fails
+        """
+        pass
+
+    @abstractmethod
+    def predict(
+        self,
+        dataset: TSDataset,
+        horizon: int,
+        quantiles: list[float] | None = None,
+    ) -> ForecastResult:
+        """Generate forecasts using the TSFM.
+
+        Args:
+            dataset: Historical data for context
+            horizon: Number of steps to forecast
+            quantiles: Optional quantile levels (e.g., [0.1, 0.5, 0.9])
+
+        Returns:
+            ForecastResult with predictions and provenance
+
+        Raises:
+            RuntimeError: If prediction fails
+            ValueError: If horizon exceeds model limits
+        """
+        pass
+
+    @abstractmethod
+    def get_model_signature(self) -> str:
+        """Return unique signature for this model configuration.
+
+        Used for provenance tracking to identify the exact model
+        configuration used for a forecast.
+
+        Returns:
+            Unique signature string (e.g., "chronos-base-cuda-v1.0")
+        """
+        pass
+
+    @property
+    def is_loaded(self) -> bool:
+        """Check if model is loaded in memory.
+
+        Returns:
+            True if model has been loaded, False otherwise
+        """
+        return self._model is not None
+
+    def unload_model(self) -> None:
+        """Unload model from memory to free resources.
+
+        Useful for managing memory when using multiple large models.
+        """
+        self._model = None
+
+    def _validate_dataset(self, dataset: TSDataset) -> None:
+        """Validate dataset compatibility.
+
+        Args:
+            dataset: Dataset to validate
+
+        Raises:
+            ValueError: If dataset is incompatible
+        """
+        # Check minimum length requirements
+        min_length = 1  # Most TSFMs need at least some context
+
+        for uid in dataset.series_ids:
+            series = dataset.get_series(uid)
+            if len(series) < min_length:
+                raise ValueError(
+                    f"Series '{uid}' has only {len(series)} observations. "
+                    f"Minimum required: {min_length}"
+                )
+
+    def _create_provenance(
+        self,
+        dataset: TSDataset,
+        horizon: int,
+        quantiles: list[float] | None = None,
+    ) -> Provenance:
+        """Create provenance for forecast.
+
+        Args:
+            dataset: Input dataset
+            horizon: Forecast horizon
+            quantiles: Quantile levels used
+
+        Returns:
+            Provenance instance
+        """
+        # Create data signature from dataset
+        import hashlib
+        from datetime import datetime
+
+        from tsagentkit.contracts import Provenance
+
+        data_hash = hashlib.sha256(
+            str(dataset.n_series).encode() +
+            str(dataset.n_observations).encode() +
+            dataset.date_range[0].isoformat().encode()
+        ).hexdigest()[:16]
+
+        return Provenance(
+            run_id=f"tsfm_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}",
+            timestamp=datetime.now(UTC).isoformat(),
+            data_signature=data_hash,
+            task_signature=f"horizon={horizon}",
+            plan_signature=self.get_model_signature(),
+            model_signature=self.get_model_signature(),
+            metadata={
+                "adapter": self.__class__.__name__,
+                "device": self._device,
+                "quantiles": quantiles,
+            },
+        )
+
+    def _batch_iterator(
+        self,
+        data: list[Any],
+        batch_size: int,
+    ):
+        """Iterate over data in batches.
+
+        Args:
+            data: List of data items
+            batch_size: Batch size
+
+        Yields:
+            Batches of data
+        """
+        for i in range(0, len(data), batch_size):
+            yield data[i : i + batch_size]
+
+    def _compute_quantiles_from_samples(
+        self,
+        samples: Any,  # numpy array or tensor
+        quantiles: list[float],
+    ) -> dict[float, float]:
+        """Compute quantiles from sampled predictions.
+
+        Args:
+            samples: Array of samples with shape (n_samples, horizon)
+            quantiles: Quantile levels to compute
+
+        Returns:
+            Dictionary mapping quantile levels to values
+        """
+        import numpy as np
+
+        results = {}
+        for q in quantiles:
+            results[q] = float(np.quantile(samples, q, axis=0))
+        return results
+
+    @classmethod
+    def _check_dependencies(cls) -> None:
+        """Check if required dependencies are installed.
+
+        Raises:
+            ImportError: If dependencies are missing
+        """
+        # Base class checks for torch only
+        try:
+            import torch  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "PyTorch is required for TSFM adapters. "
+                "Install with: pip install torch"
+            ) from e