autogluon.timeseries 1.0.1b20240323__tar.gz → 1.0.1b20240325__tar.gz

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.timeseries
-Version: 1.0.1b20240323
+Version: 1.0.1b20240325
 Summary: AutoML for Image, Text, and Tabular Data
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community
@@ -132,5 +132,6 @@ Classifier: Topic :: Scientific/Engineering :: Image Recognition
 Requires-Python: >=3.8, <3.12
 Description-Content-Type: text/markdown
 Provides-Extra: tests
-Provides-Extra: chronos-cpu
+Provides-Extra: chronos-openvino
+Provides-Extra: chronos-onnx
 Provides-Extra: all

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/setup.py

@@ -55,12 +55,17 @@ extras_require = {
         "isort>=5.10",
         "black~=23.0",
     ],
-    "chronos-cpu": [  # for faster CPU inference in pretrained models
-        "optimum[onnxruntime,openvino,nncf]>=1.17,<1.18",
+    "chronos-openvino": [  # for faster CPU inference in pretrained models with OpenVINO
+        "optimum[openvino,nncf]>=1.17,<1.18",
+    ],
+    "chronos-onnx": [  # for faster CPU inference in pretrained models with ONNX
+        "optimum[onnxruntime]>=1.17,<1.18",
     ],
 }
 
-extras_require["all"] = list(set.union(*(set(extras_require[extra]) for extra in ["chronos-cpu"])))
+extras_require["all"] = list(
+    set.union(*(set(extras_require[extra]) for extra in ["chronos-onnx", "chronos-openvino"]))
+)
 
 install_requires = ag.get_dependency_version_ranges(install_requires)
 

autogluon.timeseries-1.0.1b20240325/src/autogluon/timeseries/configs/presets_configs.py

@@ -0,0 +1,55 @@
+"""Preset configurations for autogluon.timeseries Predictors"""
+from autogluon.timeseries.models.presets import get_default_hps
+
+# TODO: change default HPO settings when other HPO strategies (e.g., Ray tune) are available
+# TODO: add refit_full arguments once refitting is available
+
+TIMESERIES_PRESETS_CONFIGS = dict(
+    best_quality={"hyperparameters": "default", "num_val_windows": 2},
+    high_quality={"hyperparameters": "default"},
+    medium_quality={"hyperparameters": "light"},
+    fast_training={"hyperparameters": "very_light"},
+    chronos_tiny={
+        "hyperparameters": {"Chronos": {"model_path": "tiny"}},
+    },
+    chronos_mini={
+        "hyperparameters": {"Chronos": {"model_path": "mini"}},
+    },
+    chronos_small={
+        "hyperparameters": {"Chronos": {"model_path": "small"}},
+    },
+    chronos_base={
+        "hyperparameters": {"Chronos": {"model_path": "base"}},
+    },
+    chronos_large={
+        "hyperparameters": {"Chronos": {"model_path": "large", "batch_size": 8}},
+    },
+    chronos_ensemble={
+        "hyperparameters": {
+            "Chronos": {"model_path": "small"},
+            **get_default_hps("default"),
+        }
+    },
+    chronos_large_ensemble={
+        "hyperparameters": {
+            "Chronos": {"model_path": "large", "batch_size": 8},
+            **get_default_hps("default"),
+        }
+    },
+)
+
+TIMESERIES_PRESETS_ALIASES = dict(
+    chronos="chronos_small",
+    best="best_quality",
+    high="high_quality",
+    medium="medium_quality",
+    bq="best_quality",
+    hq="high_quality",
+    mq="medium_quality",
+)
+
+# update with aliases
+TIMESERIES_PRESETS_CONFIGS = {
+    **TIMESERIES_PRESETS_CONFIGS,
+    **{k: TIMESERIES_PRESETS_CONFIGS[v].copy() for k, v in TIMESERIES_PRESETS_ALIASES.items()},
+}

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon/timeseries/models/__init__.py

@@ -1,4 +1,5 @@
 from .autogluon_tabular import DirectTabularModel, RecursiveTabularModel
