autogluon.timeseries 1.4.1b20250820__py3-none-any.whl → 1.4.1b20250902__py3-none-any.whl

autogluon/timeseries/models/chronos/pipeline/base.py

@@ -2,7 +2,7 @@
 
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import torch
 
@@ -18,7 +18,7 @@ class ForecastType(Enum):
 
 
 class PipelineRegistry(type):
-    REGISTRY: Dict[str, "PipelineRegistry"] = {}
+    REGISTRY: dict[str, "PipelineRegistry"] = {}
 
     def __new__(cls, name, bases, attrs):
         """See, https://github.com/faif/python-patterns."""
@@ -43,13 +43,13 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
         """
         Parameters
         ----------
-        inner_model : PreTrainedModel
+        inner_model
             A hugging-face transformers PreTrainedModel, e.g., T5ForConditionalGeneration
         """
         # for easy access to the inner HF-style model
         self.inner_model = inner_model
 
-    def _prepare_and_validate_context(self, context: Union[torch.Tensor, List[torch.Tensor]]):
+    def _prepare_and_validate_context(self, context: Union[torch.Tensor, list[torch.Tensor]]):
         if isinstance(context, list):
             context = left_pad_and_stack_1D(context)
         assert isinstance(context, torch.Tensor)
@@ -61,7 +61,7 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
 
     def predict(
         self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
+        context: Union[torch.Tensor, list[torch.Tensor]],
         prediction_length: Optional[int] = None,
         **kwargs,
     ):
