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

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)

autogluon/timeseries/models/chronos/{pipeline.py → pipeline/chronos.py}

@@ -11,14 +11,16 @@ 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 transformers import AutoConfig, AutoModelForSeq2SeqLM, GenerationConfig, PreTrainedModel
 
 from autogluon.timeseries.utils.warning_filters import set_loggers_level
 
-logger = logging.getLogger(__name__)
+from .base import BaseChronosPipeline, ForecastType
 
+logger = logging.getLogger("autogluon.timeseries.models.chronos")
 
-__all__ = ["ChronosConfig", "ChronosPipeline", "OptimizedChronosPipeline"]
+
+__all__ = ["ChronosConfig", "ChronosPipeline"]
 
 
 @dataclass
@@ -35,7 +37,7 @@ class ChronosConfig:
     pad_token_id: int
     eos_token_id: int
     use_eos_token: bool
-    model_type: Literal["causal", "seq2seq"]
+    model_type: Literal["seq2seq"]
     context_length: int
     prediction_length: int
     num_samples: int
@@ -279,18 +281,7 @@ class ChronosPretrainedModel(nn.Module):
         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:
+class ChronosPipeline(BaseChronosPipeline):
     """
     A ``ChronosPipeline`` uses the given tokenizer and model to forecast
     input time series.
@@ -308,21 +299,12 @@ class ChronosPipeline:
 
     tokenizer: ChronosTokenizer
     model: ChronosPretrainedModel
+    forecast_type: ForecastType = ForecastType.SAMPLES
 
     def __init__(self, tokenizer, model):
         self.tokenizer = tokenizer
         self.model = model
 
-    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
-
     @torch.no_grad()
     def embed(self, context: Union[torch.Tensor, List[torch.Tensor]]) -> Tuple[torch.Tensor, Any]:
         """
@@ -363,7 +345,7 @@ class ChronosPipeline:
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_p: Optional[float] = None,
-        limit_prediction_length: bool = True,
+        limit_prediction_length: bool = False,
     ) -> torch.Tensor:
         """
         Get forecasts for the given time series.
@@ -442,42 +424,33 @@ class ChronosPipeline:
 
         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),
+    def predict_quantiles(
+        self,
+        context: torch.Tensor,
+        prediction_length: int,
+        quantile_levels: List[float],
+        num_samples: Optional[int] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        num_samples = num_samples or self.model.config.num_samples
+        prediction_samples = (
+            self.predict(
+                context,
+                prediction_length=prediction_length,
+                num_samples=num_samples,
+            )
+            .detach()
+            .cpu()
+            .swapaxes(1, 2)
         )
+        mean = prediction_samples.mean(axis=-1, keepdims=True)
+        quantiles = torch.quantile(
+            prediction_samples,
+            q=torch.tensor(quantile_levels, dtype=prediction_samples.dtype),
+            dim=-1,
+        ).permute(1, 2, 0)
 
-
-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,
-    }
+        return quantiles, mean
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
@@ -498,49 +471,40 @@ class OptimizedChronosPipeline(ChronosPipeline):
             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}"
+        assert 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"
                     )
 
-                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)
+                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}
+                    )
 
         return cls(
             tokenizer=chronos_config.create_tokenizer(),