+from .chronos import ChronosModel
 from .gluonts import (
     DeepARModel,
     DLinearModel,
@@ -44,6 +45,7 @@ __all__ = [
     "DynamicOptimizedThetaModel",
     "ETSModel",
     "IMAPAModel",
+    "ChronosModel",
     "NPTSModel",
     "NaiveModel",
     "PatchTSTModel",

autogluon.timeseries-1.0.1b20240325/src/autogluon/timeseries/models/chronos/__init__.py

@@ -0,0 +1,3 @@
+from .model import ChronosModel
+
+__all__ = ["ChronosModel"]

autogluon.timeseries-1.0.1b20240325/src/autogluon/timeseries/models/chronos/chronos.py

@@ -0,0 +1,487 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Original Source: https://github.com/amazon-science/chronos-forecasting
+# Author: Lorenzo Stella <stellalo@amazon.com>
+
+import logging
+import warnings
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, GenerationConfig, PreTrainedModel
+
+from autogluon.timeseries.utils.warning_filters import set_loggers_level
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChronosConfig:
+    """
+    This class holds all the configuration parameters to be used
+    by ``ChronosTokenizer`` and ``ChronosPretrainedModel``.
+    """
+
+    tokenizer_class: str
+    tokenizer_kwargs: Dict[str, Any]
+    n_tokens: int
+    n_special_tokens: int
+    pad_token_id: int
+    eos_token_id: int
+    use_eos_token: bool
+    model_type: Literal["causal", "seq2seq"]
+    context_length: int
+    prediction_length: int
+    num_samples: int
+    temperature: float
+    top_k: int
+    top_p: float
+
+    def __post_init__(self):
+        assert (
+            self.pad_token_id < self.n_special_tokens and self.eos_token_id < self.n_special_tokens
+        ), f"Special token id's must be smaller than {self.n_special_tokens=}"
+
+    def create_tokenizer(self) -> "ChronosTokenizer":
+        if self.tokenizer_class == "MeanScaleUniformBins":
+            return MeanScaleUniformBins(**self.tokenizer_kwargs, config=self)
+        raise ValueError
+
+
+class ChronosTokenizer:
+    """
+    A ``ChronosTokenizer`` defines how time series are mapped into token IDs
+    and back.
+
+    For details, see the ``input_transform`` and ``output_transform`` methods,
+    which concrete classes must implement.
+    """
+
+    def input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Any]:
+        """
+        Turn a batch of time series into token IDs, attention map, and scale.
+
+        Parameters
+        ----------
+        context
+            A tensor shaped (batch_size, time_length), containing the
+            timeseries to forecast. Use left-padding with ``torch.nan``
+            to align time series of different lengths.
+
+        Returns
+        -------
+        token_ids
+            A tensor of integers, shaped (batch_size, time_length + 1)
+            if ``config.use_eos_token`` and (batch_size, time_length)
+            otherwise, containing token IDs for the input series.
+        attention_mask
+            A boolean tensor, same shape as ``token_ids``, indicating
+            which input observations are not ``torch.nan`` (i.e. not
+            missing nor padding).
+        decoding_context
+            An object that will be passed to ``output_transform``.
+            Contains the relevant context to decode output samples into
+            real values, such as location and scale parameters.
+        """
+        raise NotImplementedError()
+
+    def output_transform(self, samples: torch.Tensor, decoding_context: Any) -> torch.Tensor:
+        """
+        Turn a batch of sample token IDs into real values.
+
+        Parameters
+        ----------
+        samples
+            A tensor of integers, shaped (batch_size, num_samples, time_length),
+            containing token IDs of sample trajectories.
+        decoding_context
+            An object returned by ``input_transform`` containing
+            relevant context to decode samples, such as location and scale.
+            The nature of this depends on the specific tokenizer.
+
+        Returns
+        -------
+        forecasts
+            A real tensor, shaped (batch_size, num_samples, time_length),
+            containing forecasted sample paths.
+        """
+        raise NotImplementedError()
+
+
+class MeanScaleUniformBins(ChronosTokenizer):
+    def __init__(self, low_limit: float, high_limit: float, config: ChronosConfig) -> None:
+        self.config = config
+        self.centers = torch.linspace(
+            low_limit,
+            high_limit,
+            config.n_tokens - config.n_special_tokens - 1,
+        )
+        self.boundaries = torch.concat(
+            (
+                torch.tensor([-1e20], device=self.centers.device),
+                (self.centers[1:] + self.centers[:-1]) / 2,
+                torch.tensor([1e20], device=self.centers.device),
+            )
+        )
+
+    def input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, length = context.shape
+
+        if length > self.config.context_length:
+            context = context[..., -self.config.context_length :]
+        elif length < self.config.context_length:
+            padding_size = (
+                *context.shape[:-1],
+                self.config.context_length - length,
+            )
+            padding = torch.full(size=padding_size, fill_value=torch.nan)
+            context = torch.concat((padding, context), dim=-1)
+
+        attention_mask = ~torch.isnan(context)
+        scale = torch.nansum(torch.abs(context) * attention_mask, dim=-1) / torch.nansum(attention_mask, dim=-1)
+        scale[~(scale > 0)] = 1.0
+        scaled_context = context / scale.unsqueeze(dim=-1)
+        token_ids = (
+            torch.bucketize(
+                input=scaled_context,
+                boundaries=self.boundaries,
+                # buckets are open to the right, see:
+                # https://pytorch.org/docs/2.1/generated/torch.bucketize.html#torch-bucketize
+                right=True,
+            )
+            + self.config.n_special_tokens
+        )
+        token_ids[~attention_mask] = self.config.pad_token_id
+
+        if self.config.use_eos_token:
+            eos_tokens = torch.full((batch_size, 1), fill_value=self.config.eos_token_id)
+            token_ids = torch.concat((token_ids, eos_tokens), dim=1)
+            eos_mask = torch.full((batch_size, 1), fill_value=True)
+            attention_mask = torch.concat((attention_mask, eos_mask), dim=1)
+
+        return token_ids, attention_mask, scale
+
+    def output_transform(self, samples: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+        scale_unsqueezed = scale.unsqueeze(-1).unsqueeze(-1)
+        indices = torch.clamp(
+            samples - self.config.n_special_tokens,
+            min=0,
+            max=len(self.centers) - 1,
+        )
+        return self.centers[indices] * scale_unsqueezed
+
+
+class ChronosPretrainedModel(nn.Module):
+    """
+    A ``ChronosPretrainedModel`` wraps a ``PreTrainedModel`` object from ``transformers``
+    and uses it to predict sample paths for time series tokens.
+
+    Parameters
+    ----------
+    config
+        The configuration to use.
+    model
+        The pre-trained model to use.
+    """
+
+    def __init__(self, config: ChronosConfig, model: PreTrainedModel) -> None:
+        super().__init__()
+        self.config = config
+        self.model = model
+        self.device = model.device
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        prediction_length: Optional[int] = None,
+        num_samples: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+    ) -> torch.Tensor:
+        """
+        Predict future sample tokens for the given token sequences.
+
+        Arguments ``prediction_length``, ``num_samples``, ``temperature``,
+        ``top_k``, ``top_p`` can be used to customize the model inference,
+        and default to the corresponding attributes in ``self.config`` if
+        not provided.
+
+        Returns
+        -------
+        samples
+            A tensor of integers, shaped (batch_size, num_samples, time_length),
+            containing forecasted sample paths.
+        """
+        if prediction_length is None:
+            prediction_length = self.config.prediction_length
+        if num_samples is None:
+            num_samples = self.config.num_samples
+        if temperature is None:
+            temperature = self.config.temperature
+        if top_k is None:
+            top_k = self.config.top_k
+        if top_p is None:
+            top_p = self.config.top_p
+
+        preds = self.model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask.long(),  # int64 (long) type conversion needed for ONNX
+            generation_config=GenerationConfig(
+                min_new_tokens=prediction_length,
+                max_new_tokens=prediction_length,
+                do_sample=True,
+                num_return_sequences=num_samples,
+                eos_token_id=self.config.eos_token_id,
+                pad_token_id=self.config.pad_token_id,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+            ),
+        )
+
+        if self.config.model_type == "seq2seq":
+            preds = preds[..., 1:]  # remove the decoder start token
+        else:
+            assert self.config.model_type == "causal"
+            assert preds.size(-1) == input_ids.size(-1) + prediction_length
+            preds = preds[..., -prediction_length:]
+
+        return preds.reshape(input_ids.size(0), num_samples, -1)
+
+
+def left_pad_and_stack_1D(tensors: List[torch.Tensor]):
+    max_len = max(len(c) for c in tensors)
+    padded = []
+    for c in tensors:
+        assert isinstance(c, torch.Tensor)
+        assert c.ndim == 1
+        padding = torch.full(size=(max_len - len(c),), fill_value=torch.nan, device=c.device)
+        padded.append(torch.concat((padding, c), dim=-1))
+    return torch.stack(padded)
+
+
+class ChronosPipeline:
+    """
+    A ``ChronosPipeline`` uses the given tokenizer and model to forecast
+    input time series.
+
+    Use the ``from_pretrained`` class method to load serialized models.
+    Use the ``predict`` method to get forecasts.
+
+    Parameters
+    ----------
+    tokenizer
+        The tokenizer object to use.
+    model
+        The model to use.
+    """
+
+    tokenizer: ChronosTokenizer
+    model: ChronosPretrainedModel
+
+    def __init__(self, tokenizer, model):
+        self.tokenizer = tokenizer
+        self.model = model
+
+    def predict(
+        self,
+        context: Union[torch.Tensor, List[torch.Tensor]],
+        prediction_length: Optional[int] = None,
+        num_samples: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        limit_prediction_length: bool = True,
+    ) -> torch.Tensor:
+        """
+        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 what specified
+            in ``self.model.config``.
+        num_samples
+            Number of sample paths to predict. Defaults to what
+            specified in ``self.model.config``.
+        temperature
+            Temperature to use for generating sample tokens.
+            Defaults to what specified in ``self.model.config``.
+        top_k
+            Top-k parameter to use for generating sample tokens.
+            Defaults to what specified in ``self.model.config``.
+        top_p
+            Top-p parameter to use for generating sample tokens.
+            Defaults to what specified in ``self.model.config``.
+        limit_prediction_length
+            Force prediction length smaller or equal than the
+            built-in prediction length from the model. True by
+            default. When true, fail loudly if longer predictions
+            are requested, otherwise longer predictions are allowed.
+
+        Returns
+        -------
+        samples
+            Tensor of sample forecasts, of shape
+            (batch_size, num_samples, prediction_length).
+        """
+        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
+
+        if prediction_length is None:
+            prediction_length = self.model.config.prediction_length
+
+        if prediction_length > self.model.config.prediction_length:
+            msg = (
+                f"We recommend keeping prediction length <= {self.model.config.prediction_length}. "
+                f"The quality of longer predictions may degrade since the model is not optimized for it. "
+            )
+            if limit_prediction_length:
+                msg += "You can turn off this check by setting `limit_prediction_length=False`."
+                raise ValueError(msg)
+            warnings.warn(msg, stacklevel=2)
+
+        predictions = []
+        remaining = prediction_length
+
+        while remaining > 0:
+            token_ids, attention_mask, scale = self.tokenizer.input_transform(context)
+            samples = self.model(
+                token_ids.to(self.model.device),
+                attention_mask.to(self.model.device),
+                min(remaining, self.model.config.prediction_length),
+                num_samples,
+                temperature,
+                top_k,
+                top_p,
+            )
+            prediction = self.tokenizer.output_transform(samples.to(scale.device), scale)
+
+            predictions.append(prediction)
+            remaining -= prediction.shape[-1]
+
+            if remaining <= 0:
+                break
+
+            context = torch.cat([context, prediction.median(dim=1).values], dim=-1)
+
+        return torch.cat(predictions, dim=-1)
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """
+        Load the model, either from a local path or from the HuggingFace Hub.
+        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
+        from ``transformers``.
+        """
+
+        config = AutoConfig.from_pretrained(*args, **kwargs)
+
+        assert hasattr(config, "chronos_config"), "Not a Chronos config file"
+
+        chronos_config = ChronosConfig(**config.chronos_config)
+
+        if chronos_config.model_type == "seq2seq":
+            inner_model = AutoModelForSeq2SeqLM.from_pretrained(*args, **kwargs)
+        else:
+            assert config.model_type == "causal"
+            inner_model = AutoModelForCausalLM.from_pretrained(*args, **kwargs)
+
+        return cls(
+            tokenizer=chronos_config.create_tokenizer(),
+            model=ChronosPretrainedModel(config=chronos_config, model=inner_model),
+        )
+
+
+class OptimizedChronosPipeline(ChronosPipeline):
+    """A wrapper around the ChronosPipeline object for CPU-optimized model classes from
+    HuggingFace optimum.
+    """
+
+    dtypes = {
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+        "float64": torch.float64,
+    }
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """
+        Load the model, either from a local path or from the HuggingFace Hub.
+        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
+        from ``transformers``.
+        """
+        kwargs = kwargs.copy()
+
+        optimization_strategy = kwargs.pop("optimization_strategy", None)
+        context_length = kwargs.pop("context_length", None)
+
+        config = AutoConfig.from_pretrained(*args, **kwargs)
+        assert hasattr(config, "chronos_config"), "Not a Chronos config file"
+
+        if context_length is not None:
+            config.chronos_config["context_length"] = context_length
+        chronos_config = ChronosConfig(**config.chronos_config)
+
+        torch_dtype = kwargs.get("torch_dtype", "auto")
+        if torch_dtype != "auto" and isinstance(torch_dtype, str):
+            kwargs["torch_dtype"] = cls.dtypes[torch_dtype]
+
+        if chronos_config.model_type == "seq2seq":
+            if optimization_strategy is None:
+                inner_model = AutoModelForSeq2SeqLM.from_pretrained(*args, **kwargs)
+            else:
+                assert optimization_strategy in [
+                    "onnx",
+                    "openvino",
+                ], "optimization_strategy not recognized. Please provide one of `onnx` or `openvino`"
+                torch_dtype = kwargs.pop("torch_dtype", "auto")
+                if torch_dtype != "auto":
+                    logger.warning(
+                        f"\t`torch_dtype` will be ignored for optimization_strategy {optimization_strategy}"
+                    )
+
+                if optimization_strategy == "onnx":
+                    try:
+                        from optimum.onnxruntime import ORTModelForSeq2SeqLM
+                    except ImportError:
+                        raise ImportError(
+                            "Huggingface Optimum library must be installed with ONNX for using the `onnx` strategy"
+                        )
+
+                    assert kwargs.pop("device_map", "cpu") in ["cpu", "auto"], "ONNX mode only available on the CPU"
+                    with set_loggers_level(regex=r"^optimum.*", level=logging.ERROR):
+                        inner_model = ORTModelForSeq2SeqLM.from_pretrained(*args, **{**kwargs, "export": True})
+                elif optimization_strategy == "openvino":
+                    try:
+                        from optimum.intel import OVModelForSeq2SeqLM
+                    except ImportError:
+                        raise ImportError(
+                            "Huggingface Optimum library must be installed with OpenVINO for using the `openvino` strategy"
+                        )
+                    with set_loggers_level(regex=r"^optimum.*", level=logging.ERROR):
+                        inner_model = OVModelForSeq2SeqLM.from_pretrained(
+                            *args, **{**kwargs, "device_map": "cpu", "export": True}
+                        )
+        else:
+            assert config.model_type == "causal"
+            inner_model = AutoModelForCausalLM.from_pretrained(*args, **kwargs)
+
+        return cls(
+            tokenizer=chronos_config.create_tokenizer(),
+            model=ChronosPretrainedModel(config=chronos_config, model=inner_model),
+        )

