autogluon.timeseries 1.1.2b20241109__py3-none-any.whl → 1.1.2b20241112__py3-none-any.whl

autogluon/timeseries/dataset/ts_dataframe.py

@@ -921,7 +921,11 @@ class TimeSeriesDataFrame(pd.DataFrame, TimeSeriesDataFrameDeprecatedMixin):
         test_data : TimeSeriesDataFrame
             Test portion of the data. Contains the slice ``[:end_idx]`` of each time series in the original dataset.
         """
-        test_data = self.slice_by_timestep(None, end_index)
+        df = self
+        if not df.index.is_monotonic_increasing:
+            logger.warning("Sorting the dataframe index before generating the train/test split.")
+            df = df.sort_index()
+        test_data = df.slice_by_timestep(None, end_index)
         train_data = test_data.slice_by_timestep(None, -prediction_length)
 
         if suffix is not None:

autogluon/timeseries/models/abstract/abstract_timeseries_model.py

@@ -5,6 +5,8 @@ import time
 from contextlib import nullcontext
 from typing import Dict, List, Optional, Union
 
+import pandas as pd
+
 from autogluon.common import space
 from autogluon.common.loaders import load_pkl
 from autogluon.common.savers import save_pkl
@@ -13,8 +15,10 @@ from autogluon.core.hpo.executors import HpoExecutor, RayHpoExecutor
 from autogluon.core.models import AbstractModel
 from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_metric
+from autogluon.timeseries.regressor import CovariateRegressor
 from autogluon.timeseries.transforms import LocalTargetScaler, get_target_scaler_from_name
 from autogluon.timeseries.utils.features import CovariateMetadata
+from autogluon.timeseries.utils.forecast import get_forecast_horizon_index_ts_dataframe
 from autogluon.timeseries.utils.warning_filters import disable_stdout, warning_filter
 
 from .model_trial import model_trial, skip_hpo
@@ -164,6 +168,8 @@ class AbstractTimeSeriesModel(AbstractModel):
     def _initialize(self, **kwargs) -> None:
         self._init_params_aux()
         self._init_params()
+        self.target_scaler = self._create_target_scaler()
+        self.covariate_regressor = self._create_covariate_regressor()
 
     def _compute_fit_metadata(self, val_data: TimeSeriesDataFrame = None, **kwargs):
         fit_metadata = dict(
@@ -208,7 +214,11 @@ class AbstractTimeSeriesModel(AbstractModel):
         return info
 
     def fit(
-        self, train_data: TimeSeriesDataFrame, val_data: Optional[TimeSeriesDataFrame] = None, **kwargs
+        self,
+        train_data: TimeSeriesDataFrame,
+        val_data: Optional[TimeSeriesDataFrame] = None,
+        time_limit: Optional[float] = None,
+        **kwargs,
     ) -> "AbstractTimeSeriesModel":
         """Fit timeseries model.
 
