autogluon.timeseries 1.4.1b20251115__py3-none-any.whl → 1.5.0b20251221__py3-none-any.whl

autogluon/timeseries/models/chronos/chronos2.py

@@ -0,0 +1,395 @@
+import logging
+import os
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from typing_extensions import Self
+
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
+
+logger = logging.getLogger(__name__)
+
+
+class Chronos2Model(AbstractTimeSeriesModel):
+    """Chronos-2 pretrained time series forecasting model [Ansari2025]_, which provides strong zero-shot forecasting
+    capability natively taking advantage of covariates. The model can also be fine-tuned in a task specific manner.
+
+    This implementation wraps the original implementation in the `chronos-forecasting`
+    `library <https://github.com/amazon-science/chronos-forecasting/blob/main/src/chronos/chronos2/pipeline.py>`_ .
+
+    Chronos-2 can be used both on GPU and CPU. However, we recommend using a GPU for faster inference and fine-tuning.
+
+    Chronos-2 variants can be fine-tuned by setting ``fine_tune=True`` and selecting appropriate fine-tuning parameters
+    such as the learning rate (``fine_tune_lr``) and max steps (``fine_tune_steps``). By default, a low-rank adapter (LoRA)
+    will be used for fine-tuning.
+
+    References
+    ----------
+    .. [Ansari2025] Ansari, Abdul Fatir, Shchur, Oleksandr, Kuken, Jaris et al.
+        "Chronos-2: From Univariate to Universal Forecasting." (2025).
+        https://arxiv.org/abs/2510.15821
+
+    Other Parameters
+    ----------------
+    model_path : str, default = "autogluon/chronos-2"
+        Model path used for the model, i.e., a Hugging Face transformers ``name_or_path``. Can be a
+        compatible model name on Hugging Face Hub or a local path to a model directory.
+    batch_size : int, default = 256
+        Size of batches used during inference.
+    device : str, default = None
+        Device to use for inference (and fine-tuning, if enabled). If None, model will use the GPU if
+        available.
+    cross_learning : bool, default = True
+        If True, the cross-learning mode of Chronos-2 is enabled. This means that the model will make joint
+        predictions across time series in a batch, by default True
+        Note: Enabling this mode makes the results sensitive to the ``batch_size`` used.
+    context_length : int or None, default = None
+        The context length to use for inference. If None, the model will use its default context length
+        of 8192. Shorter context lengths may reduce accuracy, but result in faster inference.
+    fine_tune : bool, default = False
+        If True, the pretrained model will be fine-tuned.
+    fine_tune_mode : str, default = "lora"
+        Fine-tuning mode, either "full" for full fine-tuning or "lora" for Low Rank Adaptation (LoRA).
+        LoRA is faster and uses less memory.
+    fine_tune_lr : float, default = 1e-5
+        The learning rate used for fine-tuning. When using full fine-tuning, a lower learning rate such as 1e-6
+        is recommended.
+    fine_tune_steps : int, default = 1000
+        The number of gradient update steps to fine-tune for.
+    fine_tune_batch_size : int, default = 32
+        The batch size to use for fine-tuning.
+    fine_tune_context_length : int, default = 2048
+        The maximum context_length to use for fine-tuning
+    eval_during_fine_tune : bool, default = False
+        If True, validation will be performed during fine-tuning to select the best checkpoint. Setting this
+        argument to True may result in slower fine-tuning. This parameter is ignored if ``skip_model_selection=True``
+        in ``TimeSeriesPredictor.fit``.
+    fine_tune_eval_max_items : int, default = 256
+        The maximum number of randomly-sampled time series to use from the validation set for evaluation
+        during fine-tuning. If None, the entire validation dataset will be used.
+    fine_tune_lora_config : dict, optional
+        Configuration for LoRA fine-tuning when ``fine_tune_mode="lora"``. If None and LoRA is enabled,
+        a default configuration will be used. Example: ``{"r": 8, "lora_alpha": 16}``.
+    fine_tune_trainer_kwargs : dict, optional
+        Extra keyword arguments passed to ``transformers.TrainingArguments``
+    revision : str, default = None
+        Model revision to use (branch name or commit hash). If None, the default branch (usually "main") is used.
+    disable_known_covariates : bool, default = False
+        If True, known covariates won't be used by the model even if they are present in the dataset.
+    disable_past_covariates : bool, default = False
+        If True, past covariates won't be used by the model even if they are present in the dataset.
+    """
+
+    ag_model_aliases = ["Chronos-2"]
+    ag_priority = 75
+    fine_tuned_ckpt_name: str = "fine-tuned-ckpt"
+
+    _supports_known_covariates = True
+    _supports_past_covariates = True
+
+    def __init__(
+        self,
+        freq: str | None = None,
+        prediction_length: int = 1,
+        path: str | None = None,
+        name: str | None = None,
+        eval_metric: str | None = None,
+        hyperparameters: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            path=path,
+            freq=freq,
+            prediction_length=prediction_length,
+            name=name,
+            eval_metric=eval_metric,
+            hyperparameters=hyperparameters,
+            **kwargs,
+        )
+        self._is_fine_tuned: bool = False
+        self._model_pipeline = None
+
+    @property
+    def model_path(self) -> str:
+        default_model_path = self.get_hyperparameter("model_path")
+
+        if self._is_fine_tuned:
+            model_path = os.path.join(self.path, self.fine_tuned_ckpt_name)
+            if not os.path.exists(model_path):
+                raise ValueError("Cannot find finetuned checkpoint for Chronos-2.")
+            else:
+                return model_path
+
+        return default_model_path
+
+    def save(self, path: str | None = 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)
+
+    def _fit(
+        self,
+        train_data: TimeSeriesDataFrame,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
+        verbosity: int = 2,
+        **kwargs,
+    ) -> None:
+        self._check_fit_params()
+        self._log_unused_hyperparameters()
+        self.load_model_pipeline()
+
+        # NOTE: This must be placed after load_model_pipeline to ensure that the loggers are available in loggerDict
+        self._update_transformers_loggers(logging.ERROR if verbosity <= 3 else logging.WARNING)
+
+        if self.get_hyperparameter("fine_tune"):
+            self._fine_tune(train_data, val_data, time_limit=time_limit, verbosity=verbosity)
+
+    def get_hyperparameters(self) -> dict:
+        """Gets params that are passed to the inner model."""
+        init_args = super().get_hyperparameters()
+
+        fine_tune_trainer_kwargs = dict(disable_tqdm=True)
+        user_fine_tune_trainer_kwargs = init_args.get("fine_tune_trainer_kwargs", {})
+        fine_tune_trainer_kwargs.update(user_fine_tune_trainer_kwargs)
+        init_args["fine_tune_trainer_kwargs"] = fine_tune_trainer_kwargs
+
+        return init_args.copy()
+
+    def _get_default_hyperparameters(self) -> dict:
+        return {
+            "model_path": "autogluon/chronos-2",
+            "batch_size": 256,
+            "device": None,
+            "cross_learning": True,
+            "context_length": None,
+            "fine_tune": False,
+            "fine_tune_mode": "lora",
+            "fine_tune_lr": 1e-5,
+            "fine_tune_steps": 1000,
+            "fine_tune_batch_size": 32,
+            "fine_tune_context_length": 2048,
+            "eval_during_fine_tune": False,
+            "fine_tune_eval_max_items": 256,
+            "fine_tune_lora_config": None,
+            "revision": None,
+            "disable_known_covariates": False,
+            "disable_past_covariates": False,
+        }
+
+    @property
+    def allowed_hyperparameters(self) -> list[str]:
+        return super().allowed_hyperparameters + [
+            "model_path",
+            "batch_size",
+            "device",
+            "cross_learning",
+            "context_length",
+            "fine_tune",
+            "fine_tune_mode",
+            "fine_tune_lr",
+            "fine_tune_steps",
+            "fine_tune_batch_size",
+            "fine_tune_context_length",
+            "eval_during_fine_tune",
+            "fine_tune_eval_max_items",
+            "fine_tune_lora_config",
+            "fine_tune_trainer_kwargs",
+            "revision",
+            "disable_known_covariates",
+            "disable_past_covariates",
+        ]
+
+    def _remove_disabled_covariates(
+        self, past_df: pd.DataFrame, future_df: pd.DataFrame | None
+    ) -> tuple[pd.DataFrame, pd.DataFrame | None]:
+        """Remove covariates from dataframes based on disable flags."""
+        cols_to_remove = []
+        if self.get_hyperparameter("disable_past_covariates"):
+            cols_to_remove.extend(self.covariate_metadata.past_covariates)
+        if self.get_hyperparameter("disable_known_covariates"):
+            cols_to_remove.extend(self.covariate_metadata.known_covariates)
+            future_df = None
+
+        if cols_to_remove:
+            past_df = past_df.drop(columns=cols_to_remove)
+
+        return past_df, future_df
+
+    def _predict(
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: TimeSeriesDataFrame | None = None,
+        **kwargs,
+    ) -> TimeSeriesDataFrame:
+        from .utils import timeout_callback
+
+        if self._model_pipeline is None:
+            self.load_model_pipeline()
+        assert self._model_pipeline is not None
+
+        if max(data.num_timesteps_per_item()) < 3:
+            # If all time series have length 2 or less, we prepend 2 dummy timesteps to the first series
+            first_item_id = data.index.get_level_values(0)[0]
+            dummy_timestamps = pd.date_range(end=data.loc[first_item_id].index[0], periods=3, freq=self.freq)[:-1]
+            full_time_index_first_item = data.loc[first_item_id].index.union(dummy_timestamps)
+            new_index = (
+                pd.MultiIndex.from_product([[first_item_id], full_time_index_first_item], names=data.index.names)
+            ).union(data.index)
+            context_df = data.reindex(new_index).reset_index()
+        else:
+            context_df = data.reset_index().to_data_frame()
+
+        batch_size = self.get_hyperparameter("batch_size")
+        cross_learning = self.get_hyperparameter("cross_learning")
+        context_length = self.get_hyperparameter("context_length")
+        future_df = known_covariates.reset_index().to_data_frame() if known_covariates is not None else None
+        time_limit = kwargs.get("time_limit")
+
+        context_df, future_df = self._remove_disabled_covariates(context_df, future_df)
+
+        forecast_df = self._model_pipeline.predict_df(
+            df=context_df,
+            future_df=future_df,
+            target=self.target,
+            prediction_length=self.prediction_length,
+            quantile_levels=self.quantile_levels,
+            context_length=context_length,
+            batch_size=batch_size,
+            validate_inputs=False,
+            cross_learning=cross_learning,
+            after_batch=timeout_callback(time_limit),
+        )
+
+        forecast_df = forecast_df.rename(columns={"predictions": "mean"}).drop(columns="target_name")
+
+        return TimeSeriesDataFrame(forecast_df)
+
+    def load_model_pipeline(self):
+        from chronos.chronos2.pipeline import Chronos2Pipeline
+
+        device = (self.get_hyperparameter("device") or "cuda") if self._is_gpu_available() else "cpu"
+
+        assert self.model_path is not None
+        pipeline = Chronos2Pipeline.from_pretrained(
+            self.model_path,
+            device_map=device,
+            revision=self.get_hyperparameter("revision"),
+        )
+
+        self._model_pipeline = pipeline
+
+    def persist(self) -> Self:
+        self.load_model_pipeline()
+        return self
+
+    def _update_transformers_loggers(self, log_level: int):
+        for logger_name in logging.root.manager.loggerDict:
+            if "transformers" in logger_name:
+                transformers_logger = logging.getLogger(logger_name)
+                transformers_logger.setLevel(log_level)
+
+    def _fine_tune(
+        self,
+        train_data: TimeSeriesDataFrame,
+        val_data: TimeSeriesDataFrame | None,
+        time_limit: float | None = None,
+        verbosity: int = 2,
+    ):
+        from chronos.df_utils import convert_df_input_to_list_of_dicts_input
+
+        from .utils import LoggerCallback, TimeLimitCallback
+
+        def convert_data(df: TimeSeriesDataFrame):
+            past_df = df.reset_index().to_data_frame()
+            past_df, _ = self._remove_disabled_covariates(past_df, None)
+
+            inputs, _, _ = convert_df_input_to_list_of_dicts_input(
+                df=past_df,
+                future_df=None,
+                target_columns=[self.target],
+                prediction_length=self.prediction_length,
+                validate_inputs=False,
+            )
+
+            # The above utility will only split the dataframe into target and past_covariates, where past_covariates contains
+            # past values of both past-only and known-future covariates. We need to add future_covariates to enable fine-tuning
+            # with known covariates by indicating which covariates are known in the future.
+            if not self.get_hyperparameter("disable_known_covariates"):
+                known_covariates = self.covariate_metadata.known_covariates
+                if len(known_covariates) > 0:
+                    for input_dict in inputs:
+                        # NOTE: the covariates are empty because the actual values are not used
+                        # This only indicates which covariates are known in the future
+                        input_dict["future_covariates"] = {name: np.array([]) for name in known_covariates}
+
+            return inputs
+
+        assert self._model_pipeline is not None
+        hyperparameters = self.get_hyperparameters()
+
+        callbacks = []
+        if time_limit is not None:
+            callbacks.append(TimeLimitCallback(time_limit=time_limit))
+
+        val_inputs = None
+        if val_data is not None and hyperparameters["eval_during_fine_tune"]:
+            # evaluate on a randomly-sampled subset
+            fine_tune_eval_max_items = (
+                min(val_data.num_items, hyperparameters["fine_tune_eval_max_items"])
+                if hyperparameters["fine_tune_eval_max_items"] is not None
+                else val_data.num_items
+            )
+
+            if fine_tune_eval_max_items < val_data.num_items:
+                eval_items = np.random.choice(val_data.item_ids.values, size=fine_tune_eval_max_items, replace=False)  # noqa: F841
+                val_data = val_data.query("item_id in @eval_items")
+
+            assert isinstance(val_data, TimeSeriesDataFrame)
+            val_inputs = convert_data(val_data)
+
+        if verbosity >= 3:
+            logger.warning(
+                "Transformers logging is turned on during fine-tuning. Note that losses reported by transformers "
+                "do not correspond to those specified via `eval_metric`."
+            )
+            callbacks.append(LoggerCallback())
+
+        self._model_pipeline = self._model_pipeline.fit(
+            inputs=convert_data(train_data),
+            prediction_length=self.prediction_length,
+            validation_inputs=val_inputs,
+            finetune_mode=hyperparameters["fine_tune_mode"],
+            lora_config=hyperparameters["fine_tune_lora_config"],
+            context_length=hyperparameters["fine_tune_context_length"],
+            learning_rate=hyperparameters["fine_tune_lr"],
+            num_steps=hyperparameters["fine_tune_steps"],
+            batch_size=hyperparameters["fine_tune_batch_size"],
+            output_dir=self.path,
+            finetuned_ckpt_name=self.fine_tuned_ckpt_name,
+            callbacks=callbacks,
+            remove_printer_callback=True,
+            min_past=1,
+            **hyperparameters["fine_tune_trainer_kwargs"],
+        )
+        self._is_fine_tuned = True
+
+    def _more_tags(self) -> dict[str, Any]:
+        do_fine_tune = self.get_hyperparameter("fine_tune")
+        return {
+            "allow_nan": True,
+            "can_use_train_data": do_fine_tune,
+            "can_use_val_data": do_fine_tune,
+        }
+
+    def _is_gpu_available(self) -> bool:
+        import torch.cuda
+
+        return torch.cuda.is_available()

