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

autogluon/timeseries/models/ensemble/array_based/regressor/tabular.py

@@ -0,0 +1,141 @@
+import logging
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+from typing_extensions import Self
+
+from autogluon.tabular import TabularPredictor
+
+from .abstract import EnsembleRegressor
+
+logger = logging.getLogger(__name__)
+
+
+class TabularEnsembleRegressor(EnsembleRegressor):
+    """TabularPredictor ensemble regressor using AutoGluon-Tabular as a single
+    quantile regressor for the target.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        quantile_levels: list[float],
+        tabular_hyperparameters: Optional[dict] = None,
+    ):
+        super().__init__()
+        self.path = path
+        self.quantile_levels = quantile_levels
+        self.tabular_hyperparameters = tabular_hyperparameters or {}
+        self.predictor: Optional[TabularPredictor] = None
+
+    def set_path(self, path: str) -> None:
+        self.path = path
+
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        time_limit: Optional[int] = None,
+        **kwargs,
+    ) -> Self:
+        self.predictor = TabularPredictor(
+            path=self.path,
+            label="target",
+            problem_type="quantile",
+            quantile_levels=self.quantile_levels,
+            verbosity=1,
+        )
+
+        # get features
+        df = self._get_feature_df(base_model_mean_predictions, base_model_quantile_predictions)
+
+        # get labels
+        num_windows, num_items, prediction_length = base_model_mean_predictions.shape[:3]
+        label_series = labels.reshape(num_windows * num_items * prediction_length)
+        df["target"] = label_series
+
+        self.predictor.fit(
+            df,
+            hyperparameters=self.tabular_hyperparameters,
+            time_limit=time_limit,  # type: ignore
+        )
+
+        return self
+
+    def predict(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if self.predictor is None:
+            try:
+                self.predictor = TabularPredictor.load(self.path)
+            except FileNotFoundError:
+                raise ValueError("Model must be fitted before prediction")
+
+        num_windows, num_items, prediction_length = base_model_mean_predictions.shape[:3]
+        assert num_windows == 1, "Prediction expects a single window to be provided"
+
+        df = self._get_feature_df(base_model_mean_predictions, base_model_quantile_predictions)
+
+        pred = self.predictor.predict(df, as_pandas=False)
+
+        # Reshape back to (num_windows, num_items, prediction_length, num_quantiles)
+        pred = pred.reshape(num_windows, num_items, prediction_length, len(self.quantile_levels))
+
+        # Use median quantile as mean prediction
+        median_idx = self._get_median_quantile_index()
+        mean_pred = pred[:, :, :, median_idx : median_idx + 1]
+        quantile_pred = pred
+
+        return mean_pred, quantile_pred
+
+    def _get_feature_df(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> pd.DataFrame:
+        num_windows, num_items, prediction_length, _, num_models = base_model_mean_predictions.shape
+        num_tabular_items = num_windows * num_items * prediction_length
+
+        X = np.hstack(
+            [
+                base_model_mean_predictions.reshape(num_tabular_items, -1),
+                base_model_quantile_predictions.reshape(num_tabular_items, -1),
+            ]
+        )
+
+        df = pd.DataFrame(X, columns=self._get_feature_names(num_models))
+        return df
+
+    def _get_feature_names(self, num_models: int) -> list[str]:
+        feature_names = []
+        for mi in range(num_models):
+            feature_names.append(f"model_{mi}_mean")
+        for quantile in self.quantile_levels:
+            for mi in range(num_models):
+                feature_names.append(f"model_{mi}_q{quantile}")
+
+        return feature_names
+
+    def _get_median_quantile_index(self):
+        """Get quantile index closest to 0.5"""
+        quantile_array = np.array(self.quantile_levels)
+        median_idx = int(np.argmin(np.abs(quantile_array - 0.5)))
+        selected_quantile = quantile_array[median_idx]
+
+        if selected_quantile != 0.5:
+            logger.warning(
+                f"Selected quantile {selected_quantile} is not exactly 0.5. "
+                f"Using closest available quantile for median prediction."
+            )
+
+        return median_idx
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # Remove the predictor to avoid pickling heavy TabularPredictor objects
+        state["predictor"] = None
+        return state

autogluon/timeseries/models/ensemble/weighted/__init__.py

@@ -0,0 +1,8 @@
+from .basic import PerformanceWeightedEnsemble, SimpleAverageEnsemble
+from .greedy import GreedyEnsemble
+
+__all__ = [
+    "SimpleAverageEnsemble",
+    "PerformanceWeightedEnsemble",
+    "GreedyEnsemble",
+]

autogluon/timeseries/models/ensemble/weighted/abstract.py

@@ -0,0 +1,41 @@
+import functools
+from abc import ABC
+from typing import Optional
+
+import numpy as np
+
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+
+from ..abstract import AbstractTimeSeriesEnsembleModel
+
+
+class AbstractWeightedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
+    """Abstract class for weighted ensembles which assign one (global) weight per model."""
+
+    def __init__(self, name: Optional[str] = None, **kwargs):
+        super().__init__(name=name, **kwargs)
+        self.model_to_weight: dict[str, float] = {}
+
+    @property
+    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:
+        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)
+
+    def get_info(self) -> dict:
+        info = super().get_info()
+        info["model_weights"] = self.model_to_weight.copy()
+        return info
+
+    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)
+            updated_weights[model_full_name] = weight
+        self.model_to_weight = updated_weights

autogluon/timeseries/models/ensemble/{basic.py → weighted/basic.py}

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Any, Optional
 
 import numpy as np
 
@@ -10,16 +10,11 @@ from .abstract import AbstractWeightedTimeSeriesEnsembleModel
 class SimpleAverageEnsemble(AbstractWeightedTimeSeriesEnsembleModel):
     """Constructs a weighted ensemble using a simple average of the constituent models' predictions."""
 
