import logging
import os
from pathlib import Path
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
from autogluon.timeseries.utils.warning_filters import warning_filter
logger = logging.getLogger(__name__)
[docs]
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
_supports_export = 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):
device = self.get_hyperparameter("device") or ("cuda" if self._is_gpu_available() else "cpu")
assert self.model_path is not None
# importing chronos triggers `import tqdm.auto`, which emits a TqdmWarning if ipywidgets is missing
with warning_filter(all_warnings=True):
from chronos.chronos2.pipeline import Chronos2Pipeline
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 export_model(self, path: str | Path) -> str:
"""Export the model as a Hugging Face checkpoint that can be loaded with
``chronos.chronos2.Chronos2Pipeline.from_pretrained``.
If the model was fine-tuned with LoRA, the adapter weights are merged into the base model, so that the
exported checkpoint is self-contained.
See `TimeSeriesPredictor.export_model` for details.
"""
from chronos.chronos2.pipeline import Chronos2Pipeline
self._assert_no_transforms_for_export()
path = Path(path)
# after fine-tuning, `model_path` points at a checkpoint that only contains the LoRA adapter, and the
# pipeline in memory holds a `PeftModel`. Reloading merges the adapter into the base model, so that the
# exported checkpoint contains the full weights and is self-contained.
Chronos2Pipeline.from_pretrained(self.model_path).save_pretrained(path)
return str(path)
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.chronos2 import preprocess
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)
if (
self.get_hyperparameter("disable_known_covariates")
or len(self.covariate_metadata.known_covariates) == 0
):
known_covariates_names = None
else:
known_covariates_names = self.covariate_metadata.known_covariates
return preprocess.from_data_frame(
past_df,
target_columns=[self.target],
prediction_length=self.prediction_length,
validate_inputs=False,
known_covariates_names=known_covariates_names,
)
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()