autogluon.timeseries 1.2.1b20250224__py3-none-any.whl → 1.4.1b20251215__py3-none-any.whl

autogluon/timeseries/models/gluonts/{abstract_gluonts.py → abstract.py}

@@ -3,7 +3,7 @@ import os
 import shutil
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterator, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Type, cast, overload
 
 import gluonts
 import gluonts.core.settings
@@ -11,7 +11,7 @@ import numpy as np
 import pandas as pd
 from gluonts.core.component import from_hyperparameters
 from gluonts.dataset.common import Dataset as GluonTSDataset
-from gluonts.dataset.field_names import FieldName
+from gluonts.env import env as gluonts_env
 from gluonts.model.estimator import Estimator as GluonTSEstimator
 from gluonts.model.forecast import Forecast, QuantileForecast, SampleForecast
 from gluonts.model.predictor import Predictor as GluonTSPredictor
@@ -21,11 +21,15 @@ from autogluon.core.hpo.constants import RAY_BACKEND
 from autogluon.tabular.models.tabular_nn.utils.categorical_encoders import (
     OneHotMergeRaresHandleUnknownEncoder as OneHotEncoder,
 )
-from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TIMESTAMP, TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
-from autogluon.timeseries.utils.datetime import norm_freq_str
 from autogluon.timeseries.utils.warning_filters import disable_root_logger, warning_filter
 
+if TYPE_CHECKING:
+    from gluonts.torch.model.forecast import DistributionForecast
+
+from .dataset import SimpleGluonTSDataset
+
 # NOTE: We avoid imports for torch and lightning.pytorch at the top level and hide them inside class methods.
 # This is done to skip these imports during multiprocessing (which may cause bugs)
 
@@ -33,124 +37,25 @@ logger = logging.getLogger(__name__)
 gts_logger = logging.getLogger(gluonts.__name__)
 
 
