autogluon.timeseries 1.3.2b20250712__py3-none-any.whl → 1.4.1b20251116__py3-none-any.whl

autogluon/timeseries/models/chronos/model.py

@@ -3,14 +3,15 @@ import os
 import shutil
 import warnings
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
+from typing_extensions import Self
 
 from autogluon.common.loaders import load_pkl
 from autogluon.common.space import Space
-from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.utils.warning_filters import disable_duplicate_logs, warning_filter
 
@@ -81,72 +82,92 @@ MODEL_ALIASES = {
 
 
 class ChronosModel(AbstractTimeSeriesModel):
-    """Chronos [Ansari2024]_ pretrained time series forecasting models which can be used for zero-shot forecasting or fine-tuned
-    in a task-specific manner. 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 capable of much faster inference.
-
-    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
-    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/>`_.
-
-    Both Chronos and Chronos-Bolt 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``).
+    """Chronos [Ansari2024]_ pretrained time series forecasting models which can be used for zero-shot
+    forecasting or fine-tuned in a task-specific manner.
+
+    Models can be based on the original
+    `Chronos <https://github.com/amazon-science/chronos-forecasting/blob/main/src/chronos/chronos.py>`_
+    implementation, as well as a newer family of
+    `Chronos-Bolt <https://github.com/amazon-science/chronos-forecasting/blob/main/src/chronos/chronos_bolt.py>`_
+    models capable of much faster inference.
+
+    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 (original) ``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.
+
+    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.
+
+    Both Chronos and Chronos-Bolt 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``).
 
     References
     ----------
     .. [Ansari2024] Ansari, Abdul Fatir, Stella, Lorenzo et al.
         "Chronos: Learning the Language of Time Series."
-        http://arxiv.org/abs/2403.07815
+        Transactions on Machine Learning Research (2024).
+        https://openreview.net/forum?id=gerNCVqqtR
 
 
     Other Parameters
     ----------------
-    model_path: str, default = "autogluon/chronos-bolt-small"
+    model_path : str, default = "autogluon/chronos-bolt-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., ``autogluon/chronos-t5-{model_size}``) can be specified with aliases
         ``tiny``, ``mini`` , ``small``, ``base``, and ``large``. Chronos-Bolt models can be specified
         with ``bolt_tiny``, ``bolt_mini``, ``bolt_small``, and ``bolt_base``.
     batch_size : int, default = 256
-        Size of batches used during inference. The default ``batch_size`` is selected based on the model type. For Chronos-Bolt
-        models the ``batch_size`` is set to 256 whereas Chronos models used a ``batch_size`` of 16, except Chronos (Large) which
-        uses 8. For the Chronos-Bolt models, the ``batch_size`` is reduced by a factor of 4 when the prediction horizon is greater
-        than the model's default prediction length.
+        Size of batches used during inference.
+
+        The default ``batch_size`` is selected based on the model type. Chronos (original) models use a
+        ``batch_size`` of 16, except Chronos (Large) which uses 8.
+
+        For Chronos-Bolt models the ``batch_size`` is set to 256. However, ``batch_size`` is reduced by
+        a factor of 4 when the prediction horizon is greater than the model's
+        default prediction length.
     num_samples : int, default = 20
         Number of samples used during inference, only used for the original Chronos models
     device : str, default = None
-        Device to use for inference (and fine-tuning, if enabled). If None, model will use the GPU if available.
-        For larger Chronos model sizes ``small``, ``base``, and ``large``; inference will fail if no GPU is available.
-        For Chronos-Bolt models, inference can be done on the CPU. Although fine-tuning the smaller Chronos models
-        (``tiny`` and ``mini``) and all Chronos-Bolt is allowed on the CPU, we recommend using a GPU for faster fine-tuning.
+        Device to use for inference (and fine-tuning, if enabled). If None, model will use the GPU if
+        available. For larger Chronos model sizes ``small``, ``base``, and ``large``; inference will fail
+        if no GPU is available.
+
+        For Chronos-Bolt models, inference can be performed on the CPU. Although fine-tuning the smaller
+        Chronos models (``tiny`` and ``mini``) and all Chronos-Bolt is allowed on the CPU, we recommend
+        using a GPU for faster fine-tuning.
     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 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.