@@ -243,22 +253,33 @@ class AbstractTimeSeriesModel(AbstractModel):
         model: AbstractTimeSeriesModel
             The fitted model object
         """
+        start_time = time.monotonic()
         self.initialize(**kwargs)
-        self.target_scaler = self._create_target_scaler()
         if self.target_scaler is not None:
             train_data = self.target_scaler.fit_transform(train_data)
 
+        if self.covariate_regressor is not None:
+            train_data = self.covariate_regressor.fit_transform(
+                train_data,
+                time_limit=0.5 * time_limit if time_limit is not None else None,
+            )
+
         train_data = self.preprocess(train_data, is_train=True)
         if self._get_tags()["can_use_val_data"] and val_data is not None:
             if self.target_scaler is not None:
                 val_data = self.target_scaler.transform(val_data)
+            if self.covariate_regressor is not None:
+                val_data = self.covariate_regressor.transform(val_data)
             val_data = self.preprocess(val_data, is_train=False)
-        return super().fit(train_data=train_data, val_data=val_data, **kwargs)
+
+        if time_limit is not None:
+            time_limit = time_limit - (time.monotonic() - start_time)
+        return super().fit(train_data=train_data, val_data=val_data, time_limit=time_limit, **kwargs)
 
     @property
     def allowed_hyperparameters(self) -> List[str]:
         """List of hyperparameters allowed by the model."""
-        return ["target_scaler"]
+        return ["target_scaler", "covariate_regressor"]
 
     def _create_target_scaler(self) -> Optional[LocalTargetScaler]:
         """Create a LocalTargetScaler object based on the value of the `target_scaler` hyperparameter."""
@@ -269,6 +290,32 @@ class AbstractTimeSeriesModel(AbstractModel):
         else:
             return None
 
+    def _create_covariate_regressor(self) -> Optional[CovariateRegressor]:
+        """Create a CovariateRegressor object based on the value of the `covariate_regressor` hyperparameter."""
+        covariate_regressor = self._get_model_params().get("covariate_regressor")
+        if covariate_regressor is not None:
+            if len(self.metadata.known_covariates + self.metadata.static_features) == 0:
+                logger.debug(
+                    "Skipping CovariateRegressor since the dataset contains no covariates or static features."
+                )
+                return None
+            else:
+                if isinstance(covariate_regressor, str):
+                    return CovariateRegressor(covariate_regressor, target=self.target, metadata=self.metadata)
+                elif isinstance(covariate_regressor, CovariateRegressor):
+                    logger.warning(
+                        "Using a custom CovariateRegressor object is experimental functionality that may break in the future!"
+                    )
+                    covariate_regressor.target = self.target
+                    covariate_regressor.metadata = self.metadata
+                    return covariate_regressor
+                else:
+                    raise ValueError(
+                        f"Invalid value for covariate_regressor {covariate_regressor} of type {type(covariate_regressor)}"
+                    )
+        else:
+            return None
+
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
@@ -324,11 +371,19 @@ class AbstractTimeSeriesModel(AbstractModel):
         """
         if self.target_scaler is not None:
             data = self.target_scaler.fit_transform(data)
+        if self.covariate_regressor is not None:
+            data = self.covariate_regressor.fit_transform(data)
 
         data = self.preprocess(data, is_train=False)
         known_covariates = self.preprocess_known_covariates(known_covariates)
+
+        # FIXME: Set self.covariate_regressor=None so to avoid copying it across processes during _predict
+        # FIXME: The clean solution is to convert all methods executed in parallel to @classmethod
+        covariate_regressor = self.covariate_regressor
+        self.covariate_regressor = None
         predictions = self._predict(data=data, known_covariates=known_covariates, **kwargs)
-        logger.debug(f"Predicting with model {self.name}")
+        self.covariate_regressor = covariate_regressor
+
         # "0.5" might be missing from the quantiles if self is a wrapper (MultiWindowBacktestingModel or ensemble)
         if "0.5" in predictions.columns:
             if self.eval_metric.optimized_by_median:
@@ -336,6 +391,19 @@ class AbstractTimeSeriesModel(AbstractModel):
             if self.must_drop_median:
                 predictions = predictions.drop("0.5", axis=1)
 
+        if self.covariate_regressor is not None:
+            if known_covariates is None:
+                forecast_index = get_forecast_horizon_index_ts_dataframe(
+                    data, prediction_length=self.prediction_length, freq=self.freq
+                )
+                known_covariates = pd.DataFrame(index=forecast_index, dtype="float32")
+
+            predictions = self.covariate_regressor.inverse_transform(
+                predictions,
+                known_covariates=known_covariates,
+                static_features=data.static_features,
+            )
+
         if self.target_scaler is not None:
             predictions = self.target_scaler.inverse_transform(predictions)
         return predictions

autogluon/timeseries/models/chronos/model.py

