autogluon.timeseries 1.4.1b20250907__py3-none-any.whl → 1.5.1b20260122__py3-none-any.whl

autogluon/timeseries/models/toto/hf_pretrained_model.py

@@ -0,0 +1,200 @@
+import json
+import logging
+import os
+from pathlib import Path
+
+from transformers import PretrainedConfig, PreTrainedModel
+
+from ._internal.backbone import TotoBackbone
+
+
+class TotoConfig(PretrainedConfig):
+    model_type = "toto"
+
+    def __init__(
+        self,
+        dropout: float = 0.0,
+        embed_dim: int = 768,
+        num_heads: int = 12,
+        num_layers: int = 12,
+        output_distribution_classes: list[str] | None = None,
+        output_distribution_kwargs: dict | None = None,
+        patch_size: int = 64,
+        scale_factor_exponent: float = 10.0,
+        spacewise_every_n_layers: int = 12,
+        spacewise_first: bool = False,
+        stabilize_with_global: bool = True,
+        stride: int = 64,
+        transformers_version: str = "4.49.0",
+        use_memory_efficient_attention: bool = False,
+        **kwargs,
+    ):
+        self.dropout = dropout
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.output_distribution_classes = output_distribution_classes or ["MixtureOfStudentTsOutput"]
+        self.output_distribution_kwargs = output_distribution_kwargs or {"k_components": 24}
+        self.patch_size = patch_size
+        self.scale_factor_exponent = scale_factor_exponent
+        self.spacewise_every_n_layers = spacewise_every_n_layers
+        self.spacewise_first = spacewise_first
+        self.stabilize_with_global = stabilize_with_global
+        self.stride = stride
+        self.transformers_version = transformers_version
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+
+        super().__init__(**kwargs)
+
+
+class TotoPretrainedModel(PreTrainedModel):
+    config_class = TotoConfig
+    base_model_prefix = "model"  # optional, used for weight naming conventions
+
+    def __init__(self, config: TotoConfig):
+        super().__init__(config)
+        self.model = TotoBackbone(
+            patch_size=config.patch_size,
+            stride=config.stride,
+            embed_dim=config.embed_dim,
+            num_layers=config.num_layers,
+            num_heads=config.num_heads,
+            mlp_hidden_dim=getattr(config, "mlp_hidden_dim", 3072),
+            dropout=config.dropout,
+            spacewise_every_n_layers=config.spacewise_every_n_layers,
+            scaler_cls=getattr(config, "scaler_cls", "model.scaler.CausalPatchStdMeanScaler"),
+            output_distribution_classes=config.output_distribution_classes,
+            spacewise_first=config.spacewise_first,
+            output_distribution_kwargs=config.output_distribution_kwargs,
+            use_memory_efficient_attention=False,
+            stabilize_with_global=config.stabilize_with_global,
+            scale_factor_exponent=config.scale_factor_exponent,
+            **getattr(config, "extra_kwargs", {}),
+        )
+        self.post_init()
+
+    @staticmethod
+    def _remap_state_dict_keys(state_dict):
+        remap = {
+            "mlp.0.w12.weight": "mlp.0.weight",
+            "mlp.0.w12.bias": "mlp.0.bias",
+            "mlp.0.w3.weight": "mlp.2.weight",
+            "mlp.0.w3.bias": "mlp.2.bias",
+        }
+
+        new_state = {}
+        keys_to_remap = []
+        for key in list(state_dict.keys()):
+            for old, new in remap.items():
+                if old in key:
+                    new_key = key.replace(old, new)
+                    keys_to_remap.append((key, new_key))
+                    break
+
+        new_state = state_dict.copy()
+        for old_key, new_key in keys_to_remap:
+            new_state[new_key] = new_state.pop(old_key)
+
+        return new_state
+
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint_path,
+        device_map: str = "cpu",
+        strict=True,
+        **model_kwargs,
+    ):
+        """
+        Custom checkpoint loading. Used to load a local
+        safetensors checkpoint with an optional config.json file.
+        """
+        import safetensors.torch as safetorch
+
+        if os.path.isdir(checkpoint_path):
+            safetensors_file = os.path.join(checkpoint_path, "model.safetensors")
+        else:
+            safetensors_file = checkpoint_path
+
+        if os.path.exists(safetensors_file):
+            model_state = safetorch.load_file(safetensors_file, device=device_map)
+        else:
+            raise FileNotFoundError(f"Model checkpoint not found at: {safetensors_file}")
+
+        # Load configuration from config.json if it exists.
+        config_file = os.path.join(checkpoint_path, "config.json")
+        config = {}
+        if os.path.exists(config_file):
+            with open(config_file, "r") as f:
+                config = json.load(f)
+
+        # Merge any extra kwargs into the configuration.
+        config.update(model_kwargs)
+
+        remapped_state_dict = cls._remap_state_dict_keys(model_state)
+
+        instance = cls(**config)
+
+        # Filter out unexpected keys
+        filtered_remapped_state_dict = {
+            k: v
+            for k, v in remapped_state_dict.items()
+            if k in instance.state_dict() and not k.endswith("rotary_emb.freqs")
+        }
+
+        instance.load_state_dict(filtered_remapped_state_dict, strict=strict)
+        instance.to(device_map)  # type: ignore
+
+        return instance
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        *,
+        model_id: str,
+        revision: str | None = None,
+        cache_dir: Path | str | None = None,
+        force_download: bool = False,
+        proxies: dict | None = None,
+        resume_download: bool | None = None,
+        local_files_only: bool = False,
+        token: str | bool | None = None,
+        device_map: str = "cpu",
+        strict: bool = False,
+        **model_kwargs,
+    ):
+        """Load Pytorch pretrained weights and return the loaded model."""
+        from huggingface_hub import constants, hf_hub_download
+
+        transformers_logger = logging.getLogger("transformers.modeling_utils")
+        original_level = transformers_logger.level
+
+        try:
+            # Here we suppress transformers logger's "some weights were not initialized" error since the
+            # remapping hook is only called after the initial model loading.
+            transformers_logger.setLevel(logging.ERROR)
+
+            if os.path.isdir(model_id):
+                print("Loading weights from local directory")
+                model_file = os.path.join(model_id, constants.SAFETENSORS_SINGLE_FILE)
+                model = cls.load_from_checkpoint(model_file, device_map, strict, **model_kwargs)
+            else:
+                model_file = hf_hub_download(
+                    repo_id=model_id,
+                    filename=constants.SAFETENSORS_SINGLE_FILE,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    token=token,
+                    local_files_only=local_files_only,
+                )
+                model = cls.load_from_checkpoint(model_file, device_map, strict, **model_kwargs)
+        finally:
+            transformers_logger.setLevel(original_level)
+
+        return model
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)