-    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.
+        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 cap it at 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.
+    torch_dtype : torch.dtype or {"auto", "bfloat16", "float32"}, 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.
+        Number of worker processes to be used in the data loader. See documentation on
+        ``torch.utils.data.DataLoader`` for more information.
     fine_tune : bool, default = False
         If True, the pretrained model will be fine-tuned
-    fine_tune_lr: float, default = 1e-5
-        The learning rate used for fine-tuning. This default is suitable for Chronos-Bolt models; for the original
-        Chronos models, we recommend using a higher learning rate such as ``1e-4``
+    fine_tune_lr : float, default = 1e-5
+        The learning rate used for fine-tuning. This default is suitable for Chronos-Bolt models; for
+        the original Chronos models, we recommend using a higher learning rate such as ``1e-4``.
     fine_tune_steps : int, default = 1000
         The number of gradient update steps to fine-tune for
     fine_tune_batch_size : int, default = 32
@@ -162,12 +183,12 @@ class ChronosModel(AbstractTimeSeriesModel):
         during fine-tuning. If None, the entire validation dataset will be used.
     fine_tune_trainer_kwargs : dict, optional
         Extra keyword arguments passed to ``transformers.TrainingArguments``
-    keep_transformers_logs: bool, default = False
+    keep_transformers_logs : bool, default = False
         If True, the logs generated by transformers will NOT be removed after fine-tuning
     """
 
-    # default number of samples for prediction
-    default_num_samples: int = 20
+    ag_priority = 55
+    default_num_samples: int = 20  # default number of samples for prediction
     default_model_path = "autogluon/chronos-bolt-small"
     default_max_time_limit_ratio = 0.8
     maximum_context_length = 2048
@@ -180,13 +201,13 @@ class ChronosModel(AbstractTimeSeriesModel):
         path: Optional[str] = None,
         name: Optional[str] = None,
         eval_metric: Optional[str] = None,
-        hyperparameters: Optional[Dict[str, Any]] = None,
+        hyperparameters: Optional[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)
+        model_path_input: str = hyperparameters.get("model_path", self.default_model_path)
+        self.model_path: str = MODEL_ALIASES.get(model_path_input, model_path_input)
 
         name = name if name is not None else "Chronos"
         if not isinstance(model_path_input, Space):
@@ -216,7 +237,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         return str(path)
 
     @classmethod
-    def load(cls, path: str, reset_paths: bool = True, verbose: bool = True) -> "ChronosModel":
+    def load(cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True) -> Self:
         model = load_pkl.load(path=os.path.join(path, cls.model_file_name), verbose=verbose)
         if reset_paths:
             model.set_contexts(path)
@@ -241,7 +262,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         return self._model_pipeline
 
     @property
-    def ag_default_config(self) -> Dict[str, Any]:
+    def ag_default_config(self) -> dict[str, Any]:
         """The default configuration of the model used by AutoGluon if the model is one of those
         defined in MODEL_CONFIGS. For now, these are ``autogluon/chronos-t5-*`` family of models.
         """
@@ -271,15 +292,15 @@ 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, Union[int, float]]:
+        minimum_resources: dict[str, Union[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
         return minimum_resources
 
     def load_model_pipeline(self, is_training: bool = False):
-        from .pipeline import BaseChronosPipeline
+        from chronos import BaseChronosPipeline
 
         gpu_available = self._is_gpu_available()
 
@@ -292,6 +313,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
         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,
@@ -322,7 +344,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
         return init_args.copy()
 
-    def _get_default_hyperparameters(self) -> Dict:
+    def _get_default_hyperparameters(self) -> dict:
         return {
             "batch_size": self.default_batch_size,
             "num_samples": self.default_num_samples,
@@ -416,19 +438,23 @@ class ChronosModel(AbstractTimeSeriesModel):
         self,
         train_data: TimeSeriesDataFrame,
         val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[int] = None,
+        time_limit: Optional[float] = None,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
+        verbosity: int = 2,
         **kwargs,
     ) -> None:
         import transformers
+        from chronos import ChronosBoltPipeline, ChronosPipeline
         from packaging import version
         from transformers.trainer import PrinterCallback, Trainer, TrainingArguments
 
-        from .pipeline import ChronosBoltPipeline, ChronosPipeline
-        from .pipeline.utils import (
+        from .utils import (
             ChronosFineTuningDataset,
             EvaluateAndSaveFinalStepCallback,
             LoggerCallback,
             TimeLimitCallback,
+            update_output_quantiles,
         )
 
         # TODO: Add support for fine-tuning models with context_length longer than the pretrained model
@@ -481,9 +507,12 @@ class ChronosModel(AbstractTimeSeriesModel):
 
                 if self.prediction_length != fine_tune_prediction_length:
                     logger.debug(
-                        f"\tChronosBolt models can only be fine-tuned with a maximum prediction_length of {model_prediction_length}. "
+                        f"\tChronos-Bolt models can only be fine-tuned with a maximum prediction_length of {model_prediction_length}. "
                         f"Fine-tuning prediction_length has been changed to {fine_tune_prediction_length}."
                     )
+                if self.quantile_levels != self.model_pipeline.quantiles:
+                    update_output_quantiles(self.model_pipeline.model, self.quantile_levels)
+                    logger.info(f"\tChronos-Bolt will be fine-tuned with quantile_levels={self.quantile_levels}")
             else:
                 raise ValueError(f"Unsupported model pipeline: {type(self.model_pipeline)}")
 
@@ -513,7 +542,7 @@ class ChronosModel(AbstractTimeSeriesModel):
                 # transformers changed the argument name from `evaluation_strategy` to `eval_strategy`
                 fine_tune_trainer_kwargs["eval_strategy"] = fine_tune_trainer_kwargs.pop("evaluation_strategy")
 
-            training_args = TrainingArguments(**fine_tune_trainer_kwargs, **pipeline_specific_trainer_kwargs)
+            training_args = TrainingArguments(**fine_tune_trainer_kwargs, **pipeline_specific_trainer_kwargs)  # type: ignore
             tokenizer_train_dataset = ChronosFineTuningDataset(
                 target_df=train_data,
                 target_column=self.target,
@@ -529,6 +558,7 @@ class ChronosModel(AbstractTimeSeriesModel):
             if time_limit is not None:
                 callbacks.append(TimeLimitCallback(time_limit=time_limit))
 
+            tokenizer_val_dataset: Optional[ChronosFineTuningDataset] = None
             if val_data is not None:
                 callbacks.append(EvaluateAndSaveFinalStepCallback())
                 # evaluate on a randomly-sampled subset
@@ -544,6 +574,7 @@ class ChronosModel(AbstractTimeSeriesModel):
                     )
                     val_data = val_data.loc[eval_items]
 
+                assert isinstance(val_data, TimeSeriesDataFrame)
                 tokenizer_val_dataset = ChronosFineTuningDataset(
                     target_df=val_data,
                     target_column=self.target,
@@ -557,7 +588,7 @@ class ChronosModel(AbstractTimeSeriesModel):
                 model=self.model_pipeline.inner_model,
                 args=training_args,
                 train_dataset=tokenizer_train_dataset,
-                eval_dataset=tokenizer_val_dataset if val_data is not None else None,
+                eval_dataset=tokenizer_val_dataset,
                 callbacks=callbacks,
             )
 
@@ -590,7 +621,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         num_workers: int = 0,
         time_limit: Optional[float] = None,
     ):
-        from .pipeline.utils import ChronosInferenceDataLoader, ChronosInferenceDataset, timeout_callback
+        from .utils import ChronosInferenceDataLoader, ChronosInferenceDataset, timeout_callback
 
         chronos_dataset = ChronosInferenceDataset(
             target_df=data,
@@ -619,7 +650,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
-        from .pipeline import ChronosBoltPipeline
+        from chronos import ChronosBoltPipeline, ChronosPipeline
 
         # 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
@@ -631,20 +662,26 @@ class ChronosModel(AbstractTimeSeriesModel):
         # (according to its config.json file) of 512, it will further truncate the series during inference.
         context_length = self._get_context_length(data)
 
+        extra_predict_kwargs = (
+            {"num_samples": self.num_samples} if isinstance(self.model_pipeline, ChronosPipeline) else {}
+        )
+
+        # adapt batch size for Chronos bolt if requested prediction length is longer than model prediction length
+        batch_size = self.batch_size
+        model_prediction_length = None
+        if isinstance(self.model_pipeline, ChronosBoltPipeline):
+            model_prediction_length = self.model_pipeline.model.config.chronos_config.get("prediction_length")
+        if model_prediction_length and self.prediction_length > model_prediction_length:
+            batch_size = max(1, batch_size // 4)
+            logger.debug(
+                f"\tThe prediction_length {self.prediction_length} exceeds model's prediction_length {model_prediction_length}. "
+                f"The inference batch_size has been reduced from {self.batch_size} to {batch_size} to avoid OOM errors."
+            )
+
         with warning_filter(all_warnings=True):
             import torch
 
             self.model_pipeline.model.eval()
-            batch_size = self.batch_size
-            if (
-                isinstance(self.model_pipeline, ChronosBoltPipeline)
-                and self.prediction_length > self.model_pipeline.model_prediction_length
-            ):
-                batch_size = max(1, batch_size // 4)
-                logger.debug(
-                    f"\tThe prediction_length {self.prediction_length} exceeds model's prediction_length {self.model_pipeline.model_prediction_length}. "
-                    f"The inference batch_size has been reduced from {self.batch_size} to {batch_size} to avoid OOM errors."
-                )
 
             inference_data_loader = self._get_inference_data_loader(
                 data=data,
@@ -662,7 +699,7 @@ class ChronosModel(AbstractTimeSeriesModel):
                             batch,
                             prediction_length=self.prediction_length,
                             quantile_levels=self.quantile_levels,
-                            num_samples=self.num_samples,
+                            **extra_predict_kwargs,
                         )
                     except torch.OutOfMemoryError as ex:
                         logger.error(
@@ -687,7 +724,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
         return TimeSeriesDataFrame(df)
 
-    def _more_tags(self) -> Dict:
+    def _more_tags(self) -> dict:
         do_fine_tune = self.get_hyperparameters()["fine_tune"]
         return {
             "allow_nan": True,

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

@@ -1,26 +1,23 @@
 import logging
-import os
-import re
 import time
 from itertools import chain, cycle
-from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Iterable, Iterator, List, Literal, Optional
+from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Literal, Optional
 
 import numpy as np
 import torch
+from chronos.chronos_bolt import ChronosBoltModelForForecasting, ResidualBlock
 from gluonts.dataset.field_names import FieldName
 from gluonts.transform import ExpectedNumInstanceSampler, InstanceSplitter, ValidationSplitSampler
 from torch.utils.data import IterableDataset
 from transformers import TrainerCallback
 
-from autogluon.common.loaders.load_s3 import download, list_bucket_prefix_suffix_contains_s3
 from autogluon.core.utils.exceptions import TimeLimitExceeded
-from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.models.gluonts.dataset import SimpleGluonTSDataset
 
 if TYPE_CHECKING:
     # TODO: fix the underlying reason for this circular import, the pipeline should handle tokenization
-    from autogluon.timeseries.models.chronos.pipeline.chronos import ChronosTokenizer
+    from chronos import ChronosTokenizer
 
 
 logger = logging.getLogger("autogluon.timeseries.models.chronos")
@@ -73,19 +70,19 @@ class ChronosFineTuningDataset(IterableDataset):
 
     Parameters
     ----------
-    target_df : TimeSeriesDataFrame
+    target_df
         The ``TimeSeriesDataFrame`` to be converted
-    target_column : str, default = "target"
+    target_column
         The name of the column which contains the target time series, by default "target"
-    context_length : int, default = 512
+    context_length
         The length of the historical context
-    prediction_length : int, default = 64
+    prediction_length
         The prediction_length, i.e., length of label or target
-    tokenizer : ``ChronosTokenizer``, default = None
+    tokenizer
         When a ``ChronosTokenizer`` object is provided, data will be converted into the
         HuggingFace format accepted by the original Chronos models using this ``ChronosTokenizer``.
         If None, data will be converted into the format accepted by ChronosBolt models.
-    mode : Literal["training", "validation"], default = "training"
+    mode
         When ``training``, random slices from the time series will be returned for training purposes.
         If ``validation``, the last slice of each time series returned in the original order.
     """