@@ -9,9 +9,9 @@ from autogluon.common.loaders import load_pkl
 from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.utils.forecast import get_forecast_horizon_index_ts_dataframe
-from autogluon.timeseries.utils.warning_filters import warning_filter
+from autogluon.timeseries.utils.warning_filters import disable_duplicate_logs, warning_filter
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("autogluon.timeseries.models.chronos")
 
 
 # allowed HuggingFace model paths with custom parameter definitions
@@ -41,6 +41,21 @@ MODEL_CONFIGS = {
         "default_torch_dtype": "bfloat16",
         "default_batch_size": 8,
     },
+    "chronos-bolt-mini": {
+        "num_gpus": 0,
+        "default_torch_dtype": "auto",
+        "default_batch_size": 256,
+    },
+    "chronos-bolt-small": {
+        "num_gpus": 0,
+        "default_torch_dtype": "auto",
+        "default_batch_size": 256,
+    },
+    "chronos-bolt-base": {
+        "num_gpus": 0,
+        "default_torch_dtype": "auto",
+        "default_batch_size": 256,
+    },
 }
 
 
@@ -50,22 +65,29 @@ MODEL_ALIASES = {
     "small": "autogluon/chronos-t5-small",
     "base": "autogluon/chronos-t5-base",
     "large": "autogluon/chronos-t5-large",
+    "bolt-mini": "autogluon/chronos-bolt-mini",
+    "bolt-small": "autogluon/chronos-bolt-small",
+    "bolt-base": "autogluon/chronos-bolt-base",
 }
 
 
 class ChronosModel(AbstractTimeSeriesModel):
-    """Chronos pretrained time series forecasting models, based on the original
-    `ChronosModel <https://github.com/amazon-science/chronos-forecasting/blob/main/src/chronos/chronos.py>`_ implementation.
+    """Chronos pretrained time series forecasting models. Models can be based on the original
+    `ChronosModel <https://github.com/amazon-science/chronos-forecasting/blob/main/src/chronos/chronos.py>`_ implementation,
+    as well as a newer family of Chronos-Bolt models which are capable of much faster inference.
 
-    Chronos is family of pretrained models, based on the T5 family, with number of parameters ranging between 8M and 710M.
-    The full collection of Chronos models is available on
+    The original Chronos is a family of pretrained models, based on the T5 family, with number of parameters ranging between
+    8M and 710M. The full collection of Chronos models is available on
     `Hugging Face <https://huggingface.co/collections/amazon/chronos-models-65f1791d630a8d57cb718444>`_. For Chronos small,
-    base, and large variants a GPU is required to perform inference efficiently.
-
-    Chronos takes a minimalistic approach to pretraining time series models, by discretizing time series data directly into bins
-    which are treated as tokens, effectively performing regression by classification. This results in a simple and flexible framework
+    base, and large variants a GPU is required to perform inference efficiently. Chronos takes a minimalistic approach to
+    pretraining time series models, by discretizing time series data directly into bins which are treated as tokens,
+    effectively performing regression by classification. This results in a simple and flexible framework
     for using any language model in the context of time series forecasting. See [Ansari2024]_ for more information.
 
+    The newer Chronos-Bolt variants enable much faster inference by first "patching" the time series. The resulting
+    time series is then fed into a T5 model for forecasting. The Chronos-Bolt variants are capable of much faster inference,
+    and can all run on CPUs. Chronos-Bolt models are also available on Hugging Face <https://huggingface.co/autogluon/>`_.
+
     References
     ----------
     .. [Ansari2024] Ansari, Abdul Fatir, Stella, Lorenzo et al.
@@ -79,7 +101,8 @@ class ChronosModel(AbstractTimeSeriesModel):
         Model path used for the model, i.e., a HuggingFace transformers ``name_or_path``. Can be a
         compatible model name on HuggingFace Hub or a local path to a model directory. Original
         Chronos models (i.e., ``autogluon/chronos-t5-{model_size}``) can be specified with aliases
