refactor(typing): fix typing accross codebase.

src/stimulus/learner/raytune_learner.py

@@ -16,11 +16,10 @@
 from torch import nn, optim
 from torch.utils.data import DataLoader
 
-from stimulus.data.handlertorch import TorchDataset
-from stimulus.data.loaders import EncoderLoader
+from stimulus.data.data_handlers import TorchDataset
 from stimulus.learner.predict import PredictWrapper
 from stimulus.utils.generic_utils import set_general_seeds
-from stimulus.utils.yaml_data import YamlSplitTransformDict
+from stimulus.data.interface.data_config_schema import SplitTransformDict
 from stimulus.utils.yaml_model_schema import RayTuneModel
 
 
@@ -36,10 +35,10 @@ class TuneWrapper:
     def __init__(
         self,
         model_config: RayTuneModel,
-        data_config: YamlSplitTransformDict,
+        data_config: SplitTransformDict,
         model_class: nn.Module,
         data_path: str,
-        encoder_loader: EncoderLoader,
+        #encoder_loader: EncoderLoader,
         seed: int,
         ray_results_dir: Optional[str] = None,
         tune_run_name: Optional[str] = None,
@@ -113,15 +112,15 @@ def __init__(
         self.tuner = self.tuner_initialization(
             data_config=data_config,
             data_path=data_path,
-            encoder_loader=encoder_loader,
+            #encoder_loader=encoder_loader,
             autoscaler=autoscaler,
         )
 
     def tuner_initialization(
         self,
-        data_config: YamlSplitTransformDict,
+        data_config: SplitTransformDict,
         data_path: str,
-        encoder_loader: EncoderLoader,
+        #encoder_loader: EncoderLoader,
         *,
         autoscaler: bool = False,
     ) -> tune.Tuner:
@@ -155,13 +154,13 @@ def tuner_initialization(
         training = TorchDataset(
             data_config=data_config,
             csv_path=data_path,
-            encoder_loader=encoder_loader,
+            #encoder_loader=encoder_loader,
             split=0,
         )
         validation = TorchDataset(
             data_config=data_config,
             csv_path=data_path,
-            encoder_loader=encoder_loader,
+            #encoder_loader=encoder_loader,
             split=1,
         )
 

src/stimulus/typing/__init__.py

@@ -13,19 +13,17 @@
 from stimulus.data.data_handlers import (
     DatasetHandler,
     DatasetLoader,
-    DatasetManager,
     DatasetProcessor,
-    EncodeManager,
-    SplitManager,
-    TransformManager,
+    TorchDataset,
 )
 from stimulus.data.encoding.encoders import AbstractEncoder as Encoder
-from stimulus.data.handlertorch import TorchDataset
-from stimulus.data.loaders import EncoderLoader, SplitLoader, TransformLoader
-from stimulus.data.splitters.splitters import AbstractSplitter as Splitter
-from stimulus.data.transform.data_transformation_generators import (
-    AbstractDataTransformer as Transform,
+from stimulus.data.interface.data_config_parser import (
+    create_encoders,
+    create_splitter,
+    create_transforms,
 )
+from stimulus.data.splitting import AbstractSplitter as Splitter
+from stimulus.data.transforming.transforms import AbstractTransform as Transform
 from stimulus.learner.predict import PredictWrapper
 from stimulus.learner.raytune_learner import CheckpointDict, TuneModel, TuneWrapper
 from stimulus.learner.raytune_parser import (
@@ -35,18 +33,18 @@
     TuneParser,
 )
 from stimulus.utils.performance import Performance
-from stimulus.utils.yaml_data import (
-    YamlColumns,
-    YamlColumnsEncoder,
-    YamlConfigDict,
-    YamlGlobalParams,
-    YamlSchema,
-    YamlSplit,
-    YamlSplitConfigDict,
-    YamlSplitTransformDict,
-    YamlTransform,
-    YamlTransformColumns,
-    YamlTransformColumnsTransformation,
+from stimulus.data.interface.data_config_schema import (
+    Columns,
+    ColumnsEncoder,
+    ConfigDict,
+    GlobalParams,
+    Schema,
+    Split,
+    SplitConfigDict,
+    SplitTransformDict,
+    Transform,
+    TransformColumns,
+    TransformColumnsTransformation,
 )
 from stimulus.utils.yaml_model_schema import (
     CustomTunableParameter,
@@ -59,32 +57,24 @@
     TunableParameter,
     Tune,
     TuneParams,
-    YamlRayConfigLoader,
+    RayConfigLoader,
 )
 
-# data/data_handlers.py
-
-DataManager: TypeAlias = DatasetManager | EncodeManager | SplitManager | TransformManager
-
-# data/experiments.py
-
-Loader: TypeAlias = DatasetLoader | EncoderLoader | TransformLoader | SplitLoader
-
 # learner/raytune_parser.py
 
 RayTuneData: TypeAlias = RayTuneMetrics | RayTuneOptimizer | RayTuneResult
 
 # utils/yaml_data.py
 
-YamlData: TypeAlias = (
-    YamlColumns
-    | YamlColumnsEncoder
-    | YamlConfigDict
-    | YamlGlobalParams
-    | YamlSchema
-    | YamlSplit
-    | YamlSplitConfigDict
-    | YamlTransform
-    | YamlTransformColumns
-    | YamlTransformColumnsTransformation
+Data: TypeAlias = (
+    Columns
+    | ColumnsEncoder
+    | ConfigDict
+    | GlobalParams
+    | Schema
+    | Split
+    | SplitConfigDict
+    | Transform
+    | TransformColumns
+    | TransformColumnsTransformation
 )

src/stimulus/utils/yaml_model_schema.py

@@ -120,7 +120,7 @@ class RayTuneModel(pydantic.BaseModel):
     tune: Tune
 
 
-class YamlRayConfigLoader:
+class RayConfigLoader:
     """Load and convert YAML configurations to Ray Tune format.
 
     This class handles loading model configurations and converting them into

tests/typing/test_typing.py

@@ -13,11 +13,8 @@ def test_data_handlers_types() -> None:
         from stimulus.typing import (
             DatasetHandler,
             DatasetLoader,
-            DatasetManager,
             DatasetProcessor,
-            EncodeManager,
-            SplitManager,
-            TransformManager,
+            TorchDataset,
         )
     except ImportError:
         pytest.fail("Failed to import Data Handlers types")
@@ -34,29 +31,30 @@ def test_learner_types() -> None:
             TuneModel,
             TuneParser,
             TuneWrapper,
+            CheckpointDict,
         )
     except ImportError:
         pytest.fail("Failed to import Learner types")
 
 
-def test_yaml_data_types() -> None:
-    """Test the YAML data types."""
+def test_data_config_schema_types() -> None:
+    """Test the data config schema types."""
     try:
         from stimulus.typing import (
-            YamlColumns,
-            YamlColumnsEncoder,
-            YamlConfigDict,
-            YamlGlobalParams,
-            YamlSchema,
-            YamlSplit,
-            YamlSplitConfigDict,
-            YamlSplitTransformDict,
-            YamlTransform,
-            YamlTransformColumns,
-            YamlTransformColumnsTransformation,
+            Columns,
+            ColumnsEncoder,
+            ConfigDict,
+            GlobalParams,
+            Schema,
+            Split,
+            SplitConfigDict,
+            SplitTransformDict,
+            Transform,
+            TransformColumns,
+            TransformColumnsTransformation,
         )
     except ImportError:
-        pytest.fail("Failed to import YAML Data types")
+        pytest.fail("Failed to import Data Config Schema types")
 
 
 def test_yaml_model_schema_types() -> None:
@@ -73,7 +71,7 @@ def test_yaml_model_schema_types() -> None:
             TunableParameter,
             Tune,
             TuneParams,
-            YamlRayConfigLoader,
+            RayConfigLoader,
         )
     except ImportError:
         pytest.fail("Failed to import YAML Model Schema types")
@@ -82,6 +80,6 @@ def test_yaml_model_schema_types() -> None:
 def test_type_aliases() -> None:
     """Test the type aliases."""
     try:
-        from stimulus.typing import DataManager, Loader, RayTuneData, YamlData
+        from stimulus.typing import RayTuneData, Data
     except ImportError:
         pytest.fail("Failed to import Type Aliases")