-    def __init__(self, name: Optional[str] = None, **kwargs):
-        if name is None:
-            name = "SimpleAverageEnsemble"
-        super().__init__(name=name, **kwargs)
-
     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 = {}
@@ -47,19 +42,14 @@ class PerformanceWeightedEnsemble(AbstractWeightedTimeSeriesEnsembleModel):
         36.1 (2020): 93-97.
     """
 
-    def __init__(self, name: Optional[str] = None, **kwargs):
-        if name is None:
-            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 → weighted/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(

autogluon/timeseries/models/gluonts/abstract.py

@@ -3,7 +3,7 @@ import os
 import shutil
 from datetime import timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Optional, Type, Union, cast, overload
 
 import gluonts
 import gluonts.core.settings
@@ -21,7 +21,7 @@ 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, TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.utils.warning_filters import disable_root_logger, warning_filter
 
@@ -42,20 +42,20 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
     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.
@@ -77,7 +77,7 @@ class AbstractGluonTSModel(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
     ):
         super().__init__(
@@ -100,9 +100,9 @@ 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: Optional[str] = None, verbose: bool = True) -> str:
@@ -234,7 +234,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
         return self._get_default_hyperparameters() | init_args
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    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_hyperparameters()
@@ -277,8 +277,8 @@ 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: 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"] = 1
@@ -440,7 +440,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         self,
         time_limit: Optional[float],
         early_stopping_patience: Optional[int] = None,
-    ) -> List[Callable]:
+    ) -> list[Callable]:
         """Retrieve a list of callback objects for the GluonTS trainer"""
         from lightning.pytorch.callbacks import EarlyStopping, Timer
 
@@ -473,7 +473,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         data: TimeSeriesDataFrame,
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         num_samples: Optional[int] = None,
-    ) -> List[Forecast]:
+    ) -> 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(
@@ -483,7 +483,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             )
         )
 
-    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 = []
@@ -496,7 +496,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:
@@ -509,7 +509,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
         return pd.DataFrame(forecast_array, columns=["mean"] + [str(q) for q in self.quantile_levels])
 
     def _stack_distribution_forecasts(
-        self, forecasts: List["DistributionForecast"], item_ids: pd.Index
+        self, forecasts: list["DistributionForecast"], item_ids: pd.Index
     ) -> pd.DataFrame:
         import torch
         from gluonts.torch.distributions import AffineTransformed
@@ -523,7 +523,7 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
             "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]
 
@@ -561,18 +561,18 @@ class AbstractGluonTSModel(AbstractTimeSeriesModel):
 
     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(cast(List[SampleForecast], 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(cast(List[QuantileForecast], 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(cast(List[DistributionForecast], 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])}")
 

autogluon/timeseries/models/gluonts/dataset.py

@@ -1,11 +1,11 @@
-from typing import Any, Dict, Iterator, Optional, Type
+from typing import Any, Iterator, Optional, Type
 
 import numpy as np
 import pandas as pd
 from gluonts.dataset.common import Dataset as GluonTSDataset
 from gluonts.dataset.field_names import FieldName
 
-from autogluon.timeseries.dataset.ts_dataframe import TIMESTAMP, TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.utils.datetime import norm_freq_str
 
 
@@ -44,7 +44,7 @@ class SimpleGluonTSDataset(GluonTSDataset):
         # Replace inefficient groupby ITEMID with indptr that stores start:end of each time series
         self.item_ids = target_df.item_ids
         self.indptr = target_df.get_indptr()
-        self.start_timestamps = target_df.index[self.indptr[:-1]].to_frame(index=False)[TIMESTAMP]
+        self.start_timestamps = target_df.index[self.indptr[:-1]].to_frame(index=False)[TimeSeriesDataFrame.TIMESTAMP]
         assert len(self.item_ids) == len(self.start_timestamps)
 
     @staticmethod
@@ -76,7 +76,7 @@ class SimpleGluonTSDataset(GluonTSDataset):
     def __len__(self):
         return len(self.indptr) - 1  # noqa
 
-    def __iter__(self) -> Iterator[Dict[str, Any]]:
+    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]

autogluon/timeseries/models/gluonts/models.py

@@ -3,7 +3,7 @@ Module including wrappers for PyTorch implementations of models in GluonTS
 """
 
 import logging