-        ``tiny``, ``mini`` , ``small``, ``base``, and ``large``.
+        ``tiny``, ``mini`` , ``small``, ``base``, and ``large``. Chronos-Bolt models can be specified
+        with ``bolt-mini``, ``bolt-small``, and ``bolt-base``.
     batch_size : int, default = 16
         Size of batches used during inference
     num_samples : int, default = 20
@@ -90,11 +113,15 @@ class ChronosModel(AbstractTimeSeriesModel):
     context_length : int or None, default = None
         The context length to use in the model. Shorter context lengths will decrease model accuracy, but result
         in faster inference. If None, the model will infer context length from the data set length at inference
-        time, but set it to a maximum of 512.
+        time, but set it to a maximum of 2048. Note that this is only the context length used to pass data into
+        the model. Individual model implementations may have different context lengths specified in their configuration,
+        and may truncate the context further. For example, original Chronos models have a context length of 512, but
+        Chronos-Bolt models handle contexts up to 2048.
     optimization_strategy : {None, "onnx", "openvino"}, default = None
         Optimization strategy to use for inference on CPUs. If None, the model will use the default implementation.
         If `onnx`, the model will be converted to ONNX and the inference will be performed using ONNX. If ``openvino``,
-        inference will be performed with the model compiled to OpenVINO.
+        inference will be performed with the model compiled to OpenVINO. These optimizations are only available for
+        the original set of Chronos models, and not in Chronos-Bolt where they are not needed.
     torch_dtype : torch.dtype or {"auto", "bfloat16", "float32", "float64"}, default = "auto"
         Torch data type for model weights, provided to ``from_pretrained`` method of Hugging Face AutoModels. If
         original Chronos models are specified and the model size is ``small``, ``base``, or ``large``, the
@@ -107,7 +134,7 @@ class ChronosModel(AbstractTimeSeriesModel):
     # default number of samples for prediction
     default_num_samples: int = 20
     default_model_path = "autogluon/chronos-t5-small"
-    maximum_context_length = 512
+    maximum_context_length = 2048
 
     def __init__(
         self,
@@ -159,7 +186,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             **kwargs,
         )
 
-        self.model_pipeline: Optional[Any] = None  # of type OptimizedChronosPipeline
+        self.model_pipeline: Optional[Any] = None  # of type BaseChronosPipeline
         self.time_limit: Optional[float] = None
 
     def save(self, path: str = None, verbose: bool = True) -> str:
@@ -218,8 +245,8 @@ class ChronosModel(AbstractTimeSeriesModel):
             minimum_resources["num_gpus"] = self.min_num_gpus
         return minimum_resources
 
-    def load_model_pipeline(self, context_length: Optional[int] = None):
-        from .pipeline import OptimizedChronosPipeline
+    def load_model_pipeline(self):
+        from .pipeline import BaseChronosPipeline
 
         gpu_available = self._is_gpu_available()
 
@@ -232,18 +259,17 @@ class ChronosModel(AbstractTimeSeriesModel):
 
         device = self.device or ("cuda" if gpu_available else "cpu")
 
