autogluon.timeseries 1.2.1b20250424__py3-none-any.whl → 1.2.1b20250426__py3-none-any.whl

autogluon/timeseries/metrics/quantile.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Sequence
 
 import numpy as np
 import pandas as pd
@@ -25,6 +25,7 @@ class WQL(TimeSeriesScorer):
     - scale-dependent (time series with large absolute value contribute more to the loss)
     - equivalent to WAPE if ``quantile_levels = [0.5]``
 
+    If `horizon_weight` is provided, both the errors and the target time series in the denominator will be re-weighted.
 
     References
     ----------
@@ -34,16 +35,25 @@ class WQL(TimeSeriesScorer):
     needs_quantile = True
 
     def compute_metric(
-        self, data_future: TimeSeriesDataFrame, predictions: TimeSeriesDataFrame, target: str = "target", **kwargs
+        self,
+        data_future: TimeSeriesDataFrame,
+        predictions: TimeSeriesDataFrame,
+        target: str = "target",
+        **kwargs,
     ) -> float:
         y_true, q_pred, quantile_levels = self._get_quantile_forecast_score_inputs(data_future, predictions, target)
-        values_true = y_true.values[:, None]  # shape [N, 1]
-        values_pred = q_pred.values  # shape [N, len(quantile_levels)]
+        y_true = y_true.to_numpy()[:, None]  # shape [N, 1]
+        q_pred = q_pred.to_numpy()  # shape [N, len(quantile_levels)]
 
-        return 2 * np.mean(
-            np.nansum(np.abs((values_true - values_pred) * ((values_true <= values_pred) - quantile_levels)), axis=0)
-            / np.nansum(np.abs(values_true))
+        errors = (
+            np.abs((q_pred - y_true) * ((y_true <= q_pred) - quantile_levels))
+            .mean(axis=1)
+            .reshape([-1, self.prediction_length])
         )
+        if self.horizon_weight is not None:
+            errors *= self.horizon_weight
+            y_true = y_true.reshape([-1, self.prediction_length]) * self.horizon_weight
+        return 2 * np.nansum(errors) / np.nansum(np.abs(y_true))
 
 
 class SQL(TimeSeriesScorer):
@@ -79,7 +89,15 @@ class SQL(TimeSeriesScorer):
 
     needs_quantile = True
 
-    def __init__(self):
+    def __init__(
+        self,
+        prediction_length: int = 1,
+        seasonal_period: Optional[int] = None,
+        horizon_weight: Optional[Sequence[float]] = None,
+    ):
+        super().__init__(
+            prediction_length=prediction_length, seasonal_period=seasonal_period, horizon_weight=horizon_weight
+        )
         self._past_abs_seasonal_error: Optional[pd.Series] = None
 
     def save_past_metrics(
@@ -93,17 +111,24 @@ class SQL(TimeSeriesScorer):
         self._past_abs_seasonal_error = None
 
     def compute_metric(
-        self, data_future: TimeSeriesDataFrame, predictions: TimeSeriesDataFrame, target: str = "target", **kwargs
+        self,
+        data_future: TimeSeriesDataFrame,
+        predictions: TimeSeriesDataFrame,
+        target: str = "target",
+        **kwargs,
     ) -> float:
         if self._past_abs_seasonal_error is None:
             raise AssertionError("Call `save_past_metrics` before `compute_metric`")
 
         y_true, q_pred, quantile_levels = self._get_quantile_forecast_score_inputs(data_future, predictions, target)
-        q_pred = q_pred.values
-        values_true = y_true.values[:, None]  # shape [N, 1]
-
-        ql = np.abs((q_pred - values_true) * ((values_true <= q_pred) - quantile_levels)).mean(axis=1)
-        num_items = len(self._past_abs_seasonal_error)
-        # Reshape quantile losses values into [num_items, prediction_length] to normalize per item without groupby
-        quantile_losses = ql.reshape([num_items, -1])
-        return 2 * self._safemean(quantile_losses / self._past_abs_seasonal_error.values[:, None])
+        q_pred = q_pred.to_numpy()
+        y_true = y_true.to_numpy()[:, None]  # shape [N, 1]
+
+        errors = (
+            np.abs((q_pred - y_true) * ((y_true <= q_pred) - quantile_levels))
+            .mean(axis=1)
+            .reshape([-1, self.prediction_length])
+        )
+        if self.horizon_weight is not None:
+            errors *= self.horizon_weight
+        return 2 * self._safemean(errors / self._past_abs_seasonal_error.to_numpy()[:, None])

autogluon/timeseries/models/abstract/abstract_timeseries_model.py

@@ -57,9 +57,6 @@ class TimeSeriesModelBase(ModelBase, ABC):
         Metric by which predictions will be ultimately evaluated on future test data. This only impacts
         ``model.score()``, as eval_metric is not used during training. Available metrics can be found in
         ``autogluon.timeseries.metrics``.
-    eval_metric_seasonal_period : int, optional
-        Seasonal period used to compute some evaluation metrics such as mean absolute scaled error (MASE). Defaults to
-        ``None``, in which case the seasonal period is computed based on the data frequency.
     hyperparameters : dict, default = None
         Hyperparameters that will be used by the model (can be search spaces instead of fixed values).
         If None, model defaults are used. This is identical to passing an empty dictionary.
@@ -88,7 +85,6 @@ class TimeSeriesModelBase(ModelBase, ABC):
         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: Union[str, TimeSeriesScorer, None] = None,
-        eval_metric_seasonal_period: Optional[int] = None,
     ):
         self.name = name or re.sub(r"Model$", "", self.__class__.__name__)
 
@@ -103,8 +99,7 @@ class TimeSeriesModelBase(ModelBase, ABC):
 
         self.path = os.path.join(self.path_root, self.name)
 
-        self.eval_metric: TimeSeriesScorer = check_get_evaluation_metric(eval_metric)
-        self.eval_metric_seasonal_period = eval_metric_seasonal_period
+        self.eval_metric = check_get_evaluation_metric(eval_metric, prediction_length=prediction_length)
         self.target: str = target
         self.covariate_metadata = covariate_metadata or CovariateMetadata()
 
@@ -187,7 +182,7 @@ class TimeSeriesModelBase(ModelBase, ABC):
             )
         return hyperparameters, extra_ag_args
 
-    def save(self, path: Optional[str] = None, verbose=True) -> str:
+    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
         if path is None:
             path = self.path
 
@@ -393,8 +388,8 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, ABC):
         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: Union[str, TimeSeriesScorer, None] = None,
-        eval_metric_seasonal_period: Optional[int] = None,
     ):
+        # TODO: make freq a required argument in AbstractTimeSeriesModel
         super().__init__(
             path=path,
             name=name,
@@ -405,12 +400,10 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, ABC):
             target=target,
             quantile_levels=quantile_levels,
             eval_metric=eval_metric,
-            eval_metric_seasonal_period=eval_metric_seasonal_period,
         )
         self.target_scaler: Optional[TargetScaler]
         self.covariate_scaler: Optional[CovariateScaler]
         self.covariate_regressor: Optional[CovariateRegressor]
-        self._initialize_transforms_and_regressor()
 
     def _initialize_transforms_and_regressor(self) -> None:
         self.target_scaler = get_target_scaler(self.get_hyperparameters().get("target_scaler"), target=self.target)
@@ -474,6 +467,7 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, ABC):
             The fitted model object
         """
         start_time = time.monotonic()
+        self._initialize_transforms_and_regressor()
 
         if self.target_scaler is not None:
             train_data = self.target_scaler.fit_transform(train_data)
@@ -699,19 +693,15 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, ABC):
         self,
         data: TimeSeriesDataFrame,
         predictions: TimeSeriesDataFrame,
-        metric: Optional[str] = None,
     ) -> float:
         """Compute the score measuring how well the predictions align with the data."""
-        eval_metric = self.eval_metric if metric is None else check_get_evaluation_metric(metric)
-        return eval_metric.score(
+        return self.eval_metric.score(
             data=data,
             predictions=predictions,
-            prediction_length=self.prediction_length,
             target=self.target,
-            seasonal_period=self.eval_metric_seasonal_period,
         )
 
-    def score(self, data: TimeSeriesDataFrame, metric: Optional[str] = None) -> float:
+    def score(self, data: TimeSeriesDataFrame) -> float:
         """Return the evaluation scores for given metric and dataset. The last
         `self.prediction_length` time steps of each time series in the input data set
         will be held out and used for computing the evaluation score. Time series
@@ -721,9 +711,6 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, ABC):
         ----------
         data: TimeSeriesDataFrame
             Dataset used for scoring.
-        metric: str
-            String identifier of evaluation metric to use, from one of
-            `autogluon.timeseries.utils.metric_utils.AVAILABLE_METRICS`.
 
         Returns
         -------
@@ -735,7 +722,7 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, ABC):
             prediction_length=self.prediction_length, known_covariates_names=self.covariate_metadata.known_covariates
         )
         predictions = self.predict(past_data, known_covariates=known_covariates)
-        return self._score_with_predictions(data=data, predictions=predictions, metric=metric)
+        return self._score_with_predictions(data=data, predictions=predictions)
 
     def score_and_cache_oof(
         self,

autogluon/timeseries/models/autogluon_tabular/mlforecast.py

@@ -2,15 +2,17 @@ import logging
 import math
 import os
 import time
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Collection, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
 from sklearn.base import BaseEstimator
+from typing_extensions import Self
 
 import autogluon.core as ag
 from autogluon.tabular import TabularPredictor
 from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TIMESTAMP, TimeSeriesDataFrame
+from autogluon.timeseries.metrics.abstract import TimeSeriesScorer
 from autogluon.timeseries.metrics.utils import in_sample_squared_seasonal_error
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.models.local import SeasonalNaiveModel
@@ -29,17 +31,21 @@ logger = logging.getLogger(__name__)
 class TabularEstimator(BaseEstimator):
     """Scikit-learn compatible interface for TabularPredictor."""
 
-    def __init__(self, predictor_init_kwargs: Optional[dict] = None, predictor_fit_kwargs: Optional[dict] = None):
+    def __init__(
+        self,
+        predictor_init_kwargs: Optional[Dict[str, Any]] = None,
+        predictor_fit_kwargs: Optional[Dict[str, Any]] = None,
+    ):
         self.predictor_init_kwargs = predictor_init_kwargs if predictor_init_kwargs is not None else {}
         self.predictor_fit_kwargs = predictor_fit_kwargs if predictor_fit_kwargs is not None else {}
 
-    def get_params(self, deep: bool = True) -> dict:
+    def get_params(self, deep: bool = True) -> Dict[str, Any]:
         return {
             "predictor_init_kwargs": self.predictor_init_kwargs,
             "predictor_fit_kwargs": self.predictor_fit_kwargs,
         }
 
-    def fit(self, X: pd.DataFrame, y: pd.Series) -> "TabularEstimator":
+    def fit(self, X: pd.DataFrame, y: pd.Series) -> Self:
         assert isinstance(X, pd.DataFrame) and isinstance(y, pd.Series)
         df = pd.concat([X, y.rename(MLF_TARGET).to_frame()], axis=1)
         self.predictor = TabularPredictor(**self.predictor_init_kwargs)
@@ -49,7 +55,7 @@ class TabularEstimator(BaseEstimator):
 
     def predict(self, X: pd.DataFrame) -> np.ndarray:
         assert isinstance(X, pd.DataFrame)
-        return self.predictor.predict(X).values
+        return self.predictor.predict(X).values  # type: ignore
 
 
 class AbstractMLForecastModel(AbstractTimeSeriesModel):
@@ -62,9 +68,9 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         prediction_length: int = 1,
         path: Optional[str] = None,
         name: Optional[str] = None,
-        eval_metric: str = None,
-        hyperparameters: Dict[str, Any] = None,
-        **kwargs,  # noqa
+        eval_metric: Optional[Union[str, TimeSeriesScorer]] = None,
+        hyperparameters: Optional[Dict[str, Any]] = None,
+        **kwargs,
     ):
         super().__init__(
             path=path,
@@ -80,14 +86,16 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
 
         self._sum_of_differences: int = 0  # number of time steps removed from each series by differencing
         self._max_ts_length: Optional[int] = None
-        self._target_lags: Optional[List[int]] = None
-        self._date_features: Optional[List[str]] = None
-        self._mlf: Optional[MLForecast] = None
+        self._target_lags: np.ndarray
+        self._date_features: List[Callable]
+        self._mlf: MLForecast
         self._scaler: Optional[BaseTargetTransform] = None
-        self._residuals_std_per_item: Optional[pd.Series] = None
+        self._residuals_std_per_item: pd.Series
         self._train_target_median: Optional[float] = None
         self._non_boolean_real_covariates: List[str] = []
 
+    def _initialize_transforms_and_regressor(self):
+        super()._initialize_transforms_and_regressor()
         # Do not create a scaler in the model, scaler will be passed to MLForecast
         self.target_scaler = None
 
@@ -95,20 +103,23 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
     def tabular_predictor_path(self) -> str:
         return os.path.join(self.path, "tabular_predictor")
 
-    def save(self, path: str = None, verbose: bool = True) -> str:
+    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
         assert "mean" in self._mlf.models_, "TabularPredictor must be trained before saving"
-        tabular_predictor = self._mlf.models_["mean"].predictor
-        self._mlf.models_["mean"].predictor = None
+
+        mean_estimator = self._mlf.models_["mean"]
+        assert isinstance(mean_estimator, TabularEstimator)
+
+        tabular_predictor = mean_estimator.predictor
+        mean_estimator.predictor = None  # type: ignore
         save_path = super().save(path=path, verbose=verbose)
-        self._mlf.models_["mean"].predictor = tabular_predictor
+        mean_estimator.predictor = tabular_predictor
         return save_path
 
     @classmethod
-    def load(
-        cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True
-    ) -> "AbstractTimeSeriesModel":
+    def load(cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True) -> Self:
         model = super().load(path=path, reset_paths=reset_paths, load_oof=load_oof, verbose=verbose)
         assert "mean" in model._mlf.models_, "Loaded model doesn't have a trained TabularPredictor"
+        assert isinstance(model._mlf.models_["mean"], TabularEstimator)
         model._mlf.models_["mean"].predictor = TabularPredictor.load(model.tabular_predictor_path)
         return model
 
@@ -131,24 +142,27 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
                 data[self.target] = data[self.target].fillna(value=self._train_target_median)
         return data, known_covariates
 
-    def _get_extra_tabular_init_kwargs(self) -> dict:
+    def _get_extra_tabular_init_kwargs(self) -> Dict[str, Any]:
         raise NotImplementedError
 
-    def get_hyperparameters(self) -> dict:
-        model_params = super().get_hyperparameters().copy()
-        model_params.setdefault("max_num_items", 20_000)
-        model_params.setdefault("max_num_samples", 1_000_000)
-        model_params.setdefault("tabular_hyperparameters", {"GBM": {}})
-        model_params.setdefault("tabular_fit_kwargs", {})
-        return model_params
+    def _get_default_hyperparameters(self) -> Dict[str, Any]:
+        return {
+            "max_num_items": 20_000,
+            "max_num_samples": 1_000_000,
+            "tabular_hyperparameters": {"GBM": {}},
+            "tabular_fit_kwargs": {},
+        }
 
-    def _get_mlforecast_init_args(self, train_data: TimeSeriesDataFrame, model_params: dict) -> dict:
+    def _get_mlforecast_init_args(
+        self, train_data: TimeSeriesDataFrame, model_params: Dict[str, Any]
+    ) -> Dict[str, Any]:
         from mlforecast.target_transforms import Differences
 
         from .transforms import MLForecastScaler
 
         lags = model_params.get("lags")
         if lags is None:
+            assert self.freq is not None
             lags = get_lags_for_frequency(self.freq)
         self._target_lags = np.array(sorted(set(lags)), dtype=np.int64)
 
@@ -159,6 +173,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
 
         target_transforms = []
         differences = model_params.get("differences")
+        assert isinstance(differences, Collection)
 
         ts_lengths = train_data.num_timesteps_per_item()
         required_ts_length = sum(differences) + 1
@@ -196,7 +211,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         return df
 
     @staticmethod
-    def _shorten_all_series(mlforecast_df: pd.DataFrame, max_length: int):
+    def _shorten_all_series(mlforecast_df: pd.DataFrame, max_length: int) -> pd.DataFrame:
         logger.debug(f"Shortening all series to at most {max_length}")
         return mlforecast_df.groupby(MLF_ITEMID, as_index=False, sort=False).tail(max_length)
 
@@ -231,7 +246,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         # Unless we set static_features=[], MLForecast interprets all known covariates as static features
         df = self._mlf.preprocess(mlforecast_df, dropna=False, static_features=[])
         # df.query results in 2x memory saving compared to df.dropna(subset="y")
-        df = df.query("y.notnull()")
+        df = df.query("y.notnull()")  # type: ignore
 
         df = self._mask_df(df)
 
@@ -250,12 +265,12 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         val_df = grouped_df.tail(val_rows_per_item)
         logger.debug(f"train_df shape: {train_df.shape}, val_df shape: {val_df.shape}")
 
-        return train_df.drop(columns=[MLF_TIMESTAMP]), val_df.drop(columns=[MLF_TIMESTAMP])
+        return train_df.drop(columns=[MLF_TIMESTAMP]), val_df.drop(columns=[MLF_TIMESTAMP])  # type: ignore
 
     def _to_mlforecast_df(
         self,
         data: TimeSeriesDataFrame,
-        static_features: pd.DataFrame,
+        static_features: Optional[pd.DataFrame],
         include_target: bool = True,
     ) -> pd.DataFrame:
         """Convert TimeSeriesDataFrame to a format expected by MLForecast methods `predict` and `preprocess`.
@@ -288,7 +303,9 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         self,
         train_data: TimeSeriesDataFrame,
         val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[int] = None,
+        time_limit: Optional[float] = None,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
         verbosity: int = 2,
         **kwargs,
     ) -> None:
@@ -304,6 +321,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         model_params = self.get_hyperparameters()
 
         mlforecast_init_args = self._get_mlforecast_init_args(train_data, model_params)
+        assert self.freq is not None
         self._mlf = MLForecast(models={}, freq=self.freq, **mlforecast_init_args)
 
         # We generate train/val splits from train_data and ignore val_data to avoid overfitting
@@ -327,10 +345,10 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
                 **model_params["tabular_fit_kwargs"],
             },
         )
-        self._mlf.models = {"mean": estimator}
+        self._mlf.models = {"mean": estimator}  # type: ignore
 
         with warning_filter():
-            self._mlf.fit_models(X=train_df.drop(columns=[MLF_TARGET, MLF_ITEMID]), y=train_df[MLF_TARGET])
+            self._mlf.fit_models(X=train_df.drop(columns=[MLF_TARGET, MLF_ITEMID]), y=train_df[MLF_TARGET])  # type: ignore
 
         self._save_residuals_std(val_df)
 
@@ -340,14 +358,19 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         Saves per-item residuals to `self.residuals_std_per_item`.
         """
         residuals_df = val_df[[MLF_ITEMID, MLF_TARGET]]
-        residuals_df = residuals_df.assign(y_pred=self._mlf.models_["mean"].predict(val_df))
+        mean_estimator = self._mlf.models_["mean"]
+        assert isinstance(mean_estimator, TabularEstimator)
+
+        residuals_df = residuals_df.assign(y_pred=mean_estimator.predict(val_df))
         if self._scaler is not None:
             # Scaler expects to find column MLF_TIMESTAMP even though it's not used - fill with dummy
-            residuals_df = residuals_df.assign(**{MLF_TIMESTAMP: 1})
+            residuals_df = residuals_df.assign(**{MLF_TIMESTAMP: np.datetime64("2010-01-01")})
             residuals_df = self._scaler.inverse_transform(residuals_df)
+
+        assert isinstance(residuals_df, pd.DataFrame)
         residuals = residuals_df[MLF_TARGET] - residuals_df["y_pred"]
         self._residuals_std_per_item = (
-            residuals.pow(2.0).groupby(val_df[MLF_ITEMID].values, sort=False).mean().pow(0.5)
+            residuals.pow(2.0).groupby(val_df[MLF_ITEMID].values, sort=False).mean().pow(0.5)  # type: ignore
         )
 
     def _remove_short_ts_and_generate_fallback_forecast(
@@ -395,7 +418,9 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             forecast_for_short_series = None
         return data_long, known_covariates_long, forecast_for_short_series
 
-    def _add_gaussian_quantiles(self, predictions: pd.DataFrame, repeated_item_ids: pd.Series, past_target: pd.Series):
+    def _add_gaussian_quantiles(
+        self, predictions: pd.DataFrame, repeated_item_ids: pd.Series, past_target: pd.Series
+    ) -> pd.DataFrame:
         """
         Add quantile levels assuming that residuals follow normal distribution
         """
@@ -410,9 +435,9 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         # Use in-sample seasonal error in for items not seen during fit
         items_not_seen_during_fit = residuals_std_per_timestep.index[residuals_std_per_timestep.isna()].unique()
         if len(items_not_seen_during_fit) > 0:
-            scale_for_new_items: pd.Series = np.sqrt(
-                in_sample_squared_seasonal_error(y_past=past_target.loc[items_not_seen_during_fit])
-            )
+            scale_for_new_items: pd.Series = in_sample_squared_seasonal_error(
+                y_past=past_target.loc[items_not_seen_during_fit]
+            ).pow(0.5)
             residuals_std_per_timestep = residuals_std_per_timestep.fillna(scale_for_new_items)
 
         std_per_timestep = residuals_std_per_timestep * normal_scale_per_timestep
@@ -420,7 +445,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             predictions[str(q)] = predictions["mean"] + norm.ppf(q) * std_per_timestep.to_numpy()
         return predictions
 
-    def _more_tags(self) -> dict:
+    def _more_tags(self) -> Dict[str, Any]:
         return {"allow_nan": True, "can_refit_full": True}
 
 
@@ -473,7 +498,7 @@ class DirectTabularModel(AbstractMLForecastModel):
     def is_quantile_model(self) -> bool:
         return self.eval_metric.needs_quantile
 
-    def get_hyperparameters(self) -> dict:
+    def get_hyperparameters(self) -> Dict[str, Any]:
         model_params = super().get_hyperparameters()
         model_params.setdefault("target_scaler", "mean_abs")
         if "differences" not in model_params or model_params["differences"] is None:
@@ -512,6 +537,7 @@ class DirectTabularModel(AbstractMLForecastModel):
         )
         if len(data) == 0:
             # All time series are too short for chosen differences
+            assert forecast_for_short_series is not None
             return forecast_for_short_series
 
         if known_covariates is not None:
@@ -522,15 +548,19 @@ class DirectTabularModel(AbstractMLForecastModel):
         # MLForecast raises exception of target contains NaN. We use inf as placeholder, replace them by NaN afterwards
         data_future[self.target] = float("inf")
         data_extended = pd.concat([data, data_future])
-        mlforecast_df = self._to_mlforecast_df(data_extended, data.static_features)
+        mlforecast_df = self._to_mlforecast_df(data_extended, data.static_features)  # type: ignore
         if self._max_ts_length is not None:
             # We appended `prediction_length` time steps to each series, so increase length
             mlforecast_df = self._shorten_all_series(mlforecast_df, self._max_ts_length + self.prediction_length)
         df = self._mlf.preprocess(mlforecast_df, dropna=False, static_features=[])
+        assert isinstance(df, pd.DataFrame)
+
         df = df.groupby(MLF_ITEMID, sort=False).tail(self.prediction_length)
         df = df.replace(float("inf"), float("nan"))
 
-        raw_predictions = self._mlf.models_["mean"].predict(df)
+        mean_estimator = self._mlf.models_["mean"]
+        assert isinstance(mean_estimator, TabularEstimator)
+        raw_predictions = mean_estimator.predict(df)
         predictions = self._postprocess_predictions(raw_predictions, repeated_item_ids=df[MLF_ITEMID])
         # Paste columns one by one to preserve dtypes
         predictions[MLF_ITEMID] = df[MLF_ITEMID].values
@@ -542,6 +572,7 @@ class DirectTabularModel(AbstractMLForecastModel):
             if self._max_ts_length is not None:
                 mlforecast_df_past = self._shorten_all_series(mlforecast_df_past, self._max_ts_length)
             self._mlf.preprocess(mlforecast_df_past, static_features=[], dropna=False)
+            assert self._mlf.ts.target_transforms is not None
             for tfm in self._mlf.ts.target_transforms[::-1]:
                 predictions = apply_inverse_transform(predictions, transform=tfm)
 
@@ -549,25 +580,30 @@ class DirectTabularModel(AbstractMLForecastModel):
             predictions = self._add_gaussian_quantiles(
                 predictions, repeated_item_ids=predictions[MLF_ITEMID], past_target=data[self.target]
             )
-        predictions = TimeSeriesDataFrame(predictions.rename(columns={MLF_ITEMID: ITEMID, MLF_TIMESTAMP: TIMESTAMP}))
+        predictions_tsdf: TimeSeriesDataFrame = TimeSeriesDataFrame(
+            predictions.rename(columns={MLF_ITEMID: ITEMID, MLF_TIMESTAMP: TIMESTAMP})
+        )
 
         if forecast_for_short_series is not None:
-            predictions = pd.concat([predictions, forecast_for_short_series])
-            predictions = predictions.reindex(original_item_id_order, level=ITEMID)
-        return predictions
+            predictions_tsdf = pd.concat([predictions_tsdf, forecast_for_short_series])  # type: ignore
+            predictions_tsdf = predictions_tsdf.reindex(original_item_id_order, level=ITEMID)
 
-    def _postprocess_predictions(self, predictions: np.ndarray, repeated_item_ids: pd.Series) -> pd.DataFrame:
+        return predictions_tsdf
+
+    def _postprocess_predictions(
+        self, predictions: Union[np.ndarray, pd.Series], repeated_item_ids: pd.Series
+    ) -> pd.DataFrame:
         if self.is_quantile_model:
-            predictions = pd.DataFrame(predictions, columns=[str(q) for q in self.quantile_levels])
-            predictions.values.sort(axis=1)
-            predictions["mean"] = predictions["0.5"]
+            predictions_df = pd.DataFrame(predictions, columns=[str(q) for q in self.quantile_levels])
+            predictions_df.values.sort(axis=1)
+            predictions_df["mean"] = predictions_df["0.5"]
         else:
-            predictions = pd.DataFrame(predictions, columns=["mean"])
+            predictions_df = pd.DataFrame(predictions, columns=["mean"])
 
-        column_order = ["mean"] + [col for col in predictions.columns if col != "mean"]
-        return predictions[column_order]
+        column_order = ["mean"] + [col for col in predictions_df.columns if col != "mean"]
+        return predictions_df[column_order]
 
-    def _get_extra_tabular_init_kwargs(self) -> dict:
+    def _get_extra_tabular_init_kwargs(self) -> Dict[str, Any]:
         if self.is_quantile_model:
             return {
                 "problem_type": ag.constants.QUANTILE,
@@ -622,7 +658,7 @@ class RecursiveTabularModel(AbstractMLForecastModel):
         end of each time series).
     """
 
-    def get_hyperparameters(self) -> dict:
+    def get_hyperparameters(self) -> Dict[str, Any]:
         model_params = super().get_hyperparameters()
         model_params.setdefault("target_scaler", "standard")
         if "differences" not in model_params or model_params["differences"] is None:
@@ -641,6 +677,7 @@ class RecursiveTabularModel(AbstractMLForecastModel):
         )
         if len(data) == 0:
             # All time series are too short for chosen differences
+            assert forecast_for_short_series is not None
             return forecast_for_short_series
 
         new_df = self._to_mlforecast_df(data, data.static_features)
@@ -648,7 +685,9 @@ class RecursiveTabularModel(AbstractMLForecastModel):
             new_df = self._shorten_all_series(new_df, self._max_ts_length)
         if known_covariates is None:
             future_index = self.get_forecast_horizon_index(data)
-            known_covariates = pd.DataFrame(columns=[self.target], index=future_index, dtype="float32")
+            known_covariates = TimeSeriesDataFrame(
+                pd.DataFrame(columns=[self.target], index=future_index, dtype="float32")
+            )
         X_df = self._to_mlforecast_df(known_covariates, data.static_features, include_target=False)
         # If both covariates & static features are missing, set X_df = None to avoid exception from MLForecast
         if len(X_df.columns.difference([MLF_ITEMID, MLF_TIMESTAMP])) == 0:
@@ -659,18 +698,19 @@ class RecursiveTabularModel(AbstractMLForecastModel):
                 new_df=new_df,
                 X_df=X_df,
             )
-        predictions = raw_predictions.rename(columns={MLF_ITEMID: ITEMID, MLF_TIMESTAMP: TIMESTAMP})
-        predictions = TimeSeriesDataFrame(
+        assert isinstance(raw_predictions, pd.DataFrame)
+        raw_predictions = raw_predictions.rename(columns={MLF_ITEMID: ITEMID, MLF_TIMESTAMP: TIMESTAMP})
+
+        predictions: TimeSeriesDataFrame = TimeSeriesDataFrame(
             self._add_gaussian_quantiles(
-                predictions, repeated_item_ids=predictions[ITEMID], past_target=data[self.target]
+                raw_predictions, repeated_item_ids=raw_predictions[ITEMID], past_target=data[self.target]
             )
         )
-
         if forecast_for_short_series is not None:
-            predictions = pd.concat([predictions, forecast_for_short_series])
+            predictions = pd.concat([predictions, forecast_for_short_series])  # type: ignore
         return predictions.reindex(original_item_id_order, level=ITEMID)
 
-    def _get_extra_tabular_init_kwargs(self) -> dict:
+    def _get_extra_tabular_init_kwargs(self) -> Dict[str, Any]:
         return {
             "problem_type": ag.constants.REGRESSION,
             "eval_metric": self.eval_metric.equivalent_tabular_regression_metric or "mean_absolute_error",