-from typing import Any, Dict, Type
+from typing import Any, Type
 
 from gluonts.model.estimator import Estimator as GluonTSEstimator
 
@@ -60,7 +60,7 @@ class DeepARModel(AbstractGluonTSModel):
         Distribution output object that defines how the model output is converted to a forecast, and how the loss is computed.
     scaling: bool, default = True
         If True, mean absolute scaling will be applied to each *context window* during training & prediction.
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -81,6 +81,8 @@ class DeepARModel(AbstractGluonTSModel):
 
     # TODO: Replace "scaling: bool" with "window_scaler": {"mean_abs", None} for consistency?
 
+    ag_priority = 40
+
     _supports_known_covariates = True
     _supports_static_features = True
 
@@ -89,7 +91,7 @@ class DeepARModel(AbstractGluonTSModel):
 
         return DeepAREstimator
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real
@@ -111,7 +113,7 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
     ----------------
     context_length : int, default = max(10, 2 * prediction_length)
         Number of time units that condition the predictions
-    hidden_dimensions: List[int], default = [20, 20]
+    hidden_dimensions: list[int], default = [20, 20]
         Size of hidden layers in the feedforward network
     distr_output : gluonts.torch.distributions.Output, default = StudentTOutput()
         Distribution output object that defines how the model output is converted to a forecast, and how the loss is computed.
@@ -119,7 +121,7 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
         Whether to use batch normalization
     mean_scaling : bool, default = True
         If True, mean absolute scaling will be applied to each *context window* during training & prediction.
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -138,6 +140,8 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 10
+
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.simple_feedforward import SimpleFeedForwardEstimator
 
@@ -199,6 +203,9 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 45
+    ag_model_aliases = ["TFT"]
+
     _supports_known_covariates = True
     _supports_past_covariates = True
     _supports_cat_covariates = True
@@ -214,7 +221,7 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
             "context_length": min(512, max(64, 2 * self.prediction_length)),
         }
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         if self.num_feat_dynamic_real > 0:
             init_kwargs["dynamic_dims"] = [self.num_feat_dynamic_real]
@@ -261,7 +268,7 @@ class DLinearModel(AbstractGluonTSModel):
         Scaling applied to each *context window* during training & prediction.
         One of ``"mean"`` (mean absolute scaling), ``"std"`` (standardization), ``None`` (no scaling).
 
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -282,6 +289,8 @@ class DLinearModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 10
+
     def _get_default_hyperparameters(self):
         return super()._get_default_hyperparameters() | {
             "context_length": 96,
@@ -325,7 +334,7 @@ class PatchTSTModel(AbstractGluonTSModel):
         Scaling applied to each *context window* during training & prediction.
         One of ``"mean"`` (mean absolute scaling), ``"std"`` (standardization), ``None`` (no scaling).
 
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -340,6 +349,8 @@ class PatchTSTModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 30
+
     _supports_known_covariates = True
 
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
@@ -350,7 +361,7 @@ class PatchTSTModel(AbstractGluonTSModel):
     def _get_default_hyperparameters(self):
         return super()._get_default_hyperparameters() | {"context_length": 96, "patch_len": 16}
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_dynamic_real"] = self.num_feat_dynamic_real
         return init_kwargs
@@ -416,6 +427,8 @@ class WaveNetModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 25
+
     _supports_known_covariates = True
     _supports_static_features = True
     default_num_samples: int = 100
@@ -425,7 +438,7 @@ class WaveNetModel(AbstractGluonTSModel):
 
         return WaveNetEstimator
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real
@@ -489,7 +502,7 @@ class TiDEModel(AbstractGluonTSModel):
         Scaling applied to each *context window* during training & prediction.
         One of ``"mean"`` (mean absolute scaling), ``"std"`` (standardization), ``None`` (no scaling).
 
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 256
@@ -508,6 +521,8 @@ class TiDEModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 30
+
     _supports_known_covariates = True
     _supports_static_features = True
 
@@ -532,7 +547,7 @@ class TiDEModel(AbstractGluonTSModel):
             "batch_size": 256,
         }
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real