Repository: Lightning-Universe/lightning-flash

Issue:
Onboard tabular inputs to new Input object
Tracking issue #955


Discussion:

Resolving patch:
diff --git a/flash/tabular/forecasting/cli.py b/flash/tabular/forecasting/cli.py
--- a/flash/tabular/forecasting/cli.py
+++ b/flash/tabular/forecasting/cli.py
@@ -32,9 +32,7 @@ def from_synthetic_ar_data(
     n_series: int = 100,
     max_encoder_length: int = 60,
     max_prediction_length: int = 20,
-    batch_size: int = 32,
-    num_workers: int = 0,
-    **input_transform_kwargs,
+    **data_module_kwargs,
 ) -> TabularForecastingData:
     """Creates and loads a synthetic auto-regressive (AR) data set."""
     data = generate_ar_data(seasonality=seasonality, timesteps=timesteps, n_series=n_series, seed=42)
@@ -53,9 +51,7 @@ def from_synthetic_ar_data(
         max_prediction_length=max_prediction_length,
         train_data_frame=data[lambda x: x.time_idx <= training_cutoff],
         val_data_frame=data,
-        batch_size=batch_size,
-        num_workers=num_workers,
-        **input_transform_kwargs,
+        **data_module_kwargs,
     )
 
 
@@ -72,7 +68,6 @@ def tabular_forecasting():
         },
         finetune=False,
         datamodule_attributes={"parameters"},
-        legacy=True,
     )
 
     cli.trainer.save_checkpoint("tabular_forecasting_model.pt")
diff --git a/flash/tabular/forecasting/data.py b/flash/tabular/forecasting/data.py
--- a/flash/tabular/forecasting/data.py
+++ b/flash/tabular/forecasting/data.py
@@ -13,17 +13,19 @@
 # limitations under the License.
 from copy import copy
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Mapping, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from torch.utils.data.sampler import Sampler
 
-from flash.core.data.callback import BaseDataFetcher
 from flash.core.data.data_module import DataModule
-from flash.core.data.io.input import DataKeys, Input, InputFormat
+from flash.core.data.io.input import DataKeys, InputFormat
+from flash.core.data.io.input_base import Input
 from flash.core.data.io.input_transform import InputTransform
 from flash.core.data.process import Deserializer
 from flash.core.data.properties import ProcessState
 from flash.core.utilities.imports import _FORECASTING_AVAILABLE, _PANDAS_AVAILABLE, requires
+from flash.core.utilities.stages import RunningStage
 
 if _PANDAS_AVAILABLE:
     from pandas.core.frame import DataFrame
@@ -39,31 +41,23 @@ class TimeSeriesDataSetParametersState(ProcessState):
     """A :class:`~flash.core.data.properties.ProcessState` containing ``labels``, a mapping from class index to
     label."""
 
-    time_series_dataset_parameters: Optional[Dict[str, Any]]
+    parameters: Optional[Dict[str, Any]]
 
 
-class TabularForecastingDataFrameInput(Input[DataFrame]):
+class TabularForecastingDataFrameInput(Input):
     @requires("tabular")