autogluon.timeseries-1.0.1b20240325/src/autogluon/timeseries/models/chronos/model.py

@@ -0,0 +1,319 @@
+import logging
+import os
+from typing import Any, Dict, Literal, Optional, Union
+
+import numpy as np
+import pandas as pd
+
+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
+
+logger = logging.getLogger(__name__)
+
+
+# allowed HuggingFace model paths with custom parameter definitions
+MODEL_CONFIGS = {
+    "amazon/chronos-t5-tiny": {
+        "num_gpus": 0,  # minimum number of required GPUs
+    },
+    "amazon/chronos-t5-mini": {"num_gpus": 0},
+    "amazon/chronos-t5-small": {"num_gpus": 1},
+    "amazon/chronos-t5-base": {"num_gpus": 1},
+    "amazon/chronos-t5-large": {"num_gpus": 1},
+}
+
+
+MODEL_ALIASES = {
+    "tiny": "amazon/chronos-t5-tiny",
+    "mini": "amazon/chronos-t5-mini",
+    "small": "amazon/chronos-t5-small",
+    "base": "amazon/chronos-t5-base",
+    "large": "amazon/chronos-t5-large",
+}
+
+
+class ChronosInferenceDataset:
+    """A container for time series datasets that implements the ``torch.utils.data.Dataset`` interface"""
+
+    def __init__(
+        self,
+        target_df: TimeSeriesDataFrame,
+        context_length: int,
+        target_column: str = "target",
+    ):
+        assert context_length > 0
+        self.context_length = context_length
+        self.target_array = target_df[target_column].to_numpy(dtype=np.float32)
+        self.freq = target_df.freq
+
+        # store pointer to start:end of each time series
+        cum_sizes = target_df.num_timesteps_per_item().values.cumsum()
+        self.indptr = np.append(0, cum_sizes).astype(np.int32)
+
+    def __len__(self):
+        return len(self.indptr) - 1  # noqa
+
+    def _get_context(self, a: np.ndarray, pad_value=np.nan):
+        a = a[-self.context_length :]
+        pad_size = self.context_length - len(a)
+        if pad_size > 0:
+            pad = np.full(shape=(pad_size,), fill_value=pad_value)
+            a = np.concatenate((pad, a))
+        return a
+
+    def __getitem__(self, idx) -> np.ndarray:
+        start_idx = self.indptr[idx]
+        end_idx = self.indptr[idx + 1]
+
+        return self._get_context(self.target_array[start_idx:end_idx])
+
+
+class ChronosModel(AbstractTimeSeriesModel):
+    """Chronos pretrained time series forecasting models, based on the original
+    `ChronosModel <https://github.com/amazon-science/chronos-forecasting>`_ implementation.
+
+    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
+    `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
+    for using any language model in the context of time series forecasting. See [Ansari2024]_ for more information.
+
+    References
+    ----------
+    .. [Ansari2024] Ansari, Abdul Fatir, Stella, Lorenzo et al.
+        "Chronos: Learning the Language of Time Series."
+        http://arxiv.org/abs/2403.07815
+
+
+    Other Parameters
+    ----------------
+    model_path: str, default = "amazon/chronos-t5-small"
+        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., ``amazon/chronos-t5-{model_size}``) can be specified with aliases
+        ``tiny``, ``mini`` , ``small``, ``base``, and ``large``.
+    batch_size : int, default = 16
+        Size of batches used during inference
+    num_samples : int, default = 20
+        Number of samples used during inference
+    device : str, default = None
+        Device to use for inference. If None, model will use the GPU if available. For larger model sizes
+        `small`, `base`, and `large`; inference will fail if no GPU is available.
+    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.
+    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.
+    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
+        ``torch_dtype`` will be set to ``bfloat16`` to enable inference on GPUs.
+    data_loader_num_workers : int, default = 0
+        Number of worker processes to be used in the data loader. See documentation on ``torch.utils.data.DataLoader``
+        for more information.
+    """
+
+    # default number of samples for prediction
+    default_num_samples: int = 20
+    default_batch_size: int = 16
+    default_model_path = "amazon/chronos-t5-small"
+    maximum_context_length = 512
+
+    def __init__(
+        self,
+        freq: Optional[str] = None,
+        prediction_length: int = 1,
+        path: Optional[str] = None,
+        name: Optional[str] = None,
+        eval_metric: str = None,
+        hyperparameters: Dict[str, Any] = None,
+        **kwargs,  # noqa
+    ):
+        hyperparameters = hyperparameters if hyperparameters is not None else {}
+
+        model_path_input = hyperparameters.get("model_path", self.default_model_path)
+        self.model_path = MODEL_ALIASES.get(model_path_input, model_path_input)
+
+        # TODO: automatically determine batch size based on GPU / memory availability
+        self.batch_size = hyperparameters.get("batch_size", self.default_batch_size)
+        self.num_samples = hyperparameters.get("num_samples", self.default_num_samples)
+        self.device = hyperparameters.get("device")
+
+        # if the model requires a GPU, set the torch dtype to bfloat16
+        self.torch_dtype = hyperparameters.get("torch_dtype", "auto" if self.min_num_gpus == 0 else "bfloat16")
+
+        self.data_loader_num_workers = hyperparameters.get("data_loader_num_workers", 0)
+        self.optimization_strategy: Optional[Literal["onnx", "openvino"]] = hyperparameters.get(
+            "optimization_strategy", None
+        )
+        self.context_length = hyperparameters.get("context_length")
+
+        if self.context_length is not None and self.context_length > self.maximum_context_length:
+            logger.warning(
+                f"\tContext length {self.context_length} exceeds maximum context length {self.maximum_context_length}."
+                f"Context length will be set to {self.maximum_context_length}."
+            )
+            self.context_length = self.maximum_context_length
+
+        model_path_safe = str.replace(model_path_input, "/", "__")
+        name = (name if name is not None else "Chronos") + f"[{model_path_safe}]"
+
+        super().__init__(
+            path=path,
+            freq=freq,
+            prediction_length=prediction_length,
+            name=name,
+            eval_metric=eval_metric,
+            hyperparameters=hyperparameters,
+            **kwargs,
+        )
+
+        self.model_pipeline: Optional[Any] = None  # of type OptimizedChronosPipeline
+
+    def save(self, path: str = None, verbose: bool = True) -> str:
+        pipeline = self.model_pipeline
+        self.model_pipeline = None
+        path = super().save(path=path, verbose=verbose)
+        self.model_pipeline = pipeline
+
+        return str(path)
+
+    @classmethod
+    def load(cls, path: str, reset_paths: bool = True, verbose: bool = True) -> "ChronosModel":
+        model = load_pkl.load(path=os.path.join(path, cls.model_file_name), verbose=verbose)
+        if reset_paths:
+            model.set_contexts(path)
+        return model
+
+    def _is_gpu_available(self) -> bool:
+        import torch.cuda
+
+        return torch.cuda.is_available()
+
+    @property
+    def min_num_gpus(self):
+        return MODEL_CONFIGS.get(self.model_path, {}).get("num_gpus", 0)
+
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> Dict[str, Union[int, float]]:
+        minimum_resources = {"num_cpus": 1}
+        # if GPU is available, we train with 1 GPU per trial
+        if is_gpu_available:
+            minimum_resources["num_gpus"] = self.min_num_gpus
+        return minimum_resources
+
+    def load_model_pipeline(self, context_length: Optional[int] = None):
+        from .chronos import OptimizedChronosPipeline
+
+        gpu_available = self._is_gpu_available()
+
+        if not gpu_available and self.min_num_gpus > 0:
+            raise RuntimeError(
+                f"{self.name} requires a GPU to run, but no GPU was detected. "
+                "Please make sure that you are using a computer with a CUDA-compatible GPU and "
+                "`import torch; torch.cuda.is_available()` returns `True`."
+            )
+
+        device = self.device or ("cuda" if gpu_available else "auto")
+
+        pipeline = OptimizedChronosPipeline.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,
+        )
+
+        self.model_pipeline = pipeline
+
+    def _fit(
+        self,
+        train_data: TimeSeriesDataFrame,
+        val_data: Optional[TimeSeriesDataFrame] = None,
+        time_limit: int = None,
+        **kwargs,
+    ) -> None:
+        self._check_fit_params()
+
+    def _get_inference_data_loader(
+        self,
+        data: TimeSeriesDataFrame,
+        context_length: int,
+        num_workers: int = 0,
+    ):
+        import torch
+
+        chronos_dataset = ChronosInferenceDataset(
+            target_df=data,
+            target_column=self.target,
+            context_length=context_length,
+        )
+
+        return torch.utils.data.DataLoader(
+            chronos_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+        )
+
+    def _predict(
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        **kwargs,
+    ) -> TimeSeriesDataFrame:
+        # We defer initialization of the model pipeline. i.e., the model is only loaded to device memory
+        # during inference. We also infer the maximum length of the time series in the inference data set
+        # 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.
+        context_length = self.context_length or min(
+            data.num_timesteps_per_item().max(),
+            self.maximum_context_length,
+        )
+
+        with warning_filter(all_warnings=True):
+            import torch
+
+            # load model pipeline to device memory
+            self.load_model_pipeline(context_length=context_length)
+
+            self.model_pipeline.model.eval()
+            with torch.inference_mode():
+                prediction_samples = [
+                    self.model_pipeline.predict(
+                        batch,
+                        prediction_length=self.prediction_length,
+                        num_samples=self.num_samples,
+                        limit_prediction_length=False,
+                    )
+                    .detach()
+                    .cpu()
+                    .numpy()
+                    for batch in self._get_inference_data_loader(
+                        data=data,
+                        num_workers=self.data_loader_num_workers,
+                        context_length=context_length,
+                    )
+                ]
+
+        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
+
+        df = pd.DataFrame(
+            np.concatenate([mean, quantiles], axis=1),
+            columns=["mean"] + [str(q) for q in self.quantile_levels],
+            index=get_forecast_horizon_index_ts_dataframe(data, self.prediction_length),
+        )
+
+        return TimeSeriesDataFrame(df)

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon/timeseries/models/presets.py