autogluon/timeseries/models/chronos/model.py

@@ -3,7 +3,7 @@ import os
 import shutil
 import warnings
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -176,8 +176,9 @@ class ChronosModel(AbstractTimeSeriesModel):
         The size of the shuffle buffer to shuffle the data during fine-tuning. If None, shuffling will
         be turned off.
     eval_during_fine_tune : bool, default = False
-        If True, validation will be performed during fine-tuning to select the best checkpoint.
-        Setting this argument to True may result in slower fine-tuning.
+        If True, validation will be performed during fine-tuning to select the best checkpoint. Setting this
+        argument to True may result in slower fine-tuning. This parameter is ignored if ``skip_model_selection=True``
+        in ``TimeSeriesPredictor.fit``.
     fine_tune_eval_max_items : int, default = 256
         The maximum number of randomly-sampled time series to use from the validation set for evaluation
         during fine-tuning. If None, the entire validation dataset will be used.
@@ -185,6 +186,8 @@ class ChronosModel(AbstractTimeSeriesModel):
         Extra keyword arguments passed to ``transformers.TrainingArguments``
     keep_transformers_logs : bool, default = False
         If True, the logs generated by transformers will NOT be removed after fine-tuning
+    revision : str, default = None
+        Model revision to use (branch name or commit hash). If None, the default branch (usually "main") is used.
     """
 
     ag_priority = 55
@@ -196,12 +199,12 @@ class ChronosModel(AbstractTimeSeriesModel):
 
     def __init__(
         self,
-        freq: Optional[str] = None,
+        freq: str | None = None,
         prediction_length: int = 1,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        eval_metric: Optional[str] = None,
-        hyperparameters: Optional[dict[str, Any]] = None,
+        path: str | None = None,
+        name: str | None = None,
+        eval_metric: str | None = None,
+        hyperparameters: dict[str, Any] | None = None,
         **kwargs,  # noqa
     ):
         hyperparameters = hyperparameters if hyperparameters is not None else {}
@@ -226,9 +229,9 @@ class ChronosModel(AbstractTimeSeriesModel):
             **kwargs,
         )
 
-        self._model_pipeline: Optional[Any] = None  # of type BaseChronosPipeline
+        self._model_pipeline: Any | None = None  # of type BaseChronosPipeline
 
-    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
+    def save(self, path: str | None = None, verbose: bool = True) -> str:
         pipeline = self._model_pipeline
         self._model_pipeline = None
         path = super().save(path=path, verbose=verbose)
@@ -292,8 +295,8 @@ class ChronosModel(AbstractTimeSeriesModel):
         """
         return self.ag_default_config.get("default_torch_dtype", "auto")
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
-        minimum_resources: dict[str, Union[int, float]] = {"num_cpus": 1}
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, int | float]:
+        minimum_resources: dict[str, int | float] = {"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
@@ -311,13 +314,14 @@ class ChronosModel(AbstractTimeSeriesModel):
                 "`import torch; torch.cuda.is_available()` returns `True`."
             )
 