-        pipeline = OptimizedChronosPipeline.from_pretrained(
+        pipeline = BaseChronosPipeline.from_pretrained(
             self.model_path,
             device_map=device,
-            optimization_strategy=self.optimization_strategy,
             torch_dtype=self.torch_dtype,
-            context_length=context_length or self.context_length,
+            optimization_strategy=self.optimization_strategy,
         )
 
         self.model_pipeline = pipeline
 
     def persist(self) -> "ChronosModel":
-        self.load_model_pipeline(context_length=self.context_length or self.maximum_context_length)
+        self.load_model_pipeline()
         return self
 
     def _fit(
@@ -263,7 +289,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         num_workers: int = 0,
         time_limit: Optional[float] = None,
     ):
-        from .utils import ChronosInferenceDataLoader, ChronosInferenceDataset, timeout_callback
+        from .pipeline.utils import ChronosInferenceDataLoader, ChronosInferenceDataset, timeout_callback
 
         chronos_dataset = ChronosInferenceDataset(
             target_df=data,
@@ -290,6 +316,9 @@ class ChronosModel(AbstractTimeSeriesModel):
         # and use that to determine the context length of the model. If the context length is specified
         # during initialization, this is always used. If not, the context length is set to the longest
         # item length. The context length is always capped by self.maximum_context_length.
+        # Note that this is independent of the model's own context length set in the model's config file.
+        # For example, if the context_length is set to 2048 here but the model expects context length
+        # (according to its config.json file) of 512, it will further truncate the series during inference.
         context_length = self.context_length or min(
             data.num_timesteps_per_item().max(),
             self.maximum_context_length,
@@ -300,7 +329,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
             if self.model_pipeline is None:
                 # load model pipeline to device memory
-                self.load_model_pipeline(context_length=context_length)
+                self.load_model_pipeline()
 
             inference_data_loader = self._get_inference_data_loader(
                 data=data,
@@ -308,28 +337,28 @@ class ChronosModel(AbstractTimeSeriesModel):
                 context_length=context_length,
                 time_limit=kwargs.get("time_limit"),
             )
+
             self.model_pipeline.model.eval()
-            with torch.inference_mode():
-                prediction_samples = [
-                    self.model_pipeline.predict(
+            with torch.inference_mode(), disable_duplicate_logs(logger):
+                batch_quantiles, batch_means = [], []
+                for batch in inference_data_loader:
+                    qs, mn = self.model_pipeline.predict_quantiles(
                         batch,
                         prediction_length=self.prediction_length,
+                        quantile_levels=self.quantile_levels,
                         num_samples=self.num_samples,
-                        limit_prediction_length=False,
                     )
-                    .detach()
-                    .cpu()
-                    .numpy()
-                    for batch in inference_data_loader
-                ]
-
-        samples = np.concatenate(prediction_samples, axis=0).swapaxes(1, 2).reshape(-1, self.num_samples)
-
-        mean = samples.mean(axis=-1, keepdims=True)
-        quantiles = np.quantile(samples, self.quantile_levels, axis=-1).T
+                    batch_quantiles.append(qs.numpy())
+                    batch_means.append(mn.numpy())
 
         df = pd.DataFrame(
-            np.concatenate([mean, quantiles], axis=1),
+            np.concatenate(
+                [
+                    np.concatenate(batch_means, axis=0).reshape(-1, 1),
+                    np.concatenate(batch_quantiles, axis=0).reshape(-1, len(self.quantile_levels)),
+                ],
+                axis=1,
+            ),
             columns=["mean"] + [str(q) for q in self.quantile_levels],
             index=get_forecast_horizon_index_ts_dataframe(data, self.prediction_length, freq=self.freq),
         )

autogluon/timeseries/models/chronos/pipeline/__init__.py

@@ -0,0 +1,11 @@
+from .chronos import ChronosPipeline
+from .chronos_bolt import ChronosBoltPipeline
+from .base import BaseChronosPipeline, ForecastType
+
+
+__all__ = [
+    "BaseChronosPipeline",
+    "ChronosBoltPipeline",
+    "ChronosPipeline",
+    "ForecastType",
+]

autogluon/timeseries/models/chronos/pipeline/base.py

@@ -0,0 +1,146 @@
+# Authors: Lorenzo Stella <stellalo@amazon.com>, Caner Turkmen <atturkm@amazon.com>
+
+from enum import Enum
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+
+from .utils import left_pad_and_stack_1D
+
+
+class ForecastType(Enum):
+    SAMPLES = "samples"
+    QUANTILES = "quantiles"
+
+
+class PipelineRegistry(type):
+    REGISTRY: Dict[str, "PipelineRegistry"] = {}
+
+    def __new__(cls, name, bases, attrs):
+        """See, https://github.com/faif/python-patterns."""
+        new_cls = type.__new__(cls, name, bases, attrs)
+        if name is not None:
+            cls.REGISTRY[name] = new_cls
+        if aliases := attrs.get("_aliases"):
+            for alias in aliases:
+                cls.REGISTRY[alias] = new_cls
+        return new_cls
+
+
+class BaseChronosPipeline(metaclass=PipelineRegistry):
+    forecast_type: ForecastType
+    dtypes = {
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+        "float64": torch.float64,
+    }
+
+    def _prepare_and_validate_context(self, context: Union[torch.Tensor, List[torch.Tensor]]):
+        if isinstance(context, list):
+            context = left_pad_and_stack_1D(context)
+        assert isinstance(context, torch.Tensor)
+        if context.ndim == 1:
+            context = context.unsqueeze(0)
+        assert context.ndim == 2
+
+        return context
+
+    def predict(
+        self,
+        context: Union[torch.Tensor, List[torch.Tensor]],
+        prediction_length: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Get forecasts for the given time series.
+
+        Parameters
+        ----------
+        context
+            Input series. This is either a 1D tensor, or a list
+            of 1D tensors, or a 2D tensor whose first dimension
+            is batch. In the latter case, use left-padding with
+            ``torch.nan`` to align series of different lengths.
+        prediction_length
+            Time steps to predict. Defaults to a model-dependent
+            value if not given.
+
+        Returns
+        -------
+        forecasts
+            Tensor containing forecasts. The layout and meaning
+            of the forecasts values depends on ``self.forecast_type``.
+        """
+        raise NotImplementedError()
+
+    def predict_quantiles(
+        self, context: torch.Tensor, prediction_length: int, quantile_levels: List[float], **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Get quantile and mean forecasts for given time series. All
+        predictions are returned on the CPU.
+
+        Parameters
+        ----------
+        context
+            Input series. This is either a 1D tensor, or a list
+            of 1D tensors, or a 2D tensor whose first dimension
+            is batch. In the latter case, use left-padding with
+            ``torch.nan`` to align series of different lengths.
+        prediction_length
+            Time steps to predict. Defaults to a model-dependent
+            value if not given.
+        quantile_levels: List[float]
+            Quantile levels to compute
+
+        Returns
+        -------
+        quantiles
+            Tensor containing quantile forecasts. Shape
+            (batch_size, prediction_length, num_quantiles)
+        mean
+            Tensor containing mean (point) forecasts. Shape
+            (batch_size, prediction_length)
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+        *model_args,
+        force=False,
+        **kwargs,
+    ):
+        """
+        Load the model, either from a local path or from the HuggingFace Hub.
+        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
+        from ``transformers``.
+
+        When a local path is provided, supports both a folder or a .tar.gz archive.
+        """
+        from transformers import AutoConfig
+
+        if str(pretrained_model_name_or_path).startswith("s3://"):
+            from .utils import cache_model_from_s3
+
+            local_model_path = cache_model_from_s3(str(pretrained_model_name_or_path), force=force)
+            return cls.from_pretrained(local_model_path, *model_args, **kwargs)
+
+        torch_dtype = kwargs.get("torch_dtype", "auto")
+        if torch_dtype != "auto" and isinstance(torch_dtype, str):
+            kwargs["torch_dtype"] = cls.dtypes[torch_dtype]
+
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        is_valid_config = hasattr(config, "chronos_pipeline_class") or hasattr(config, "chronos_config")
+
+        if not is_valid_config:
+            raise ValueError("Not a Chronos config file")
+
+        pipeline_class_name = getattr(config, "chronos_pipeline_class", "ChronosPipeline")
+        class_ = PipelineRegistry.REGISTRY.get(pipeline_class_name)
+        if class_ is None:
+            raise ValueError(f"Trying to load unknown pipeline class: {pipeline_class_name}")
+
+        return class_.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)