-class SimpleGluonTSDataset(GluonTSDataset):
-    """Wrapper for TimeSeriesDataFrame that is compatible with the GluonTS Dataset API."""
-
-    def __init__(
-        self,
-        target_df: TimeSeriesDataFrame,
-        freq: str,
-        target_column: str = "target",
-        feat_static_cat: Optional[np.ndarray] = None,
-        feat_static_real: Optional[np.ndarray] = None,
-        feat_dynamic_cat: Optional[np.ndarray] = None,
-        feat_dynamic_real: Optional[np.ndarray] = None,
-        past_feat_dynamic_cat: Optional[np.ndarray] = None,
-        past_feat_dynamic_real: Optional[np.ndarray] = None,
-        includes_future: bool = False,
-        prediction_length: int = None,
-    ):
-        assert target_df is not None
-        # Convert TimeSeriesDataFrame to pd.Series for faster processing
-        self.target_array = target_df[target_column].to_numpy(np.float32)
-        self.feat_static_cat = self._astype(feat_static_cat, dtype=np.int64)
-        self.feat_static_real = self._astype(feat_static_real, dtype=np.float32)
-        self.feat_dynamic_cat = self._astype(feat_dynamic_cat, dtype=np.int64)
-        self.feat_dynamic_real = self._astype(feat_dynamic_real, dtype=np.float32)
-        self.past_feat_dynamic_cat = self._astype(past_feat_dynamic_cat, dtype=np.int64)
-        self.past_feat_dynamic_real = self._astype(past_feat_dynamic_real, dtype=np.float32)
-        self.freq = self._get_freq_for_period(freq)
-
-        # Necessary to compute indptr for known_covariates at prediction time
-        self.includes_future = includes_future
-        self.prediction_length = prediction_length
-
-        # Replace inefficient groupby ITEMID with indptr that stores start:end of each time series
-        item_id_index = target_df.index.get_level_values(ITEMID)
-        indices_sizes = item_id_index.value_counts(sort=False)
-        self.item_ids = indices_sizes.index  # shape [num_items]
-        cum_sizes = indices_sizes.to_numpy().cumsum()
-        self.indptr = np.append(0, cum_sizes).astype(np.int32)
-        self.start_timestamps = target_df.reset_index(TIMESTAMP).groupby(level=ITEMID, sort=False).first()[TIMESTAMP]
-        assert len(self.item_ids) == len(self.start_timestamps)
-
-    @staticmethod
-    def _astype(array: Optional[np.ndarray], dtype: np.dtype) -> Optional[np.ndarray]:
-        if array is None:
-            return None
-        else:
-            return array.astype(dtype)
-
-    @staticmethod
-    def _get_freq_for_period(freq: str) -> str:
-        """Convert freq to format compatible with pd.Period.
-
-        For example, ME freq must be converted to M when creating a pd.Period.
-        """
-        offset = pd.tseries.frequencies.to_offset(freq)
-        freq_name = norm_freq_str(offset)
-        if freq_name == "SME":
-            # Replace unsupported frequency "SME" with "2W"
-            return "2W"
-        elif freq_name == "bh":
-            # Replace unsupported frequency "bh" with dummy value "Y"
-            return "Y"
-        else:
-            freq_name_for_period = {"YE": "Y", "QE": "Q", "ME": "M"}.get(freq_name, freq_name)
-            return f"{offset.n}{freq_name_for_period}"
-
-    def __len__(self):
-        return len(self.indptr) - 1  # noqa
-
-    def __iter__(self) -> Iterator[Dict[str, Any]]:
-        for j in range(len(self.indptr) - 1):
-            start_idx = self.indptr[j]
-            end_idx = self.indptr[j + 1]
-            # GluonTS expects item_id to be a string
-            ts = {
-                FieldName.ITEM_ID: str(self.item_ids[j]),
-                FieldName.START: pd.Period(self.start_timestamps.iloc[j], freq=self.freq),
-                FieldName.TARGET: self.target_array[start_idx:end_idx],
-            }
-            if self.feat_static_cat is not None:
-                ts[FieldName.FEAT_STATIC_CAT] = self.feat_static_cat[j]
-            if self.feat_static_real is not None:
-                ts[FieldName.FEAT_STATIC_REAL] = self.feat_static_real[j]
-            if self.past_feat_dynamic_cat is not None:
-                ts[FieldName.PAST_FEAT_DYNAMIC_CAT] = self.past_feat_dynamic_cat[start_idx:end_idx].T
-            if self.past_feat_dynamic_real is not None:
-                ts[FieldName.PAST_FEAT_DYNAMIC_REAL] = self.past_feat_dynamic_real[start_idx:end_idx].T
-
-            # Dynamic features that may extend into the future
-            if self.includes_future:
-                start_idx = start_idx + j * self.prediction_length
-                end_idx = end_idx + (j + 1) * self.prediction_length
-            if self.feat_dynamic_cat is not None:
-                ts[FieldName.FEAT_DYNAMIC_CAT] = self.feat_dynamic_cat[start_idx:end_idx].T
-            if self.feat_dynamic_real is not None:
-                ts[FieldName.FEAT_DYNAMIC_REAL] = self.feat_dynamic_real[start_idx:end_idx].T
-            yield ts
-
-
 class AbstractGluonTSModel(AbstractTimeSeriesModel):
     """Abstract class wrapping GluonTS estimators for use in autogluon.timeseries.
 
     Parameters
     ----------
-    path: str
+    path
         directory to store model artifacts.
-    freq: str
+    freq
         string representation (compatible with GluonTS frequency strings) for the data provided.
         For example, "1D" for daily data, "1H" for hourly data, etc.
-    prediction_length: int
+    prediction_length
         Number of time steps ahead (length of the forecast horizon) the model will be optimized
         to predict. At inference time, this will be the number of time steps the model will
         predict.
-    name: str
+    name
         Name of the model. Also, name of subdirectory inside path where model will be saved.
-    eval_metric: str
+    eval_metric
         objective function the model intends to optimize, will use WQL by default.
-    hyperparameters:
+    hyperparameters
         various hyperparameters that will be used by model (can be search spaces instead of
         fixed values). See *Other Parameters* in each inheriting model's documentation for
         possible values.
@@ -167,12 +72,12 @@ class AbstractGluonTSModel(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: str = None,
-        hyperparameters: 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
     ):
         super().__init__(
@@ -184,9 +89,9 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             hyperparameters=hyperparameters,
             **kwargs,
         )
-        self.gts_predictor: Optional[GluonTSPredictor] = None
-        self._ohe_generator_known: Optional[OneHotEncoder] = None
-        self._ohe_generator_past: Optional[OneHotEncoder] = None
+        self.gts_predictor: GluonTSPredictor | None = None
+        self._ohe_generator_known: OneHotEncoder | None = None
+        self._ohe_generator_past: OneHotEncoder | None = None
         self.callbacks = []
         # Following attributes may be overridden during fit() based on train_data & model parameters
         self.num_feat_static_cat = 0
@@ -195,30 +100,32 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         self.num_feat_dynamic_real = 0
         self.num_past_feat_dynamic_cat = 0
         self.num_past_feat_dynamic_real = 0
-        self.feat_static_cat_cardinality: List[int] = []
-        self.feat_dynamic_cat_cardinality: List[int] = []
-        self.past_feat_dynamic_cat_cardinality: List[int] = []
+        self.feat_static_cat_cardinality: list[int] = []
+        self.feat_dynamic_cat_cardinality: list[int] = []
+        self.past_feat_dynamic_cat_cardinality: list[int] = []
         self.negative_data = True
 
-    def save(self, path: str = None, verbose: bool = True) -> str:
+    def save(self, path: str | None = None, verbose: bool = True) -> str:
         # we flush callbacks instance variable if it has been set. it can keep weak references which breaks training
         self.callbacks = []
         # The GluonTS predictor is serialized using custom logic
         predictor = self.gts_predictor
         self.gts_predictor = None
-        path = Path(super().save(path=path, verbose=verbose))
+        saved_path = Path(super().save(path=path, verbose=verbose))
 
         with disable_root_logger():
             if predictor:
-                Path.mkdir(path / self.gluonts_model_path, exist_ok=True)
-                predictor.serialize(path / self.gluonts_model_path)
+                Path.mkdir(saved_path / self.gluonts_model_path, exist_ok=True)
+                predictor.serialize(saved_path / self.gluonts_model_path)
 
         self.gts_predictor = predictor
 
-        return str(path)
+        return str(saved_path)
 
     @classmethod
-    def load(cls, path: str, reset_paths: bool = True, verbose: bool = True) -> "AbstractGluonTSModel":
+    def load(
+        cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True
+    ) -> "AbstractGluonTSModel":
         from gluonts.torch.model.predictor import PyTorchPredictor
 
         with warning_filter():
@@ -235,31 +142,33 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
     def _get_hpo_backend(self):
         return RAY_BACKEND
 
-    def _deferred_init_params_aux(self, dataset: TimeSeriesDataFrame) -> None:
-        """Update GluonTS specific parameters with information available only at training time."""
-        model_params = self._get_model_params()
+    def _deferred_init_hyperparameters(self, dataset: TimeSeriesDataFrame) -> None:
+        """Update GluonTS specific hyperparameters with information available only at training time."""
+        model_params = self.get_hyperparameters()
         disable_static_features = model_params.get("disable_static_features", False)
         if not disable_static_features:
-            self.num_feat_static_cat = len(self.metadata.static_features_cat)
-            self.num_feat_static_real = len(self.metadata.static_features_real)
+            self.num_feat_static_cat = len(self.covariate_metadata.static_features_cat)
+            self.num_feat_static_real = len(self.covariate_metadata.static_features_real)
             if self.num_feat_static_cat > 0:
-                feat_static_cat = dataset.static_features[self.metadata.static_features_cat]
-                self.feat_static_cat_cardinality = feat_static_cat.nunique().tolist()
+                assert dataset.static_features is not None, (
+                    "Static features must be provided if num_feat_static_cat > 0"
+                )
+                self.feat_static_cat_cardinality = list(self.covariate_metadata.static_cat_cardinality.values())
 
         disable_known_covariates = model_params.get("disable_known_covariates", False)
         if not disable_known_covariates and self.supports_known_covariates:
-            self.num_feat_dynamic_cat = len(self.metadata.known_covariates_cat)
-            self.num_feat_dynamic_real = len(self.metadata.known_covariates_real)
+            self.num_feat_dynamic_cat = len(self.covariate_metadata.known_covariates_cat)
+            self.num_feat_dynamic_real = len(self.covariate_metadata.known_covariates_real)
             if self.num_feat_dynamic_cat > 0:
-                feat_dynamic_cat = dataset[self.metadata.known_covariates_cat]
                 if self.supports_cat_covariates:
-                    self.feat_dynamic_cat_cardinality = feat_dynamic_cat.nunique().tolist()
+                    self.feat_dynamic_cat_cardinality = list(self.covariate_metadata.known_cat_cardinality.values())
                 else:
+                    feat_dynamic_cat = dataset[self.covariate_metadata.known_covariates_cat]
                     # If model doesn't support categorical covariates, convert them to real via one hot encoding
                     self._ohe_generator_known = OneHotEncoder(
                         max_levels=model_params.get("max_cat_cardinality", 100),
                         sparse=False,
-                        dtype="float32",
+                        dtype="float32",  # type: ignore
                     )
                     feat_dynamic_cat_ohe = self._ohe_generator_known.fit_transform(pd.DataFrame(feat_dynamic_cat))
                     self.num_feat_dynamic_cat = 0
@@ -267,18 +176,20 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         disable_past_covariates = model_params.get("disable_past_covariates", False)
         if not disable_past_covariates and self.supports_past_covariates:
-            self.num_past_feat_dynamic_cat = len(self.metadata.past_covariates_cat)
-            self.num_past_feat_dynamic_real = len(self.metadata.past_covariates_real)
+            self.num_past_feat_dynamic_cat = len(self.covariate_metadata.past_covariates_cat)
+            self.num_past_feat_dynamic_real = len(self.covariate_metadata.past_covariates_real)
             if self.num_past_feat_dynamic_cat > 0:
-                past_feat_dynamic_cat = dataset[self.metadata.past_covariates_cat]
                 if self.supports_cat_covariates:
-                    self.past_feat_dynamic_cat_cardinality = past_feat_dynamic_cat.nunique().tolist()
+                    self.past_feat_dynamic_cat_cardinality = list(
+                        self.covariate_metadata.past_cat_cardinality.values()
+                    )
                 else:
+                    past_feat_dynamic_cat = dataset[self.covariate_metadata.past_covariates_cat]
                     # If model doesn't support categorical covariates, convert them to real via one hot encoding
                     self._ohe_generator_past = OneHotEncoder(
                         max_levels=model_params.get("max_cat_cardinality", 100),
                         sparse=False,
-                        dtype="float32",
+                        dtype="float32",  # type: ignore
                     )
                     past_feat_dynamic_cat_ohe = self._ohe_generator_past.fit_transform(
                         pd.DataFrame(past_feat_dynamic_cat)
@@ -288,7 +199,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         self.negative_data = (dataset[self.target] < 0).any()
 
-    def _get_default_params(self):
+    def _get_default_hyperparameters(self):
         """Gets default parameters for GluonTS estimator initialization that are available after
         AbstractTimeSeriesModel initialization (i.e., before deferred initialization). Models may
         override this method to update default parameters.