-        device = self.device or ("cuda" if gpu_available else "cpu")
+        device = (self.device or "cuda") if gpu_available else "cpu"
 
         assert self.model_path is not None
         pipeline = BaseChronosPipeline.from_pretrained(
             self.model_path,
             device_map=device,
             torch_dtype=self.torch_dtype,
+            revision=self.get_hyperparameter("revision"),
         )
 
         self._model_pipeline = pipeline
@@ -360,6 +364,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             "eval_during_fine_tune": False,
             "fine_tune_eval_max_items": 256,
             "fine_tune_shuffle_buffer_size": 10_000,
+            "revision": None,
         }
 
     @property
@@ -381,6 +386,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             "fine_tune_eval_max_items",
             "fine_tune_trainer_kwargs",
             "keep_transformers_logs",
+            "revision",
         ]
 
     def _get_fine_tune_trainer_kwargs(self, init_args, eval_during_fine_tune: bool):
@@ -437,10 +443,10 @@ class ChronosModel(AbstractTimeSeriesModel):
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-        num_cpus: Optional[int] = None,
-        num_gpus: Optional[int] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
         verbosity: int = 2,
         **kwargs,
     ) -> None:
@@ -462,11 +468,10 @@ class ChronosModel(AbstractTimeSeriesModel):
         # verbosity < 3: all logs and warnings from transformers will be suppressed
         # verbosity >= 3: progress bar and loss logs will be logged
         # verbosity 4: everything will be logged