@@ -88,8 +88,8 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
         raise NotImplementedError()
 
     def predict_quantiles(
-        self, context: torch.Tensor, prediction_length: int, quantile_levels: List[float], **kwargs
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, context: torch.Tensor, prediction_length: int, quantile_levels: list[float], **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Get quantile and mean forecasts for given time series. All
         predictions are returned on the CPU.
@@ -104,7 +104,7 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
         prediction_length
             Time steps to predict. Defaults to a model-dependent
             value if not given.
-        quantile_levels: List[float]
+        quantile_levels
             Quantile levels to compute
 
         Returns

autogluon/timeseries/models/chronos/pipeline/chronos.py

@@ -7,7 +7,7 @@
 import logging
 import warnings
 from dataclasses import dataclass
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Literal, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -29,7 +29,7 @@ class ChronosConfig:
     """
 
     tokenizer_class: str
-    tokenizer_kwargs: Dict[str, Any]
+    tokenizer_kwargs: dict[str, Any]
     n_tokens: int
     n_special_tokens: int
     pad_token_id: int
@@ -66,7 +66,7 @@ class ChronosTokenizer:
     def context_input_transform(
         self,
         context: torch.Tensor,
-    ) -> Tuple:
+    ) -> tuple:
         """
         Turn a batch of time series into token IDs, attention mask, and tokenizer_state.
 
@@ -95,7 +95,7 @@ class ChronosTokenizer:
         """
         raise NotImplementedError()
 
-    def label_input_transform(self, label: torch.Tensor, tokenizer_state: Any) -> Tuple:
+    def label_input_transform(self, label: torch.Tensor, tokenizer_state: Any) -> tuple:
         """
         Turn a batch of label slices of time series into token IDs and attention mask
         using the ``tokenizer_state`` provided by ``context_input_transform``.
@@ -171,7 +171,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
     def _input_transform(
         self, context: torch.Tensor, scale: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         attention_mask = ~torch.isnan(context)
 
         if scale is None:
@@ -196,7 +196,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
     def _append_eos_token(
         self, token_ids: torch.Tensor, attention_mask: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size = token_ids.shape[0]
         eos_tokens = torch.full((batch_size, 1), fill_value=self.config.eos_token_id)
         token_ids = torch.concat((token_ids, eos_tokens), dim=1)
@@ -205,7 +205,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
         return token_ids, attention_mask
 
-    def context_input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def context_input_transform(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         length = context.shape[-1]
 
         if length > self.config.context_length:
@@ -218,7 +218,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
         return token_ids, attention_mask, scale
 
-    def label_input_transform(self, label: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def label_input_transform(self, label: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         token_ids, attention_mask, _ = self._input_transform(context=label, scale=scale)
 
         if self.config.use_eos_token:
@@ -371,7 +371,7 @@ class ChronosPipeline(BaseChronosPipeline):
         self.model = model
 
     @torch.no_grad()
-    def embed(self, context: Union[torch.Tensor, List[torch.Tensor]]) -> Tuple[torch.Tensor, Any]:
+    def embed(self, context: Union[torch.Tensor, list[torch.Tensor]]) -> tuple[torch.Tensor, Any]:
         """
         Get encoder embeddings for the given time series.
 
@@ -404,7 +404,7 @@ class ChronosPipeline(BaseChronosPipeline):
 
     def predict(
         self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
+        context: Union[torch.Tensor, list[torch.Tensor]],
         prediction_length: Optional[int] = None,
         num_samples: Optional[int] = None,
         temperature: Optional[float] = None,
@@ -494,10 +494,10 @@ class ChronosPipeline(BaseChronosPipeline):
         self,
         context: torch.Tensor,
         prediction_length: int,
-        quantile_levels: List[float],
+        quantile_levels: list[float],
         num_samples: Optional[int] = None,
         **kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         num_samples = num_samples or self.model.config.num_samples
         prediction_samples = (
             self.predict(

autogluon/timeseries/models/chronos/pipeline/chronos_bolt.py

@@ -7,7 +7,7 @@ import copy
 import logging
 import warnings
 from dataclasses import dataclass, fields
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -32,7 +32,7 @@ class ChronosBoltConfig:
     prediction_length: int
     input_patch_size: int
     input_patch_stride: int
-    quantiles: List[float]
+    quantiles: list[float]
     use_reg_token: bool = False
 
 
@@ -77,8 +77,8 @@ class InstanceNorm(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        loc_scale: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        loc_scale: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         if loc_scale is None:
             loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0)
             scale = torch.nan_to_num((x - loc).square().nanmean(dim=-1, keepdim=True).sqrt(), nan=1.0)
@@ -88,7 +88,7 @@ class InstanceNorm(nn.Module):
 
         return (x - loc) / scale, (loc, scale)
 
-    def inverse(self, x: torch.Tensor, loc_scale: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    def inverse(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
         loc, scale = loc_scale
         return x * scale + loc
 
@@ -343,11 +343,11 @@ class ChronosBoltModelForForecasting(T5PreTrainedModel):
         """
         Parameters
         ----------
-        input_embeds: torch.Tensor
+        input_embeds
             Patched and embedded inputs. Shape (batch_size, patched_context_length, d_model)
-        attention_mask: torch.Tensor
+        attention_mask
             Attention mask for the patched context. Shape (batch_size, patched_context_length), type: torch.int64
-        hidden_states: torch.Tensor
+        hidden_states
             Hidden states returned by the encoder. Shape (batch_size, patched_context_length, d_model)
 
         Returns
@@ -385,12 +385,12 @@ class ChronosBoltPipeline(BaseChronosPipeline):
         self.model_prediction_length: int = self.model.config.chronos_config["prediction_length"]
 
     @property
-    def quantiles(self) -> List[float]:
+    def quantiles(self) -> list[float]:
         return self.model.config.chronos_config["quantiles"]
 
     def predict(  # type: ignore[override]
         self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
+        context: Union[torch.Tensor, list[torch.Tensor]],
         prediction_length: Optional[int] = None,
         limit_prediction_length: bool = False,
     ):
@@ -458,8 +458,8 @@ class ChronosBoltPipeline(BaseChronosPipeline):
         return torch.cat(predictions, dim=-1)[..., :prediction_length]
 
     def predict_quantiles(
-        self, context: torch.Tensor, prediction_length: int, quantile_levels: List[float], **kwargs
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, context: torch.Tensor, prediction_length: int, quantile_levels: list[float], **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # shape (batch_size, prediction_length, len(training_quantile_levels))
         predictions = (
             self.predict(

autogluon/timeseries/models/chronos/pipeline/utils.py

@@ -4,7 +4,7 @@ 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
@@ -73,19 +73,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.
     """
@@ -145,7 +145,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 +172,7 @@ class ChronosFineTuningDataset(IterableDataset):
 
         Parameters
         ----------
-        entry : dict
+        entry
             time series data entry in GluonTS format with ``past_target`` and ``future_target`` keys
 
         Returns
@@ -200,7 +200,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 +209,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:
@@ -303,7 +303,7 @@ class TimeLimitCallback(TrainerCallback):
 
         Parameters
         ----------
-        time_limit: int
+        time_limit
             maximum time allowed for training in seconds.
         """
         self.time_limit = time_limit

autogluon/timeseries/models/ensemble/abstract.py

@@ -1,7 +1,7 @@
 import functools
 import logging
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from typing import Optional
 
 import numpy as np
 from typing_extensions import final
@@ -18,32 +18,32 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
 
     @property
     @abstractmethod
-    def model_names(self) -> List[str]:
+    def model_names(self) -> list[str]:
         """Names of base models included in the ensemble."""
         pass
 
     @final
     def fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
     ):
         """Fit ensemble model given predictions of candidate base models and the true data.
 
         Parameters
         ----------
-        predictions_per_window : Dict[str, List[TimeSeriesDataFrame]]
+        predictions_per_window
             Dictionary that maps the names of component models to their respective predictions for each validation
             window.
-        data_per_window : List[TimeSeriesDataFrame]
+        data_per_window
             Observed ground truth data used to train the ensemble for each validation window. Each entry in the list
             includes both the forecast horizon (for which the predictions are given in ``predictions``), as well as the
             "history".
-        model_scores : Optional[Dict[str, float]]
+        model_scores
             Scores (higher is better) for the models that will constitute the ensemble.
-        time_limit : Optional[float]
+        time_limit
             Maximum allowed time for training in seconds.
         """
         if time_limit is not None and time_limit <= 0:
@@ -67,9 +67,9 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
 
     def _fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
     ):
         """Private method for `fit`. See `fit` for documentation of arguments. Apart from the model
@@ -78,7 +78,7 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         raise NotImplementedError
 
     @final
-    def predict(self, data: Dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+    def predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
         if not set(self.model_names).issubset(set(data.keys())):
             raise ValueError(
                 f"Set of models given for prediction in {self.name} differ from those provided during initialization."
@@ -93,11 +93,11 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         return self._predict(data=data, **kwargs)
 
     @abstractmethod
-    def _predict(self, data: Dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+    def _predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
         pass
 
     @abstractmethod
-    def remap_base_models(self, model_refit_map: Dict[str, str]) -> None:
+    def remap_base_models(self, model_refit_map: dict[str, str]) -> None:
         """Update names of the base models based on the mapping in model_refit_map.
 
         This method should be called after performing refit_full to point to the refitted base models, if necessary.
@@ -112,17 +112,17 @@ class AbstractWeightedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, A
         if name is None:
             name = "WeightedEnsemble"
         super().__init__(name=name, **kwargs)
-        self.model_to_weight: Dict[str, float] = {}
+        self.model_to_weight: dict[str, float] = {}
 
     @property
-    def model_names(self) -> List[str]:
+    def model_names(self) -> list[str]:
         return list(self.model_to_weight.keys())
 
     @property
     def model_weights(self) -> np.ndarray:
         return np.array(list(self.model_to_weight.values()), dtype=np.float64)
 
-    def _predict(self, data: Dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+    def _predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
         weighted_predictions = [data[model_name] * weight for model_name, weight in self.model_to_weight.items()]
         return functools.reduce(lambda x, y: x + y, weighted_predictions)
 
@@ -131,7 +131,7 @@ class AbstractWeightedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, A
         info["model_weights"] = self.model_to_weight.copy()
         return info
 
-    def remap_base_models(self, model_refit_map: Dict[str, str]) -> None:
+    def remap_base_models(self, model_refit_map: dict[str, str]) -> None:
         updated_weights = {}
         for model, weight in self.model_to_weight.items():
             model_full_name = model_refit_map.get(model, model)

autogluon/timeseries/models/ensemble/basic.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Any, Optional
 
 import numpy as np
 
@@ -17,9 +17,9 @@ class SimpleAverageEnsemble(AbstractWeightedTimeSeriesEnsembleModel):
 
     def _fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
     ):
         self.model_to_weight = {}
@@ -52,14 +52,14 @@ class PerformanceWeightedEnsemble(AbstractWeightedTimeSeriesEnsembleModel):
             name = "PerformanceWeightedEnsemble"
         super().__init__(name=name, **kwargs)
 
-    def _get_default_hyperparameters(self) -> Dict:
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
         return {"weight_scheme": "sqrt"}
 
     def _fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
     ):
         assert model_scores is not None

autogluon/timeseries/models/ensemble/greedy.py

@@ -1,7 +1,7 @@
 import copy
 import logging
 import pprint
-from typing import Dict, List, Optional
+from typing import Any, Optional
 
 import numpy as np
 
@@ -47,14 +47,14 @@ class TimeSeriesEnsembleSelection(EnsembleSelection):
         self.dummy_pred_per_window = []
         self.scorer_per_window = []
 
-        self.dummy_pred_per_window: Optional[List[TimeSeriesDataFrame]]
-        self.scorer_per_window: Optional[List[TimeSeriesScorer]]
-        self.data_future_per_window: Optional[List[TimeSeriesDataFrame]]
+        self.dummy_pred_per_window: Optional[list[TimeSeriesDataFrame]]
+        self.scorer_per_window: Optional[list[TimeSeriesScorer]]
+        self.data_future_per_window: Optional[list[TimeSeriesDataFrame]]
 
     def fit(  # type: ignore
         self,
-        predictions: List[List[TimeSeriesDataFrame]],
-        labels: List[TimeSeriesDataFrame],
+        predictions: list[list[TimeSeriesDataFrame]],
+        labels: list[TimeSeriesDataFrame],
         time_limit: Optional[float] = None,
     ):
         return super().fit(
@@ -65,10 +65,10 @@ class TimeSeriesEnsembleSelection(EnsembleSelection):
 
     def _fit(  # type: ignore
         self,
-        predictions: List[List[TimeSeriesDataFrame]],
-        labels: List[TimeSeriesDataFrame],
+        predictions: list[list[TimeSeriesDataFrame]],
+        labels: list[TimeSeriesDataFrame],
         time_limit: Optional[float] = None,
-        sample_weight: Optional[List[float]] = None,
+        sample_weight: Optional[list[float]] = None,
     ):
         # Stack predictions for each model into a 3d tensor of shape [num_val_windows, num_rows, num_cols]
         stacked_predictions = [np.stack(preds) for preds in predictions]
@@ -157,14 +157,14 @@ class GreedyEnsemble(AbstractWeightedTimeSeriesEnsembleModel):
             name = "WeightedEnsemble"
         super().__init__(name=name, **kwargs)
 
-    def _get_default_hyperparameters(self) -> Dict:
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
         return {"ensemble_size": 100}
 
     def _fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
     ):
         ensemble_selection = TimeSeriesEnsembleSelection(