@@ -132,11 +129,11 @@ class ChronosFineTuningDataset(IterableDataset):
     def _create_training_data(self, data: Iterable[dict]):
         data = chain.from_iterable(cycle([data]))
         split_transform = self._create_instance_splitter("training")
-        data = split_transform.apply(data, is_train=True)
+        data = split_transform.apply(data, is_train=True)  # type: ignore
         return data
 
     def _create_validation_data(self, data: Iterable[dict]):
-        data = self._create_instance_splitter("validation").apply(data, is_train=False)
+        data = self._create_instance_splitter("validation").apply(data, is_train=False)  # type: ignore
         return data
 
     def to_chronos_format(self, entry: dict) -> dict:
@@ -145,7 +142,7 @@ class ChronosFineTuningDataset(IterableDataset):
 
         Parameters
         ----------
-        entry : dict
+        entry
             time series data entry in GluonTS format with ``past_target`` and ``future_target`` keys
 
         Returns
@@ -172,7 +169,7 @@ class ChronosFineTuningDataset(IterableDataset):
 
         Parameters
         ----------
-        entry : dict
+        entry
             time series data entry in GluonTS format with ``past_target`` and ``future_target`` keys
 
         Returns
@@ -190,6 +187,8 @@ class ChronosFineTuningDataset(IterableDataset):
             iterable = self._create_training_data(self.gluonts_dataset)
         elif self.mode == "validation":
             iterable = self._create_validation_data(self.gluonts_dataset)