-        verbosity = kwargs.get("verbosity", 2)
         for logger_name in logging.root.manager.loggerDict:
             if "transformers" in logger_name:
                 transformers_logger = logging.getLogger(logger_name)
-                transformers_logger.setLevel(logging.ERROR if verbosity <= 3 else logging.INFO)
+                transformers_logger.setLevel(logging.ERROR if verbosity <= 3 else logging.WARNING)
 
         self._check_fit_params()
         self._log_unused_hyperparameters()
@@ -558,7 +563,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             if time_limit is not None:
                 callbacks.append(TimeLimitCallback(time_limit=time_limit))
 
-            tokenizer_val_dataset: Optional[ChronosFineTuningDataset] = None
+            tokenizer_val_dataset: ChronosFineTuningDataset | None = None
             if val_data is not None:
                 callbacks.append(EvaluateAndSaveFinalStepCallback())
                 # evaluate on a randomly-sampled subset
@@ -599,7 +604,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             if verbosity >= 3:
                 logger.warning(
                     "Transformers logging is turned on during fine-tuning. Note that losses reported by transformers "
-                    "may not correspond to those specified via `eval_metric`."
+                    "do not correspond to those specified via `eval_metric`."
                 )
                 trainer.add_callback(LoggerCallback())
 
@@ -619,7 +624,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         context_length: int,
         batch_size: int,
         num_workers: int = 0,
-        time_limit: Optional[float] = None,
+        time_limit: float | None = None,
     ):
         from .utils import ChronosInferenceDataLoader, ChronosInferenceDataset, timeout_callback
 
@@ -634,7 +639,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             batch_size=batch_size,
             shuffle=False,
             num_workers=num_workers,
-            on_batch=timeout_callback(seconds=time_limit),
+            after_batch=timeout_callback(seconds=time_limit),
         )
 
     def _get_context_length(self, data: TimeSeriesDataFrame) -> int:
@@ -647,7 +652,7 @@ class ChronosModel(AbstractTimeSeriesModel):
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         from chronos import ChronosBoltPipeline, ChronosPipeline
@@ -725,7 +730,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         return TimeSeriesDataFrame(df)
 
     def _more_tags(self) -> dict:
-        do_fine_tune = self.get_hyperparameters()["fine_tune"]
+        do_fine_tune = self.get_hyperparameter("fine_tune")
         return {
             "allow_nan": True,
             "can_use_train_data": do_fine_tune,

autogluon/timeseries/models/chronos/utils.py

@@ -1,7 +1,7 @@
 import logging
 import time
 from itertools import chain, cycle
-from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Literal, Optional
+from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Literal
 
 import numpy as np
 import torch
@@ -93,7 +93,7 @@ class ChronosFineTuningDataset(IterableDataset):
         target_column: str = "target",
         context_length: int = 512,
         prediction_length: int = 64,
-        tokenizer: Optional["ChronosTokenizer"] = None,
+        tokenizer: "ChronosTokenizer | None" = None,
         mode: Literal["training", "validation"] = "training",
     ) -> None:
         super().__init__()