@@ -306,7 +217,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             "covariate_scaler": "global",
         }
 
-    def _get_model_params(self) -> dict:
+    def get_hyperparameters(self) -> dict:
         """Gets params that are passed to the inner model."""
         # for backward compatibility with the old GluonTS MXNet API
         parameter_name_aliases = {
@@ -314,7 +225,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             "learning_rate": "lr",
         }
 
-        init_args = super()._get_model_params()
+        init_args = super().get_hyperparameters()
         for alias, actual in parameter_name_aliases.items():
             if alias in init_args:
                 if actual in init_args:
@@ -322,12 +233,12 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
                 else:
                     init_args[actual] = init_args.pop(alias)
 
-        return self._get_default_params() | init_args
+        return self._get_default_hyperparameters() | init_args
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
-        """Get GluonTS specific constructor arguments for estimator objects, an alias to `self._get_model_params`
+    def _get_estimator_init_args(self) -> dict[str, Any]:
+        """Get GluonTS specific constructor arguments for estimator objects, an alias to `self.get_hyperparameters`
         for better readability."""
-        return self._get_model_params()
+        return self.get_hyperparameters()
 
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         raise NotImplementedError
@@ -367,25 +278,39 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         return torch.cuda.is_available()
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> Dict[str, Union[int, float]]:
-        minimum_resources = {"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"] = 1
         return minimum_resources
 
+    @overload
+    def _to_gluonts_dataset(self, time_series_df: None, known_covariates=None) -> None: ...
+    @overload
+    def _to_gluonts_dataset(self, time_series_df: TimeSeriesDataFrame, known_covariates=None) -> GluonTSDataset: ...
     def _to_gluonts_dataset(
-        self, time_series_df: Optional[TimeSeriesDataFrame], known_covariates: Optional[TimeSeriesDataFrame] = None
-    ) -> Optional[GluonTSDataset]:
+        self, time_series_df: TimeSeriesDataFrame | None, known_covariates: TimeSeriesDataFrame | None = None
+    ) -> GluonTSDataset | None:
         if time_series_df is not None:
             # TODO: Preprocess real-valued features with StdScaler?
             if self.num_feat_static_cat > 0:
-                feat_static_cat = time_series_df.static_features[self.metadata.static_features_cat].to_numpy()
+                assert time_series_df.static_features is not None, (
+                    "Static features must be provided if num_feat_static_cat > 0"
+                )
+                feat_static_cat = time_series_df.static_features[
+                    self.covariate_metadata.static_features_cat
+                ].to_numpy()
             else:
                 feat_static_cat = None
 
             if self.num_feat_static_real > 0:
-                feat_static_real = time_series_df.static_features[self.metadata.static_features_real].to_numpy()
+                assert time_series_df.static_features is not None, (
+                    "Static features must be provided if num_feat_static_real > 0"
+                )
+                feat_static_real = time_series_df.static_features[
+                    self.covariate_metadata.static_features_real
+                ].to_numpy()
             else:
                 feat_static_real = None
 
@@ -393,31 +318,33 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             # Convert TSDF -> DF to avoid overhead / input validation
             df = pd.DataFrame(time_series_df)
             if known_covariates is not None:
-                known_covariates = pd.DataFrame(known_covariates)
+                known_covariates = pd.DataFrame(known_covariates)  # type: ignore
             if self.num_feat_dynamic_cat > 0:
-                feat_dynamic_cat = df[self.metadata.known_covariates_cat].to_numpy()
+                feat_dynamic_cat = df[self.covariate_metadata.known_covariates_cat].to_numpy()
                 if known_covariates is not None:
                     feat_dynamic_cat = np.concatenate(
-                        [feat_dynamic_cat, known_covariates[self.metadata.known_covariates_cat].to_numpy()]
+                        [feat_dynamic_cat, known_covariates[self.covariate_metadata.known_covariates_cat].to_numpy()]
                     )
                     assert len(feat_dynamic_cat) == expected_known_covariates_len
             else:
                 feat_dynamic_cat = None
 
             if self.num_feat_dynamic_real > 0:
-                feat_dynamic_real = df[self.metadata.known_covariates_real].to_numpy()
+                feat_dynamic_real = df[self.covariate_metadata.known_covariates_real].to_numpy()
                 # Append future values of known covariates
                 if known_covariates is not None:
                     feat_dynamic_real = np.concatenate(
-                        [feat_dynamic_real, known_covariates[self.metadata.known_covariates_real].to_numpy()]
+                        [feat_dynamic_real, known_covariates[self.covariate_metadata.known_covariates_real].to_numpy()]
                     )
                     assert len(feat_dynamic_real) == expected_known_covariates_len
                 # Categorical covariates are one-hot-encoded as real
                 if self._ohe_generator_known is not None:
-                    feat_dynamic_cat_ohe = self._ohe_generator_known.transform(df[self.metadata.known_covariates_cat])
+                    feat_dynamic_cat_ohe: np.ndarray = self._ohe_generator_known.transform(
+                        df[self.covariate_metadata.known_covariates_cat]
+                    )  # type: ignore
                     if known_covariates is not None:
-                        future_dynamic_cat_ohe = self._ohe_generator_known.transform(
-                            known_covariates[self.metadata.known_covariates_cat]
+                        future_dynamic_cat_ohe: np.ndarray = self._ohe_generator_known.transform(  # type: ignore
+                            known_covariates[self.covariate_metadata.known_covariates_cat]
                         )
                         feat_dynamic_cat_ohe = np.concatenate([feat_dynamic_cat_ohe, future_dynamic_cat_ohe])
                         assert len(feat_dynamic_cat_ohe) == expected_known_covariates_len
@@ -426,15 +353,15 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
                 feat_dynamic_real = None
 
             if self.num_past_feat_dynamic_cat > 0:
-                past_feat_dynamic_cat = df[self.metadata.past_covariates_cat].to_numpy()
+                past_feat_dynamic_cat = df[self.covariate_metadata.past_covariates_cat].to_numpy()
             else:
                 past_feat_dynamic_cat = None
 
             if self.num_past_feat_dynamic_real > 0:
-                past_feat_dynamic_real = df[self.metadata.past_covariates_real].to_numpy()
+                past_feat_dynamic_real = df[self.covariate_metadata.past_covariates_real].to_numpy()
                 if self._ohe_generator_past is not None:
-                    past_feat_dynamic_cat_ohe = self._ohe_generator_past.transform(
-                        df[self.metadata.past_covariates_cat]
+                    past_feat_dynamic_cat_ohe: np.ndarray = self._ohe_generator_past.transform(  # type: ignore
+                        df[self.covariate_metadata.past_covariates_cat]
                     )
                     past_feat_dynamic_real = np.concatenate(
                         [past_feat_dynamic_real, past_feat_dynamic_cat_ohe], axis=1
@@ -442,8 +369,9 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             else:
                 past_feat_dynamic_real = None
 
+            assert self.freq is not None
             return SimpleGluonTSDataset(
-                target_df=time_series_df[[self.target]],
+                target_df=time_series_df[[self.target]],  # type: ignore
                 freq=self.freq,
                 target_column=self.target,
                 feat_static_cat=feat_static_cat,
@@ -461,14 +389,16 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: 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:
         # necessary to initialize the loggers
         import lightning.pytorch  # noqa
 
-        verbosity = kwargs.get("verbosity", 2)
         for logger_name in logging.root.manager.loggerDict:
             if "lightning" in logger_name:
                 pl_logger = logging.getLogger(logger_name)
@@ -489,18 +419,18 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             time_limit=time_limit,
             early_stopping_patience=None if val_data is None else init_args["early_stopping_patience"],
         )
-        self._deferred_init_params_aux(train_data)
+        self._deferred_init_hyperparameters(train_data)
 
         estimator = self._get_estimator()
-        with warning_filter(), disable_root_logger(), gluonts.core.settings.let(gluonts.env.env, use_tqdm=False):
+        with warning_filter(), disable_root_logger(), gluonts.core.settings.let(gluonts_env, use_tqdm=False):
             self.gts_predictor = estimator.train(
                 self._to_gluonts_dataset(train_data),
                 validation_data=self._to_gluonts_dataset(val_data),
-                cache_data=True,
+                cache_data=True,  # type: ignore
             )
             # Increase batch size during prediction to speed up inference
             if init_args["predict_batch_size"] is not None:
-                self.gts_predictor.batch_size = init_args["predict_batch_size"]
+                self.gts_predictor.batch_size = init_args["predict_batch_size"]  # type: ignore
 
         lightning_logs_dir = Path(self.path) / "lightning_logs"
         if not keep_lightning_logs and lightning_logs_dir.exists() and lightning_logs_dir.is_dir():
@@ -509,9 +439,9 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
     def _get_callbacks(
         self,
-        time_limit: int,
-        early_stopping_patience: Optional[int] = None,
-    ) -> List[Callable]:
+        time_limit: float | None,
+        early_stopping_patience: int | None = None,
+    ) -> list[Callable]:
         """Retrieve a list of callback objects for the GluonTS trainer"""
         from lightning.pytorch.callbacks import EarlyStopping, Timer
 
@@ -525,14 +455,14 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         if self.gts_predictor is None:
             raise ValueError("Please fit the model before predicting.")
 
-        with warning_filter(), gluonts.core.settings.let(gluonts.env.env, use_tqdm=False):
-            predicted_targets = self._predict_gluonts_forecasts(data, known_covariates=known_covariates, **kwargs)
+        with warning_filter(), gluonts.core.settings.let(gluonts_env, use_tqdm=False):
+            predicted_targets = self._predict_gluonts_forecasts(data, known_covariates=known_covariates)
             df = self._gluonts_forecasts_to_data_frame(
                 predicted_targets,
                 forecast_index=self.get_forecast_horizon_index(data),
@@ -540,16 +470,21 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         return df
 
     def _predict_gluonts_forecasts(
-        self, data: TimeSeriesDataFrame, known_covariates: Optional[TimeSeriesDataFrame] = None, **kwargs
-    ) -> List[Forecast]:
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: TimeSeriesDataFrame | None = None,
+        num_samples: int | None = None,
+    ) -> list[Forecast]:
+        assert self.gts_predictor is not None, "GluonTS models must be fit before predicting."
         gts_data = self._to_gluonts_dataset(data, known_covariates=known_covariates)
+        return list(
+            self.gts_predictor.predict(
+                dataset=gts_data,
+                num_samples=num_samples or self.default_num_samples,
+            )
+        )
 
-        predictor_kwargs = dict(dataset=gts_data)
-        predictor_kwargs["num_samples"] = kwargs.get("num_samples", self.default_num_samples)
-
-        return list(self.gts_predictor.predict(**predictor_kwargs))
-
-    def _stack_quantile_forecasts(self, forecasts: List[QuantileForecast], item_ids: pd.Index) -> pd.DataFrame:
+    def _stack_quantile_forecasts(self, forecasts: list[QuantileForecast], item_ids: pd.Index) -> pd.DataFrame:
         # GluonTS always saves item_id as a string
         item_id_to_forecast = {str(f.item_id): f for f in forecasts}
         result_dfs = []
@@ -562,7 +497,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         columns_order = ["mean"] + [str(q) for q in self.quantile_levels]
         return forecast_df[columns_order]
 
-    def _stack_sample_forecasts(self, forecasts: List[SampleForecast], item_ids: pd.Index) -> pd.DataFrame:
+    def _stack_sample_forecasts(self, forecasts: list[SampleForecast], item_ids: pd.Index) -> pd.DataFrame:
         item_id_to_forecast = {str(f.item_id): f for f in forecasts}
         samples_per_item = []
         for item_id in item_ids:
@@ -574,17 +509,25 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         forecast_array = np.concatenate([mean, quantiles], axis=1)
         return pd.DataFrame(forecast_array, columns=["mean"] + [str(q) for q in self.quantile_levels])
 
-    def _stack_distribution_forecasts(self, forecasts: List[Forecast], item_ids: pd.Index) -> pd.DataFrame:
+    def _stack_distribution_forecasts(
+        self, forecasts: list["DistributionForecast"], item_ids: pd.Index
+    ) -> pd.DataFrame:
         import torch
         from gluonts.torch.distributions import AffineTransformed
         from torch.distributions import Distribution
 
         # Sort forecasts in the same order as in the dataset
         item_id_to_forecast = {str(f.item_id): f for f in forecasts}
-        forecasts = [item_id_to_forecast[str(item_id)] for item_id in item_ids]
+        dist_forecasts = [item_id_to_forecast[str(item_id)] for item_id in item_ids]
+
+        assert all(isinstance(f.distribution, AffineTransformed) for f in dist_forecasts), (
+            "Expected forecast.distribution to be an instance of AffineTransformed"
+        )
 
-        def stack_distributions(distributions: List[Distribution]) -> Distribution:
+        def stack_distributions(distributions: list[Distribution]) -> Distribution:
             """Stack multiple torch.Distribution objects into a single distribution"""
+            last_dist: Distribution = distributions[-1]
+
             params_per_dist = []
             for dist in distributions:
                 params = {name: getattr(dist, name) for name in dist.arg_constraints.keys()}
@@ -593,22 +536,19 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             assert len(set(tuple(p.keys()) for p in params_per_dist)) == 1
 
             stacked_params = {}
-            for key in dist.arg_constraints.keys():
+            for key in last_dist.arg_constraints.keys():
                 stacked_params[key] = torch.cat([p[key] for p in params_per_dist])
-            return dist.__class__(**stacked_params)
-
-        if not isinstance(forecasts[0].distribution, AffineTransformed):
-            raise AssertionError("Expected forecast.distribution to be an instance of AffineTransformed")
+            return last_dist.__class__(**stacked_params)
 
         # We stack all forecast distribution into a single Distribution object.
         # This dramatically speeds up the quantiles calculation.
-        stacked_base_dist = stack_distributions([f.distribution.base_dist for f in forecasts])
+        stacked_base_dist = stack_distributions([f.distribution.base_dist for f in dist_forecasts])  # type: ignore
 
-        stacked_loc = torch.cat([f.distribution.loc for f in forecasts])
+        stacked_loc = torch.cat([f.distribution.loc for f in dist_forecasts])  # type: ignore
         if stacked_loc.shape != stacked_base_dist.batch_shape:
             stacked_loc = stacked_loc.repeat_interleave(self.prediction_length)
 
-        stacked_scale = torch.cat([f.distribution.scale for f in forecasts])
+        stacked_scale = torch.cat([f.distribution.scale for f in dist_forecasts])  # type: ignore
         if stacked_scale.shape != stacked_base_dist.batch_shape:
             stacked_scale = stacked_scale.repeat_interleave(self.prediction_length)
 
@@ -616,24 +556,24 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         mean_prediction = stacked_dist.mean.cpu().detach().numpy()
         quantiles = torch.tensor(self.quantile_levels, device=stacked_dist.mean.device).reshape(-1, 1)
-        quantile_predictions = stacked_dist.icdf(quantiles).cpu().detach().numpy()
+        quantile_predictions = stacked_dist.icdf(quantiles).cpu().detach().numpy()  # type: ignore
         forecast_array = np.vstack([mean_prediction, quantile_predictions]).T
         return pd.DataFrame(forecast_array, columns=["mean"] + [str(q) for q in self.quantile_levels])
 
     def _gluonts_forecasts_to_data_frame(
         self,
-        forecasts: List[Forecast],
+        forecasts: list[Forecast],
         forecast_index: pd.MultiIndex,
     ) -> TimeSeriesDataFrame:
         from gluonts.torch.model.forecast import DistributionForecast
 
-        item_ids = forecast_index.unique(level=ITEMID)
+        item_ids = forecast_index.unique(level=TimeSeriesDataFrame.ITEMID)
         if isinstance(forecasts[0], SampleForecast):
-            forecast_df = self._stack_sample_forecasts(forecasts, item_ids)
+            forecast_df = self._stack_sample_forecasts(cast(list[SampleForecast], forecasts), item_ids)
         elif isinstance(forecasts[0], QuantileForecast):
-            forecast_df = self._stack_quantile_forecasts(forecasts, item_ids)
+            forecast_df = self._stack_quantile_forecasts(cast(list[QuantileForecast], forecasts), item_ids)
         elif isinstance(forecasts[0], DistributionForecast):
-            forecast_df = self._stack_distribution_forecasts(forecasts, item_ids)
+            forecast_df = self._stack_distribution_forecasts(cast(list[DistributionForecast], forecasts), item_ids)
         else:
             raise ValueError(f"Unrecognized forecast type {type(forecasts[0])}")