@@ -14,6 +14,7 @@ from . import (
     AutoCESModel,
     AutoETSModel,
     AverageModel,
+    ChronosModel,
     CrostonSBAModel,
     DeepARModel,
     DirectTabularModel,
@@ -68,6 +69,7 @@ MODEL_TYPES = dict(
     ADIDA=ADIDAModel,
     CrostonSBA=CrostonSBAModel,
     IMAPA=IMAPAModel,
+    Chronos=ChronosModel,
 )
 
 DEFAULT_MODEL_NAMES = {v: k for k, v in MODEL_TYPES.items()}
@@ -85,6 +87,7 @@ DEFAULT_MODEL_PRIORITY = dict(
     AutoETS=80,
     AutoARIMA=70,
     RecursiveTabular=60,
+    Chronos=50,
     DirectTabular=50,
     DeepAR=40,
     TemporalFusionTransformer=30,

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon/timeseries/predictor.py

@@ -496,11 +496,23 @@ class TimeSeriesPredictor(TimeSeriesPredictorDeprecatedMixin):
 
             Available presets:
 
-            - ``"fast_training"``: fit simple statistical models (``ETS``, ``Theta``, ``Naive``, ``SeasonalNaive``) + fast tree-based models ``RecursiveTabular`` and ``DirectTabular``. These models are fast to train but may not be very accurate.
-            - ``"medium_quality"``: all models mentioned above + deep learning model ``TemporalFusionTransformer``. Default setting that produces good forecasts with reasonable training time.
-            - ``"high_quality"``: All ML models available in AutoGluon + additional statistical models (``NPTS``, ``AutoETS``, ``AutoARIMA``, ``CrostonSBA``, ``DynamicOptimizedTheta``). Much more accurate than ``medium_quality``, but takes longer to train.
+            - ``"fast_training"``: fit simple statistical models (``ETS``, ``Theta``, ``Naive``, ``SeasonalNaive``) + fast tree-based models ``RecursiveTabular``
+              and ``DirectTabular``. These models are fast to train but may not be very accurate.
+            - ``"medium_quality"``: all models mentioned above + deep learning model ``TemporalFusionTransformer``. Default setting that produces good forecasts
+              with reasonable training time.
+            - ``"high_quality"``: All ML models available in AutoGluon + additional statistical models (``NPTS``, ``AutoETS``, ``AutoARIMA``, ``CrostonSBA``,
+              ``DynamicOptimizedTheta``). Much more accurate than ``medium_quality``, but takes longer to train.
             - ``"best_quality"``: Same models as in ``"high_quality"`, but performs validation with multiple backtests. Usually better than ``high_quality``, but takes even longer to train.
 
+            Available presets with the `Chronos <https://github.com/amazon-science/chronos-forecasting>`_ model:
+
+            - ``"chronos_{model_size}"``: where model size is one of ``tiny,mini,small,base,large``. Uses the Chronos pretrained model for zero-shot forecasting.
+              See the documentation for ``ChronosModel`` or see `Hugging Face <https://huggingface.co/collections/amazon/chronos-models-65f1791d630a8d57cb718444>`_ for more information.
+              Note that a GPU is required for model sizes ``small``, ``base`` and ``large``.
+            - ``"chronos"``: alias for ``"chronos_small"``.
+            - ``"chronos_ensemble"``: builds an ensemble of the models specified in ``"high_quality"`` and ``"chronos_small"``.
+            - ``"chronos_large_ensemble"``: builds an ensemble of the models specified in ``"high_quality"`` and ``"chronos_large"``.
+
             Details for these presets can be found in ``autogluon/timeseries/configs/presets_configs.py``. If not
             provided, user-provided values for ``hyperparameters`` and ``hyperparameter_tune_kwargs`` will be used
             (defaulting to their default values specified below).

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon/timeseries/utils/warning_filters.py

@@ -3,6 +3,7 @@ import functools
 import io
 import logging
 import os
+import re
 import sys
 import warnings
 
@@ -12,10 +13,13 @@ __all__ = ["warning_filter", "disable_root_logger", "disable_tqdm"]
 
 
 @contextlib.contextmanager
-def warning_filter():
+def warning_filter(all_warnings: bool = False):
+    categories = [RuntimeWarning, UserWarning, ConvergenceWarning, ValueWarning, FutureWarning]
+    if all_warnings:
+        categories.append(Warning)
     with warnings.catch_warnings():
         env_py_warnings = os.environ.get("PYTHONWARNINGS", "")
-        for warning_category in [RuntimeWarning, UserWarning, ConvergenceWarning, ValueWarning, FutureWarning]:
+        for warning_category in categories:
             warnings.simplefilter("ignore", category=warning_category)
         try:
             os.environ["PYTHONWARNINGS"] = "ignore"
@@ -25,14 +29,28 @@ def warning_filter():
 
 
 @contextlib.contextmanager
-def disable_root_logger():
+def disable_root_logger(root_log_level=logging.ERROR):
     try:
-        logging.getLogger().setLevel(logging.ERROR)
+        logging.getLogger().setLevel(root_log_level)
         yield
     finally:
         logging.getLogger().setLevel(logging.INFO)
 
 
+@contextlib.contextmanager
+def set_loggers_level(regex: str, level=logging.ERROR):
+    log_levels = {}
+    try:
+        for logger_name in logging.root.manager.loggerDict:
+            if re.match(regex, logger_name):
+                log_levels[logger_name] = logging.getLogger(logger_name).level
+                logging.getLogger(logger_name).setLevel(level)
+        yield
+    finally:
+        for logger_name, level in log_levels.items():
+            logging.getLogger(logger_name).setLevel(level)
+
+
 @contextlib.contextmanager
 def disable_tqdm():
     """monkey-patch tqdm to disable it within context"""

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon/timeseries/version.py

@@ -1,3 +1,3 @@
 """This is the autogluon version file."""
-__version__ = '1.0.1b20240323'
+__version__ = '1.0.1b20240325'
 __lite__ = False

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon.timeseries.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.timeseries
-Version: 1.0.1b20240323
+Version: 1.0.1b20240325
 Summary: AutoML for Image, Text, and Tabular Data
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community
@@ -132,5 +132,6 @@ Classifier: Topic :: Scientific/Engineering :: Image Recognition
 Requires-Python: >=3.8, <3.12
 Description-Content-Type: text/markdown
 Provides-Extra: tests
-Provides-Extra: chronos-cpu
+Provides-Extra: chronos-openvino
+Provides-Extra: chronos-onnx
 Provides-Extra: all

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon.timeseries.egg-info/SOURCES.txt

@@ -29,6 +29,9 @@ src/autogluon/timeseries/models/abstract/model_trial.py
 src/autogluon/timeseries/models/autogluon_tabular/__init__.py
 src/autogluon/timeseries/models/autogluon_tabular/mlforecast.py
 src/autogluon/timeseries/models/autogluon_tabular/utils.py
+src/autogluon/timeseries/models/chronos/__init__.py
+src/autogluon/timeseries/models/chronos/chronos.py
+src/autogluon/timeseries/models/chronos/model.py
 src/autogluon/timeseries/models/ensemble/__init__.py
 src/autogluon/timeseries/models/ensemble/abstract_timeseries_ensemble.py
 src/autogluon/timeseries/models/ensemble/greedy_ensemble.py

{autogluon.timeseries-1.0.1b20240323 → autogluon.timeseries-1.0.1b20240325}/src/autogluon.timeseries.egg-info/requires.txt

@@ -16,15 +16,19 @@ utilsforecast<0.0.11,>=0.0.10
 tqdm<5,>=4.38
 orjson~=3.9
 tensorboard<3,>=2.9
-autogluon.core[raytune]==1.0.1b20240323
-autogluon.common==1.0.1b20240323
-autogluon.tabular[catboost,lightgbm,xgboost]==1.0.1b20240323
+autogluon.core[raytune]==1.0.1b20240325
+autogluon.common==1.0.1b20240325
+autogluon.tabular[catboost,lightgbm,xgboost]==1.0.1b20240325
 
 [all]
-optimum[nncf,onnxruntime,openvino]<1.18,>=1.17
+optimum[nncf,openvino]<1.18,>=1.17
+optimum[onnxruntime]<1.18,>=1.17
 
-[chronos-cpu]
-optimum[nncf,onnxruntime,openvino]<1.18,>=1.17
+[chronos-onnx]
+optimum[onnxruntime]<1.18,>=1.17
+
+[chronos-openvino]
+optimum[nncf,openvino]<1.18,>=1.17
 
 [tests]
 pytest

autogluon.timeseries-1.0.1b20240323/src/autogluon/timeseries/configs/presets_configs.py

@@ -1,11 +0,0 @@
-"""Preset configurations for autogluon.timeseries Predictors"""
-
-# TODO: change default HPO settings when other HPO strategies (e.g., Ray tune) are available
-# TODO: add refit_full arguments once refitting is available
-
-TIMESERIES_PRESETS_CONFIGS = dict(
-    best_quality={"hyperparameters": "default", "num_val_windows": 2},
-    high_quality={"hyperparameters": "default"},
-    medium_quality={"hyperparameters": "light"},
-    fast_training={"hyperparameters": "very_light"},
-)