+        else:
+            raise ValueError(f"Unknown mode {self.mode}")
 
         format_transform_fn = self.to_chronos_format if self.tokenizer is not None else self.to_chronos_bolt_format
         for entry in iterable:
@@ -200,7 +199,7 @@ class ChronosFineTuningDataset(IterableDataset):
 
         Parameters
         ----------
-        shuffle_buffer_size : int, optional, default = None
+        shuffle_buffer_size
             The shuffle buffer size used for pseudo shuffling
         """
         assert shuffle_buffer_size is None or shuffle_buffer_size >= 0
@@ -209,7 +208,7 @@ class ChronosFineTuningDataset(IterableDataset):
         return PseudoShuffledIterableDataset(self, shuffle_buffer_size)
 
 
-def left_pad_and_stack_1D(tensors: List[torch.Tensor]) -> torch.Tensor:
+def left_pad_and_stack_1D(tensors: list[torch.Tensor]) -> torch.Tensor:
     max_len = max(len(c) for c in tensors)
     padded = []
     for c in tensors:
@@ -220,27 +219,6 @@ def left_pad_and_stack_1D(tensors: List[torch.Tensor]) -> torch.Tensor:
     return torch.stack(padded)
 
 
-def cache_model_from_s3(s3_uri: str, force=False):
-    if re.match("^s3://([^/]+)/(.*?([^/]+)/?)$", s3_uri) is None:
-        raise ValueError(f"Not a valid S3 URI: {s3_uri}")
-
-    # we expect the prefix to point to a "directory" on S3
-    if not s3_uri.endswith("/"):
-        s3_uri += "/"
-
-    cache_home = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
-    bucket, prefix = s3_uri.replace("s3://", "").split("/", 1)
-    bucket_cache_path = cache_home / "autogluon" / "timeseries" / bucket
-
-    for obj_path in list_bucket_prefix_suffix_contains_s3(bucket=bucket, prefix=prefix):
-        destination_path = bucket_cache_path / obj_path
-        if not force and destination_path.exists():
-            continue
-        download(bucket, obj_path, local_path=str(destination_path))
-
-    return str(bucket_cache_path / prefix)
-
-
 class ChronosInferenceDataset:
     """A container for time series datasets that implements the ``torch.utils.data.Dataset`` interface"""
 
@@ -280,7 +258,7 @@ class ChronosInferenceDataLoader(torch.utils.data.DataLoader):
         self.callback: Callable = kwargs.pop("on_batch", lambda: None)
         super().__init__(*args, **kwargs)
 
-    def __iter__(self):
+    def __iter__(self):  # type: ignore
         for item in super().__iter__():
             yield item
             self.callback()
@@ -297,13 +275,13 @@ class EvaluateAndSaveFinalStepCallback(TrainerCallback):
 
 
 class TimeLimitCallback(TrainerCallback):
-    def __init__(self, time_limit: int):
+    def __init__(self, time_limit: float):
         """
         Callback to stop training once a specified time has elapsed.
 
         Parameters
         ----------