-    def __init__(
+    def load_data(
         self,
+        data: DataFrame,
         time_idx: Optional[str] = None,
         target: Optional[Union[str, List[str]]] = None,
         group_ids: Optional[List[str]] = None,
         parameters: Optional[Dict[str, Any]] = None,
-        **input_kwargs: Any,
+        **time_series_dataset_kwargs: Any,
     ):
-        super().__init__()
-        self.time_idx = time_idx
-        self.target = target
-        self.group_ids = group_ids
-        self.input_kwargs = input_kwargs
-
-        self.set_state(TimeSeriesDataSetParametersState(parameters))
-
-    def load_data(self, data: DataFrame, dataset: Optional[Any] = None):
         if self.training:
             time_series_dataset = TimeSeriesDataSet(
-                data, time_idx=self.time_idx, group_ids=self.group_ids, target=self.target, **self.input_kwargs
+                data, time_idx=time_idx, group_ids=group_ids, target=target, **time_series_dataset_kwargs
             )
             parameters = time_series_dataset.get_parameters()
 
@@ -71,9 +65,11 @@ def load_data(self, data: DataFrame, dataset: Optional[Any] = None):
             parameters["data_sample"] = data.iloc[[0]]
 
             self.set_state(TimeSeriesDataSetParametersState(parameters))
-            dataset.parameters = parameters
+            self.parameters = parameters
         else:
-            parameters = copy(self.get_state(TimeSeriesDataSetParametersState).time_series_dataset_parameters)
+            parameters_state = self.get_state(TimeSeriesDataSetParametersState)
+            parameters = parameters or (parameters_state.parameters if parameters_state is not None else None)
+            parameters = copy(parameters)
             if parameters is None:
                 raise MisconfigurationException(
                     "Loading data for evaluation or inference requires parameters from the train data. Either "
@@ -87,10 +83,9 @@ def load_data(self, data: DataFrame, dataset: Optional[Any] = None):
                 predict=True,
                 stop_randomization=True,
             )
-        dataset.time_series_dataset = time_series_dataset
         return time_series_dataset
 
-    def load_sample(self, sample: Mapping[str, Any], dataset: Optional[Any] = None) -> Any:
+    def load_sample(self, sample: Tuple) -> Any:
         return {DataKeys.INPUT: sample[0], DataKeys.TARGET: sample[1]}
 
 
@@ -102,23 +97,21 @@ def __init__(
         test_transform: Optional[Dict[str, Callable]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
         deserializer: Optional[Deserializer] = None,
-        **input_kwargs: Any,
     ):
-        self.input_kwargs = input_kwargs
         super().__init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
             predict_transform=predict_transform,
             inputs={
-                InputFormat.DATAFRAME: TabularForecastingDataFrameInput(**input_kwargs),
+                InputFormat.DATAFRAME: TabularForecastingDataFrameInput,
             },
             deserializer=deserializer,
             default_input=InputFormat.DATAFRAME,
         )
 
     def get_state_dict(self, strict: bool = False) -> Dict[str, Any]:
-        return {**self.transforms, **self.input_kwargs}
+        return {**self.transforms}
 
     @classmethod
     def load_state_dict(cls, state_dict: Dict[str, Any], strict: bool = True) -> "InputTransform":
@@ -151,13 +144,13 @@ def from_data_frame(
         val_transform: Optional[Dict[str, Callable]] = None,
         test_transform: Optional[Dict[str, Callable]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
-        data_fetcher: Optional[BaseDataFetcher] = None,
-        input_transform: Optional[InputTransform] = None,
+        # TODO: Update these when DataModule is updated
         val_split: Optional[float] = None,
         batch_size: int = 4,
-        num_workers: Optional[int] = None,
-        **input_transform_kwargs: Any,
-    ):
+        num_workers: int = 0,
+        sampler: Optional[Type[Sampler]] = None,
+        **time_series_dataset_kwargs: Any,
+    ) -> "TabularForecastingData":
         """Creates a :class:`~flash.tabular.forecasting.data.TabularForecastingData` object from the given data
         frames.
 
@@ -167,69 +160,24 @@ def from_data_frame(
             instead. These can be obtained from the
             :attr:`~flash.tabular.forecasting.data.TabularForecastingData.parameters` attribute of the
             :class:`~flash.tabular.forecasting.data.TabularForecastingData` object that contains your training data.
-
-        Args:
-            time_idx:
-            target: Column denoting the target or list of columns denoting the target.
-            group_ids: List of column names identifying a time series. This means that the group_ids identify a sample
-                together with the time_idx. If you have only one timeseries, set this to the name of column that is
-                constant.
-            parameters: Parameters to use for the timeseries if ``time_idx``, ``target``, and ``group_ids`` are not
-                provided (e.g. when loading data for inference or validation).
-            train_data_frame: The pandas ``DataFrame`` containing the training data.
-            val_data_frame: The pandas ``DataFrame`` containing the validation data.
-            test_data_frame: The pandas ``DataFrame`` containing the testing data.
-            predict_data_frame: The pandas ``DataFrame`` containing the data to use when predicting.
-            train_transform: The dictionary of transforms to use during training which maps
-                :class:`~flash.core.data.io.input_transform.InputTransform` hook names to callable transforms.
-            val_transform: The dictionary of transforms to use during validation which maps
-                :class:`~flash.core.data.io.input_transform.InputTransform` hook names to callable transforms.
-            test_transform: The dictionary of transforms to use during testing which maps
-                :class:`~flash.core.data.io.input_transform.InputTransform` hook names to callable transforms.
-            predict_transform: The dictionary of transforms to use during predicting which maps
-                :class:`~flash.core.data.io.input_transform.InputTransform` hook names to callable transforms.
-            data_fetcher: The :class:`~flash.core.data.callback.BaseDataFetcher` to pass to the
-                :class:`~flash.core.data.data_module.DataModule`.
-            input_transform: The :class:`~flash.core.data.data.InputTransform` to pass to the
-                :class:`~flash.core.data.data_module.DataModule`. If ``None``, ``cls.input_transform_cls``
-                will be constructed and used.
-            val_split: The ``val_split`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
-            batch_size: The ``batch_size`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
-            num_workers: The ``num_workers`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
-            input_transform_kwargs: Additional keyword arguments to use when constructing the input_transform.
-                Will only be used if ``input_transform = None``.
-
-        Returns:
-            The constructed data module.
-
-        Examples::
-
-            data_module = TabularForecastingData.from_data_frame(
-                time_idx="time_idx",
-                target="value",
-                group_ids=["series"],
-                train_data_frame=train_data,
-            )
         """
-
-        return cls.from_input(
+        train_input = TabularForecastingDataFrameInput(
+            RunningStage.TRAINING,
+            train_data_frame,
             time_idx=time_idx,
-            target=target,
             group_ids=group_ids,
-            parameters=parameters,
-            input=InputFormat.DATAFRAME,
-            train_data=train_data_frame,
-            val_data=val_data_frame,
-            test_data=test_data_frame,
-            predict_data=predict_data_frame,
-            train_transform=train_transform,
-            val_transform=val_transform,
-            test_transform=test_transform,
-            predict_transform=predict_transform,
-            data_fetcher=data_fetcher,
-            input_transform=input_transform,
+            target=target,
+            **time_series_dataset_kwargs,
+        )
+        parameters = train_input.parameters if train_input else parameters
+        return cls(
+            train_input,
+            TabularForecastingDataFrameInput(RunningStage.VALIDATING, val_data_frame, parameters=parameters),
+            TabularForecastingDataFrameInput(RunningStage.TESTING, test_data_frame, parameters=parameters),
+            TabularForecastingDataFrameInput(RunningStage.PREDICTING, predict_data_frame, parameters=parameters),
+            input_transform=cls.input_transform_cls(train_transform, val_transform, test_transform, predict_transform),
             val_split=val_split,
             batch_size=batch_size,
             num_workers=num_workers,
-            **input_transform_kwargs,
+            sampler=sampler,
         )


Test patch: