[ENH] Test framework for ptf-v2

pytorch_forecasting/models/base/_base_object.py

@@ -17,10 +17,6 @@ class _BasePtForecaster(_BaseObject):
     This class points to model objects and contains metadata as tags.
     """
 
-    _tags = {
-        "object_type": "forecaster_pytorch",
-    }
-
     @classmethod
     def get_model_cls(cls):
         """Get model class."""
@@ -112,3 +108,19 @@ def create_test_instances_and_names(cls, parameter_set="default"):
             names = [cls.__name__]
 
         return objs, names
+
+
+class _BasePtForecasterV1(_BasePtForecaster):
+    """Base class for PyTorch Forecasting v1 forecasters."""
+
+    _tags = {
+        "object_type": "forecaster_pytorch_v1",
+    }
+
+
+class _BasePtForecasterV2(_BasePtForecaster):
+    """Base class for PyTorch Forecasting v2 forecasters."""
+
+    _tags = {
+        "object_type": "forecaster_pytorch_v2",
+    }

pytorch_forecasting/models/deepar/_deepar_metadata.py

@@ -1,9 +1,9 @@
 """DeepAR metadata container."""
 
-from pytorch_forecasting.models.base._base_object import _BasePtForecaster
+from pytorch_forecasting.models.base._base_object import _BasePtForecasterV1
 
 
-class DeepARMetadata(_BasePtForecaster):
+class DeepARMetadata(_BasePtForecasterV1):
     """DeepAR metadata container."""
 
     _tags = {

pytorch_forecasting/models/mlp/_decodermlp_metadata.py

@@ -1,9 +1,9 @@
 """DecoderMLP metadata container."""
 
-from pytorch_forecasting.models.base._base_object import _BasePtForecaster
+from pytorch_forecasting.models.base._base_object import _BasePtForecasterV1
 
 
-class DecoderMLPMetadata(_BasePtForecaster):
+class DecoderMLPMetadata(_BasePtForecasterV1):
     """DecoderMLP metadata container."""
 
     _tags = {

pytorch_forecasting/models/nbeats/_nbeats_metadata.py

@@ -1,9 +1,9 @@
 """NBeats metadata container."""
 
-from pytorch_forecasting.models.base._base_object import _BasePtForecaster
+from pytorch_forecasting.models.base._base_object import _BasePtForecasterV1
 
 
-class NBeatsMetadata(_BasePtForecaster):
+class NBeatsMetadata(_BasePtForecasterV1):
     """NBeats metadata container."""
 
     _tags = {

pytorch_forecasting/models/temporal_fusion_transformer/tft_v2_metadata.py

@@ -0,0 +1,137 @@
+"""TFT metadata container."""
+
+from pytorch_forecasting.models.base._base_object import _BasePtForecasterV2
+
+
+class TFTMetadata(_BasePtForecasterV2):
+    """TFT metadata container."""
+
+    _tags = {
+        "info:name": "TFT",
+        "authors": ["phoeenniixx"],
+        "capability:exogenous": True,
+        "capability:multivariate": True,
+        "capability:pred_int": True,
+        "capability:flexible_history_length": False,
+    }
+
+    @classmethod
+    def get_model_cls(cls):
+        """Get model class."""
+        from pytorch_forecasting.models.temporal_fusion_transformer._tft_v2 import TFT
+
+        return TFT
+
+    @classmethod
+    def _get_test_datamodule_from(cls, trainer_kwargs):
+        """Create test dataloaders from trainer_kwargs - following v1 pattern."""
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
+        from pytorch_forecasting.tests._conftest import make_datasets_v2
+        from pytorch_forecasting.tests._data_scenarios import data_with_covariates_v2
+
+        data_with_covariates = data_with_covariates_v2()
+
+        data_loader_default_kwargs = dict(
+            target="target",
+            group_ids=["agency_encoded", "sku_encoded"],
+            add_relative_time_idx=True,
+        )
+
+        data_loader_kwargs = trainer_kwargs.get("data_loader_kwargs", {})
+        data_loader_default_kwargs.update(data_loader_kwargs)
+
+        datasets_info = make_datasets_v2(
+            data_with_covariates, **data_loader_default_kwargs
+        )
+
+        training_dataset = datasets_info["training_dataset"]
+        validation_dataset = datasets_info["validation_dataset"]
+        training_max_time_idx = datasets_info["training_max_time_idx"]
+
+        max_encoder_length = data_loader_kwargs.get("max_encoder_length", 4)
+        max_prediction_length = data_loader_kwargs.get("max_prediction_length", 3)
+        add_relative_time_idx = data_loader_kwargs.get("add_relative_time_idx", True)
+        batch_size = data_loader_kwargs.get("batch_size", 2)
+
+        train_datamodule = EncoderDecoderTimeSeriesDataModule(
+            time_series_dataset=training_dataset,
+            max_encoder_length=max_encoder_length,
+            max_prediction_length=max_prediction_length,
+            add_relative_time_idx=add_relative_time_idx,
+            batch_size=batch_size,
+            train_val_test_split=(0.8, 0.2, 0.0),
+        )
+
+        val_datamodule = EncoderDecoderTimeSeriesDataModule(
+            time_series_dataset=validation_dataset,
+            max_encoder_length=max_encoder_length,
+            max_prediction_length=max_prediction_length,
+            min_prediction_idx=training_max_time_idx,
+            add_relative_time_idx=add_relative_time_idx,
+            batch_size=batch_size,
+            train_val_test_split=(0.0, 1.0, 0.0),
+        )
+
+        test_datamodule = EncoderDecoderTimeSeriesDataModule(
+            time_series_dataset=validation_dataset,
+            max_encoder_length=max_encoder_length,
+            max_prediction_length=max_prediction_length,
+            min_prediction_idx=training_max_time_idx,
+            add_relative_time_idx=add_relative_time_idx,
+            batch_size=1,
+            train_val_test_split=(0.0, 0.0, 1.0),
+        )
+
+        train_datamodule.setup("fit")
+        val_datamodule.setup("fit")
+        test_datamodule.setup("test")
+
+        train_dataloader = train_datamodule.train_dataloader()
+        val_dataloader = val_datamodule.val_dataloader()
+        test_dataloader = test_datamodule.test_dataloader()
+
+        return {
+            "train": train_dataloader,
+            "val": val_dataloader,
+            "test": test_dataloader,
+            "data_module": train_datamodule,
+        }
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Return testing parameter settings for the trainer.
+
+        Returns
+        -------
+        params : dict or list of dict, default = {}
+            Parameters to create testing instances of the class
+            Each dict are parameters to construct an "interesting" test instance, i.e.,
+            `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
+            `create_test_instance` uses the first (or only) dictionary in `params`
+        """
+        return [
+            {},
+            dict(
+                hidden_size=25,
+                attention_head_size=5,
+            ),
+            dict(
+                data_loader_kwargs=dict(max_encoder_length=5, max_prediction_length=3)
+            ),
+            dict(
+                hidden_size=24,
+                attention_head_size=8,
+                data_loader_kwargs=dict(
+                    max_encoder_length=5,
+                    max_prediction_length=3,
+                    add_relative_time_idx=False,
+                ),
+            ),
+            dict(
+                hidden_size=12,
+                data_loader_kwargs=dict(max_encoder_length=7, max_prediction_length=10),
+            ),
+            dict(attention_head_size=2),
+        ]

pytorch_forecasting/models/tide/_tide_metadata.py

@@ -1,9 +1,9 @@
 """TiDE metadata container."""
 
-from pytorch_forecasting.models.base._base_object import _BasePtForecaster
+from pytorch_forecasting.models.base._base_object import _BasePtForecasterV1
 
 
-class TiDEModelMetadata(_BasePtForecaster):
+class TiDEModelMetadata(_BasePtForecasterV1):
     """Metadata container for TiDE Model."""
 
     _tags = {