autogluon/timeseries/models/toto/model.py

@@ -0,0 +1,249 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Any, Sequence
+
+import numpy as np
+import pandas as pd
+from typing_extensions import Self
+
+from autogluon.common.loaders import load_pkl
+from autogluon.timeseries import TimeSeriesDataFrame
+from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
+from autogluon.timeseries.utils.features import CovariateMetadata
+
+if TYPE_CHECKING:
+    from ._internal import TotoForecaster
+
+logger = logging.getLogger(__name__)
+
+
+class TotoModel(AbstractTimeSeriesModel):
+    """Toto (Time-Series-Optimized Transformer for Observability) [CohenKhwajaetal2025]_ pretrained time series forecasting model.
+
+    Toto is a 151M parameter model trained on over 1T data points from DataDog's internal observability systems, as well as
+    the GIFT-eval pretrain, Chronos pretraining, and synthetically generated time series corpora. It is a decoder-only
+    architecture that autoregressively outputs parametric distribution forecasts. More details can be found on
+    `Hugging Face <https://huggingface.co/Datadog/Toto-Open-Base-1.0>`_ and `GitHub <https://github.com/DataDog/toto>`_.
+
+    The AutoGluon implementation of Toto is on a port of the original implementation. AutoGluon supports Toto for
+    **inference only**, i.e., the model will not be trained or fine-tuned on the provided training data. Toto is optimized
+    for easy maintenance with the rest of the AutoGluon model zoo, and does not feature some important optimizations such
+    as xformers and flash-attention available in the original model repository. The AutoGluon implementation of Toto
+    requires a CUDA-compatible GPU.
+
+    References
+    ----------
+    .. [CohenKhwajaetal2025] Cohen, Ben, Khwaja, Emaad et al.
+        "This Time is Different: An Observability Perspective on Time Series Foundation Models."
+        https://arxiv.org/abs/2505.14766
+
+
+    Other Parameters
+    ----------------
+    model_path : str, default = "Datadog/Toto-Open-Base-1.0"
+        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.
+    batch_size : int, default = 24
+        Size of batches used during inference.
+    num_samples : int, default = 256
+        Number of samples used during inference.
+    device : str, default = "cuda"
+        Device to use for inference. Toto requires a CUDA-compatible GPU to run.
+    context_length : int or None, default = 4096
+        The context length to use in the model. Shorter context lengths will decrease model accuracy, but result
+        in faster inference.
+    compile_model : bool, default = True
+        Whether to compile the model using torch.compile() for faster inference. May increase initial loading time
+        but can provide speedups during inference.
+    """
+
+    default_model_path: str = "Datadog/Toto-Open-Base-1.0"
+
+    def __init__(
+        self,
+        path: str | None = None,
+        name: str | None = None,
+        hyperparameters: dict[str, Any] | None = None,
+        freq: str | None = None,
+        prediction_length: int = 1,
+        covariate_metadata: CovariateMetadata | None = None,
+        target: str = "target",
+        quantile_levels: Sequence[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
+        eval_metric: Any = None,
+    ):
+        hyperparameters = hyperparameters if hyperparameters is not None else {}
+
+        self.model_path = hyperparameters.get("model_path", self.default_model_path)
+
+        super().__init__(
+            path=path,
+            name=name,
+            hyperparameters=hyperparameters,
+            freq=freq,
+            prediction_length=prediction_length,
+            covariate_metadata=covariate_metadata,
+            target=target,
+            quantile_levels=quantile_levels,
+            eval_metric=eval_metric,
+        )
+
+        self._forecaster: TotoForecaster | None = None
+
+    def save(self, path: str | None = None, verbose: bool = True) -> str:
+        forecaster = self._forecaster
+        self._forecaster = None
+        path = super().save(path=path, verbose=verbose)
+        self._forecaster = forecaster
+
+        return str(path)
+
+    @classmethod
+    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)
+
+        return model
+
+    def _is_gpu_available(self) -> bool:
+        import torch.cuda
+
+        return torch.cuda.is_available()
+
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, int | float]:
+        return {"num_cpus": 1, "num_gpus": 1}
+
+    def load_forecaster(self):
+        from ._internal import TotoForecaster
+        from .hf_pretrained_model import TotoConfig, TotoPretrainedModel
+
+        if not self._is_gpu_available():
+            raise RuntimeError(
+                f"{self.name} requires a GPU to run, but no GPU was detected. "
+                "Please make sure that you are using a computer with a CUDA-compatible GPU and "
+                "`import torch; torch.cuda.is_available()` returns `True`."
+            )
+
+        hyperparameters = self.get_hyperparameters()
+        pretrained_model = TotoPretrainedModel.from_pretrained(
+            model_id=self.model_path,
+            config=TotoConfig.from_pretrained(self.model_path),
+            device_map=hyperparameters["device"],
+        )
+
+        if hyperparameters["compile_model"]:
+            pretrained_model.model.compile()
+
+        self._forecaster = TotoForecaster(model=pretrained_model.model)
+
+    def persist(self) -> Self:
+        if self._forecaster is None:
+            self.load_forecaster()
+        return self
+
+    def _get_default_hyperparameters(self) -> dict:
+        return {
+            "batch_size": 24,
+            "num_samples": 256,
+            "device": "cuda",
+            "context_length": 4096,
+            "compile_model": False,
+        }
+
+    def _get_sample_batch_size(self) -> int:
+        num_samples = self.get_hyperparameter("num_samples")
+        batch_size = num_samples
+        while batch_size > 32:
+            for factor in range(2, int(batch_size**0.5) + 1):
+                if batch_size % factor == 0:
+                    batch_size //= factor
+                    break
+            else:  # batch_size is prime
+                return batch_size
+        return batch_size
+
+    @property
+    def allowed_hyperparameters(self) -> list[str]:
+        return super().allowed_hyperparameters + [
+            "model_path",
+            "batch_size",
+            "num_samples",
+            "device",
+            "context_length",
+            "compile_model",
+        ]
+
+    def _more_tags(self) -> dict:
+        return {
+            "allow_nan": True,
+            "can_use_train_data": False,
+            "can_use_val_data": False,
+        }
+
+    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.load_forecaster()
+
+    def _predict(
+        self, data: TimeSeriesDataFrame, known_covariates: TimeSeriesDataFrame | None = None, **kwargs
+    ) -> TimeSeriesDataFrame:
+        import torch
+
+        from .dataloader import TotoDataLoader, TotoInferenceDataset
+
+        hyperparameters = self.get_hyperparameters()
+
+        if self._forecaster is None:
+            self.load_forecaster()
+        assert self._forecaster, "Toto model failed to load"
+        device = self._forecaster.model.device
+
+        dataset = TotoInferenceDataset(
+            target_df=data.fill_missing_values("auto"),
+            max_context_length=hyperparameters["context_length"],
+            target_column=self.target,
+        )
+        loader = TotoDataLoader(
+            dataset,
+            freq=self.freq,
+            batch_size=hyperparameters["batch_size"],
+            time_limit=kwargs.get("time_limit"),
+            device=device,
+        )
+
+        batch_means, batch_quantiles = [], []
+        with torch.inference_mode():
+            for masked_timeseries in loader:
+                forecast = self._forecaster.forecast(
+                    masked_timeseries,
+                    prediction_length=self.prediction_length,
+                    num_samples=hyperparameters["num_samples"],
+                    samples_per_batch=self._get_sample_batch_size(),
+                )
+
+                batch_means.append(forecast.mean.cpu().numpy())
+                qs = np.array([forecast.quantile(q).cpu().numpy() for q in self.quantile_levels])
+                batch_quantiles.append(qs.squeeze(2).transpose(1, 2, 0))
+
+        df = pd.DataFrame(
+            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=self.get_forecast_horizon_index(data),
+        )
+
+        return TimeSeriesDataFrame(df)