Remove old input transform and data module (#1058)

Lightning-Universe · Dec 13, 2021 · 3e1612b · 3e1612b
1 parent 445b7bc
commit 3e1612b
Show file tree

Hide file tree

Showing 58 changed files with 638 additions and 2,787 deletions.
diff --git a/docs/source/api/data.rst b/docs/source/api/data.rst
@@ -90,12 +90,10 @@ _______________________
     :nosignatures:
     :template: classtemplate.rst
 
-    ~flash.core.data.io.input_transform.BaseInputTransform
-    ~flash.core.data.io.input_transform.DefaultInputTransform
     ~flash.core.data.process.DeserializerMapping
     ~flash.core.data.process.Deserializer
     ~flash.core.data.io.output_transform.OutputTransform
-    ~flash.core.data.io.input_transform.InputTransform
+    ~flash.core.data.input_transform.InputTransform
 
 flash.core.data.properties
 __________________________
@@ -144,9 +142,6 @@ _____________________
     :nosignatures:
     :template: classtemplate.rst
 
-    ~flash.core.data.utils.CurrentFuncContext
-    ~flash.core.data.utils.CurrentRunningStageContext
-    ~flash.core.data.utils.CurrentRunningStageFuncContext
     ~flash.core.data.utils.FuncModule
 
 .. autosummary::

diff --git a/docs/source/api/flash.rst b/docs/source/api/flash.rst
@@ -12,6 +12,6 @@ flash
     ~flash.core.data.callback.FlashCallback
     ~flash.core.data.io.output_transform.OutputTransform
     ~flash.core.data.io.output.Output
-    ~flash.core.data.io.input_transform.InputTransform
+    ~flash.core.data.input_transform.InputTransform
     ~flash.core.model.Task
     ~flash.core.trainer.Trainer
diff --git a/flash/__init__.py b/flash/__init__.py
@@ -21,8 +21,8 @@
 
     from flash.core.data.callback import FlashCallback
     from flash.core.data.data_module import DataModule
+    from flash.core.data.input_transform import InputTransform
     from flash.core.data.io.input import DataKeys, Input
-    from flash.core.data.io.input_transform import InputTransform
     from flash.core.data.io.output import Output
     from flash.core.data.io.output_transform import OutputTransform
     from flash.core.data.process import Serializer

diff --git a/flash/audio/classification/data.py b/flash/audio/classification/data.py
@@ -27,10 +27,10 @@
 )
 from flash.audio.classification.input_transform import AudioClassificationInputTransform
 from flash.core.data.callback import BaseDataFetcher
+from flash.core.data.data_module import DataModule
 from flash.core.data.data_pipeline import DataPipelineState
 from flash.core.data.input_transform import INPUT_TRANSFORM_TYPE
 from flash.core.data.io.input import Input
-from flash.core.data.new_data_module import DataModule
 from flash.core.data.utilities.paths import PATH_TYPE
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.stages import RunningStage

diff --git a/flash/audio/speech_recognition/data.py b/flash/audio/speech_recognition/data.py
@@ -22,10 +22,10 @@
     SpeechRecognitionPathsInput,
 )
 from flash.audio.speech_recognition.output_transform import SpeechRecognitionOutputTransform
+from flash.core.data.data_module import DataModule
 from flash.core.data.data_pipeline import DataPipelineState
 from flash.core.data.input_transform import INPUT_TRANSFORM_TYPE, InputTransform
 from flash.core.data.io.input import Input
-from flash.core.data.new_data_module import DataModule
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.stages import RunningStage
 

diff --git a/flash/core/data/batch.py b/flash/core/data/batch.py
@@ -17,43 +17,14 @@
 from torch import Tensor
 
 from flash.core.data.callback import ControlFlow, FlashCallback
-from flash.core.data.utils import convert_to_modules, CurrentFuncContext, CurrentRunningStageContext
+from flash.core.data.utils import convert_to_modules
 from flash.core.utilities.stages import RunningStage
 
 if TYPE_CHECKING:
-    from flash.core.data.io.input_transform import InputTransform
+    from flash.core.data.input_transform import InputTransform
     from flash.core.data.process import Deserializer
 
 
-class _DeserializeProcessor(torch.nn.Module):
-    def __init__(
-        self,
-        deserializer: "Deserializer",
-        input_transform: "InputTransform",
-        per_sample_transform: Callable,
-        callbacks: Optional[List[FlashCallback]] = None,
-    ):
-        super().__init__()
-        self.input_transform = input_transform
-        self.callback = ControlFlow(callbacks or [])
-        self.deserializer = convert_to_modules(deserializer)
-        self.per_sample_transform = convert_to_modules(per_sample_transform)
-
-        self._current_stage_context = CurrentRunningStageContext(RunningStage.PREDICTING, input_transform, reset=False)
-        self._per_sample_transform_context = CurrentFuncContext("per_sample_transform", input_transform)
-
-    def forward(self, sample: str):
-
-        sample = self.deserializer(sample)
-
-        with self._current_stage_context:
-            with self._per_sample_transform_context:
-                sample = self.per_sample_transform(sample)
-                self.callback.on_per_sample_transform(sample, RunningStage.PREDICTING)
-
-        return sample
-
-
 class _DeserializeProcessorV2(torch.nn.Module):
     def __init__(
         self,