@@ -194,7 +194,7 @@ class ChronosFineTuningDataset(IterableDataset):
         for entry in iterable:
             yield format_transform_fn(entry)
 
-    def shuffle(self, shuffle_buffer_size: Optional[int] = None):
+    def shuffle(self, shuffle_buffer_size: int | None = None):
         """Returns a (pseudo) shuffled version of this iterable dataset.
 
         Parameters
@@ -255,7 +255,7 @@ class ChronosInferenceDataset:
 
 class ChronosInferenceDataLoader(torch.utils.data.DataLoader):
     def __init__(self, *args, **kwargs):
-        self.callback: Callable = kwargs.pop("on_batch", lambda: None)
+        self.callback: Callable = kwargs.pop("after_batch", lambda: None)
         super().__init__(*args, **kwargs)
 
     def __iter__(self):  # type: ignore
@@ -305,7 +305,7 @@ class LoggerCallback(TrainerCallback):
             logger.info(logs)
 
 
-def timeout_callback(seconds: Optional[float]) -> Callable:
+def timeout_callback(seconds: float | None) -> Callable:
     """Return a callback object that raises an exception if time limit is exceeded."""
     start_time = time.monotonic()
 

autogluon/timeseries/models/ensemble/__init__.py

@@ -1,28 +1,35 @@
 from .abstract import AbstractTimeSeriesEnsembleModel
-from .array_based import MedianEnsemble, PerQuantileTabularEnsemble, TabularEnsemble
+from .array_based import LinearStackerEnsemble, MedianEnsemble, PerQuantileTabularEnsemble, TabularEnsemble
+from .per_item_greedy import PerItemGreedyEnsemble
 from .weighted import GreedyEnsemble, PerformanceWeightedEnsemble, SimpleAverageEnsemble
 
 
 def get_ensemble_class(name: str):
     mapping = {
-        "GreedyEnsemble": GreedyEnsemble,
-        "PerformanceWeightedEnsemble": PerformanceWeightedEnsemble,
-        "SimpleAverageEnsemble": SimpleAverageEnsemble,
-        "WeightedEnsemble": GreedyEnsemble,  # old alias for this model
-        "MedianEnsemble": MedianEnsemble,
-        "TabularEnsemble": TabularEnsemble,
-        "PerQuantileTabularEnsemble": PerQuantileTabularEnsemble,
+        "Greedy": GreedyEnsemble,
+        "PerItemGreedy": PerItemGreedyEnsemble,
+        "PerformanceWeighted": PerformanceWeightedEnsemble,
+        "SimpleAverage": SimpleAverageEnsemble,
+        "Weighted": GreedyEnsemble,  # old alias for this model
+        "Median": MedianEnsemble,
+        "Tabular": TabularEnsemble,
+        "PerQuantileTabular": PerQuantileTabularEnsemble,
+        "LinearStacker": LinearStackerEnsemble,
     }
-    if name not in mapping:
+
+    name_clean = name.removesuffix("Ensemble")
+    if name_clean not in mapping:
         raise ValueError(f"Unknown ensemble type: {name}. Available: {list(mapping.keys())}")
-    return mapping[name]
+    return mapping[name_clean]
 
 
 __all__ = [
     "AbstractTimeSeriesEnsembleModel",
     "GreedyEnsemble",
+    "LinearStackerEnsemble",
     "MedianEnsemble",
     "PerformanceWeightedEnsemble",
+    "PerItemGreedyEnsemble",
     "PerQuantileTabularEnsemble",
     "SimpleAverageEnsemble",
     "TabularEnsemble",