-        time_limit: int
+        time_limit
             maximum time allowed for training in seconds.
         """
         self.time_limit = time_limit
@@ -321,7 +299,8 @@ class TimeLimitCallback(TrainerCallback):
 
 class LoggerCallback(TrainerCallback):
     def on_log(self, args, state, control, logs=None, **kwargs):
-        logs.pop("total_flos", None)
+        if logs:
+            logs.pop("total_flos", None)
         if state.is_local_process_zero:
             logger.info(logs)
 
@@ -335,3 +314,56 @@ def timeout_callback(seconds: Optional[float]) -> Callable:
             raise TimeLimitExceeded
 
     return callback
+
+
+def update_output_quantiles(model: ChronosBoltModelForForecasting, new_quantiles: list[float]) -> None:
+    """In-place updates model's output layer to support only the specified new quantiles by copying
+    weights from closest existing quantiles.
+    """
+    old_quantiles = model.chronos_config.quantiles
+    new_quantiles = sorted(new_quantiles)
+
+    if new_quantiles == old_quantiles:
+        return
+
+    model.chronos_config.quantiles = new_quantiles
+    model.num_quantiles = len(new_quantiles)
+    model.register_buffer("quantiles", torch.tensor(new_quantiles, dtype=model.dtype), persistent=False)
+
+    old_output_layer = model.output_patch_embedding
+    new_output_layer = ResidualBlock(
+        in_dim=model.config.d_model,
+        h_dim=model.config.d_ff,
+        out_dim=len(new_quantiles) * model.chronos_config.prediction_length,
+        act_fn_name=model.config.dense_act_fn,
+        dropout_p=model.config.dropout_rate,
+    )
+
+    # hidden_layer is shared across all quantiles
+    new_output_layer.hidden_layer.weight.data.copy_(old_output_layer.hidden_layer.weight.data)
+    if old_output_layer.hidden_layer.bias is not None:
+        new_output_layer.hidden_layer.bias.data.copy_(old_output_layer.hidden_layer.bias.data)
+
+    def copy_quantile_weights(src_idx: int, dst_idx: int):
+        """Copy weights for one quantile from src_idx to dst_idx"""
+        prediction_length = model.chronos_config.prediction_length
+        src_start, src_end = src_idx * prediction_length, (src_idx + 1) * prediction_length
+        dst_start, dst_end = dst_idx * prediction_length, (dst_idx + 1) * prediction_length
+
+        for layer_name in ["output_layer", "residual_layer"]:
+            old_layer_attr = getattr(old_output_layer, layer_name)
+            new_layer_attr = getattr(new_output_layer, layer_name)
+
+            new_layer_attr.weight[dst_start:dst_end] = old_layer_attr.weight[src_start:src_end]
+            if old_layer_attr.bias is not None:
+                new_layer_attr.bias[dst_start:dst_end] = old_layer_attr.bias[src_start:src_end]
+
+    with torch.no_grad():
+        for new_idx, new_q in enumerate(new_quantiles):
+            closest_q = min(old_quantiles, key=lambda x: abs(x - new_q))
+            closest_idx = old_quantiles.index(closest_q)
+            copy_quantile_weights(closest_idx, new_idx)
+
+    model.output_patch_embedding = new_output_layer
+    model.config.chronos_config["quantiles"] = new_quantiles
+    model.chronos_config.quantiles = new_quantiles

autogluon/timeseries/models/ensemble/__init__.py

@@ -1,3 +1,30 @@
 from .abstract import AbstractTimeSeriesEnsembleModel
-from .basic import PerformanceWeightedEnsemble, SimpleAverageEnsemble
-from .greedy import GreedyEnsemble
+from .array_based import MedianEnsemble, PerQuantileTabularEnsemble, TabularEnsemble
+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,
+    }
+    if name not in mapping:
+        raise ValueError(f"Unknown ensemble type: {name}. Available: {list(mapping.keys())}")
+    return mapping[name]
+
+
+__all__ = [
+    "AbstractTimeSeriesEnsembleModel",
+    "GreedyEnsemble",
+    "MedianEnsemble",
+    "PerformanceWeightedEnsemble",
+    "PerQuantileTabularEnsemble",
+    "SimpleAverageEnsemble",
+    "TabularEnsemble",
+    "get_ensemble_class",
+]