Repository: mathysgrapotte/stimulus-py

Issue:
change: Major refactoring on encoder -> loader -> manager flow to align with Unix philosophy
### Is your change request related to a problem? Please describe.
The current data processing pipeline suffers from:
1. Overly complex class hierarchies (Loader -> Manager -> Handler)
2. Tight coupling between configuration parsing and component creation
3. Hidden state management in loader classes
4. Obscured data flow through multiple abstraction layers

This violates principles of:
- Single responsibility
- Composability
- Clarity over cleverness

!!! adapt documentation !!!



### Describe the solution you'd like
**Refactor the architecture to:**
1. Replace class-based loaders with pure factory functions
2. Remove intermediate manager classes
3. Use explicit functional composition
4. Separate configuration parsing from component creation
5. Establish clear data pipelines without hidden state

### Describe alternatives you've considered
<!-- A clear and concise description of any alternative solutions you've considered. -->

### Additional context
**Affected Components:**

```
src/stimulus/data/
├── encoders.py       # No changes
├── loaders.py        # Major rewrite
├── data_handlers.py  # Major refactor
└── config_parser.py  # New module

src/stimulus/utils/
└── yaml_data.py      # Validation updates
```


Discussion:

Patch:
diff --git a/src/stimulus/cli/check_model.py b/src/stimulus/cli/check_model.py
index 1314893..2fec396 100755
--- a/src/stimulus/cli/check_model.py
+++ b/src/stimulus/cli/check_model.py
@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader
 
 from stimulus.data import handlertorch, loaders
 from stimulus.learner import raytune_learner
-from stimulus.utils import launch_utils, yaml_data, yaml_model_schema
+from stimulus.utils import model_file_interface, yaml_data, yaml_model_schema
 
 logger = logging.getLogger(__name__)
 
@@ -134,7 +134,7 @@ def main(
 
     logger.info("Dataset loaded successfully.")
 
-    model_class = launch_utils.import_class_from_file(model_path)
+    model_class = model_file_interface.import_class_from_file(model_path)
 
     logger.info("Model class loaded successfully.")
 
diff --git a/src/stimulus/cli/predict.py b/src/stimulus/cli/predict.py
index 7834151..f1468e6 100755
--- a/src/stimulus/cli/predict.py
+++ b/src/stimulus/cli/predict.py
@@ -10,9 +10,9 @@ import polars as pl
 import torch
 from torch.utils.data import DataLoader
 
-from stimulus.data.handlertorch import TorchDataset
+from stimulus.data.data_handlers import TorchDataset
 from stimulus.learner.predict import PredictWrapper
-from stimulus.utils.launch_utils import get_experiment, import_class_from_file
+from stimulus.utils.model_file_interface import get_experiment, import_class_from_file
 
 
 def get_args() -> argparse.Namespace:
diff --git a/src/stimulus/cli/tuning.py b/src/stimulus/cli/tuning.py
index 65802ef..0581887 100755
--- a/src/stimulus/cli/tuning.py
+++ b/src/stimulus/cli/tuning.py
@@ -12,7 +12,7 @@ import yaml
 
 from stimulus.data import loaders
 from stimulus.learner import raytune_learner, raytune_parser
-from stimulus.utils import launch_utils, yaml_data, yaml_model_schema
+from stimulus.utils import model_file_interface, yaml_data, yaml_model_schema
 
 logger = logging.getLogger(__name__)
 
@@ -186,7 +186,7 @@ def main(
         column_config=data_config.columns,
     )
 
-    model_class = launch_utils.import_class_from_file(model_path)
+    model_class = model_file_interface.import_class_from_file(model_path)
 
     ray_config_loader = yaml_model_schema.YamlRayConfigLoader(model=model_config)
     ray_config_model = ray_config_loader.get_config()
diff --git a/src/stimulus/data/data_handlers.py b/src/stimulus/data/data_handlers.py
index 4b53089..34ad46c 100644
--- a/src/stimulus/data/data_handlers.py
+++ b/src/stimulus/data/data_handlers.py
@@ -22,249 +22,15 @@ The data handling pipeline consists of:
 See titanic.yaml in tests/test_data/titanic/ for an example configuration file format.
 """
 
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import numpy as np
 import polars as pl
 import torch
-import yaml
 
-from stimulus.data import loaders
-from stimulus.utils import yaml_data
-
-
-class DatasetManager:
-    """Class for managing the dataset.
-
-    This class handles loading and organizing dataset configuration from YAML files.
-    It manages column categorization into input, label and meta types based on the config.
-
-    Attributes:
-        config (dict): The loaded configuration dictionary from YAML
-        column_categories (dict): Dictionary mapping column types to lists of column names
-
-    Methods:
-        _load_config(config_path: str) -> dict: Loads the config from a YAML file.
-        categorize_columns_by_type() -> dict: Organizes the columns into input, label, meta based on the config.
-    """
-
-    def __init__(
-        self,
-        config_dict: yaml_data.YamlSplitConfigDict,
-    ) -> None:
-        """Initialize the DatasetManager."""
-        # self.config = self._load_config(config_path)
-        self.config: yaml_data.YamlSplitTransformDict = config_dict
-        self.column_categories = self.categorize_columns_by_type()
-
-    def categorize_columns_by_type(self) -> dict:
-        """Organizes columns from config into input, label, and meta categories.
-
-        Reads the column definitions from the config and sorts them into categories
-        based on their column_type field.
-
-        Returns:
-            dict: Dictionary containing lists of column names for each category:
-                {
-                    "input": ["col1", "col2"],  # Input columns
-                    "label": ["target"],        # Label/output columns
-                    "meta": ["id"]     # Metadata columns
-                }
-
-        Example:
-            >>> manager = DatasetManager("config.yaml")
-            >>> categories = manager.categorize_columns_by_type()
-            >>> print(categories)
-            {
-                'input': ['hello', 'bonjour'],
-                'label': ['ciao'],
-                'meta': ["id"]
-            }
-        """
-        input_columns = []
-        label_columns = []
-        meta_columns = []
-        for column in self.config.columns:
-            if column.column_type == "input":
-                input_columns.append(column.column_name)
-            elif column.column_type == "label":
-                label_columns.append(column.column_name)
-            elif column.column_type == "meta":
-                meta_columns.append(column.column_name)
-
-        return {"input": input_columns, "label": label_columns, "meta": meta_columns}
-
-    # TODO: Remove or change this function as the config is now preloaded
-    def _load_config(self, config_path: str) -> yaml_data.YamlSplitConfigDict:
-        """Loads and parses a YAML configuration file.
-
-        Args:
-            config_path (str): Path to the YAML config file
-
-        Returns:
-            dict: Parsed configuration dictionary
-
-        Example:
-            >>> manager = DatasetManager()
-            >>> config = manager._load_config("config.yaml")
-            >>> print(config["columns"][0]["column_name"])
-            'hello'
-        """
-        with open(config_path) as file:
-            # FIXME: cette fonction est appellé pour test_shuffle_csv et test_tune
-            return yaml_data.YamlSplitConfigDict(**yaml.safe_load(file))
-            return yaml_data.YamlSplitTransformDict(**yaml.safe_load(file))
-
-    def get_split_columns(self) -> list[str]:
-        """Get the columns that are used for splitting."""
-        return self.config.split.split_input_columns
-
-    def get_transform_logic(self) -> dict:
-        """Get the transformation logic.
-
-        Returns a dictionary in the following structure :
-        {
-            "transformation_name": str,
-            "transformations": list[tuple[str, str, dict]]
-        }
-        """
-        transformation_logic = {
-            "transformation_name": self.config.transforms.transformation_name,
-            "transformations": [],
-        }
-        for column in self.config.transforms.columns:
-            for transformation in column.transformations:
-                transformation_logic["transformations"].append(
-                    (column.column_name, transformation.name, transformation.params),
-                )
-        return transformation_logic
-
-
-class EncodeManager:
-    """Manages the encoding of data columns using configured encoders.
-
-    This class handles encoding of data columns based on the encoders specified in the
-    configuration. It uses an EncoderLoader to get the appropriate encoder for each column
-    and applies the encoding.
-
-    Attributes:
-        encoder_loader (experiments.EncoderLoader): Loader that provides encoders based on config.
-
-    Example:
-        >>> encoder_loader = EncoderLoader(config)
-        >>> encode_manager = EncodeManager(encoder_loader)
-        >>> data = ["ACGT", "TGCA", "GCTA"]
-        >>> encoded = encode_manager.encode_column("dna_seq", data)
-        >>> print(encoded.shape)
-        torch.Size([3, 4, 4])  # 3 sequences, length 4, one-hot encoded
-    """
-
-    def __init__(
-        self,
-        encoder_loader: loaders.EncoderLoader,
-    ) -> None:
-        """Initialize the EncodeManager.
-
-        Args:
-            encoder_loader: Loader that provides encoders based on configuration.
-        """
-        self.encoder_loader = encoder_loader
-
-    def encode_column(self, column_name: str, column_data: list) -> torch.Tensor:
-        """Encodes a column of data using the configured encoder.
-
-        Gets the appropriate encoder for the column from the encoder_loader and uses it
-        to encode all the data in the column.
-
-        Args:
-            column_name: Name of the column to encode.
-            column_data: List of data values from the column to encode.
-
-        Returns:
-            Encoded data as a torch.Tensor. The exact shape depends on the encoder used.
-
-        Example:
-            >>> data = ["ACGT", "TGCA"]
-            >>> encoded = encode_manager.encode_column("dna_seq", data)
-            >>> print(encoded.shape)
-            torch.Size([2, 4, 4])  # 2 sequences, length 4, one-hot encoded
-        """
-        encode_all_function = self.encoder_loader.get_function_encode_all(column_name)
-        return encode_all_function(column_data)
-
-    def encode_columns(self, column_data: dict) -> dict:
-        """Encodes multiple columns of data using the configured encoders.
-
-        Gets the appropriate encoder for each column from the encoder_loader and encodes
-        all data values in those columns.
-
-        Args:
-            column_data: Dict mapping column names to lists of data values to encode.
-
-        Returns:
-            Dict mapping column names to their encoded tensors. The exact shape of each
-            tensor depends on the encoder used for that column.
-
-        Example:
-            >>> data = {"dna_seq": ["ACGT", "TGCA"], "labels": ["1", "2"]}
-            >>> encoded = encode_manager.encode_columns(data)
-            >>> print(encoded["dna_seq"].shape)
-            torch.Size([2, 4, 4])  # 2 sequences, length 4, one-hot encoded
-        """
-        return {col: self.encode_column(col, values) for col, values in column_data.items()}
-
-    def encode_dataframe(self, dataframe: pl.DataFrame) -> dict[str, torch.Tensor]:
-        """Encode the dataframe using the encoders."""
-        return {col: self.encode_column(col, dataframe[col].to_list()) for col in dataframe.columns}
-
-
-class TransformManager:
-    """Class for managing the transformations."""
-
-    def __init__(
-        self,
-        transform_loader: loaders.TransformLoader,
-    ) -> None:
-        """Initialize the TransformManager."""
-        self.transform_loader = transform_loader
-
-    def transform_column(
-        self,
-        column_name: str,
-        transform_name: str,
-        column_data: list,
-    ) -> tuple[list, bool]:
-        """Transform a column of data using the specified transformation.
-
-        Args:
-            column_name (str): The name of the column to transform.
-            transform_name (str): The name of the transformation to use.
-            column_data (list): The data to transform.
-
-        Returns:
-            list: The transformed data.
-            bool: Whether the transformation added new rows to the data.
-        """
-        transformer = self.transform_loader.__getattribute__(column_name)[transform_name]
-        return transformer.transform_all(column_data), transformer.add_row
-
-
-class SplitManager:
-    """Class for managing the splitting."""
-
-    def __init__(
-        self,
-        split_loader: loaders.SplitLoader,
-    ) -> None:
-        """Initialize the SplitManager."""
-        self.split_loader = split_loader
-
-    def get_split_indices(
-        self,
-        data: dict,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        """Get the indices for train, validation, and test splits."""
-        return self.split_loader.get_function_split()(data)
+import stimulus.data.encoding.encoders as encoders_module
+import stimulus.data.splitting.splitters as splitters_module
+import stimulus.data.transforming.transforms as transforms_module
 
 
 class DatasetHandler:
@@ -277,21 +43,17 @@ class DatasetHandler:
         encoder_manager (EncodeManager): Manager for handling data encoding operations.
         transform_manager (TransformManager): Manager for handling data transformations.
         split_manager (SplitManager): Manager for handling dataset splitting.
-        dataset_manager (DatasetManager): Manager for organizing dataset columns and config.
     """
 
     def __init__(
         self,
-        data_config: yaml_data.YamlSplitTransformDict,
         csv_path: str,
     ) -> None:
         """Initialize the DatasetHandler with required config.
 
         Args:
-            data_config (yaml_data.YamlSplitTransformDict): A YamlSplitTransformDict object holding the config.
             csv_path (str): Path to the CSV data file.
         """
-        self.dataset_manager = DatasetManager(data_config)
         self.columns = self.read_csv_header(csv_path)
         self.data = self.load_csv(csv_path)
 
@@ -343,17 +105,25 @@ class DatasetHandler:
 class DatasetProcessor(DatasetHandler):
     """Class for loading dataset, applying transformations and splitting."""
 
-    def __init__(self, data_config: yaml_data.YamlSplitTransformDict, csv_path: str) -> None:
+    def __init__(
+        self,
+        csv_path: str,
+        transforms: dict[str, list[transforms_module.AbstractTransform]],
+        split_columns: list[str],
+        splitter: splitters_module.AbstractSplitter,
+    ) -> None:
         """Initialize the DatasetProcessor."""
-        super().__init__(data_config, csv_path)
+        super().__init__(csv_path)
+        self.transforms = transforms
+        self.split_columns = split_columns
+        self.splitter = splitter
 
-    def add_split(self, split_manager: SplitManager, *, force: bool = False) -> None:
+    def add_split(self, *, force: bool = False) -> None:
         """Add a column specifying the train, validation, test splits of the data.
 
         An error exception is raised if the split column is already present in the csv file. This behaviour can be overriden by setting force=True.
 
         Args:
-            split_manager (SplitManager): Manager for handling dataset splitting
             force (bool): If True, the split column present in the csv file will be overwritten.
         """
         if ("split" in self.columns) and (not force):
@@ -361,11 +131,10 @@ class DatasetProcessor(DatasetHandler):
                 "The category split is already present in the csv file. If you want to still use this function, set force=True",
             )
         # get relevant split columns from the dataset_manager
-        split_columns = self.dataset_manager.get_split_columns()
-        split_input_data = self.select_columns(split_columns)
+        split_input_data = self.select_columns(self.split_columns)
 
         # get the split indices
-        train, validation, test = split_manager.get_split_indices(split_input_data)
+        train, validation, test = self.splitter.get_split_indexes(split_input_data)
 
         # add the split column to the data
         split_column = np.full(len(self.data), -1).astype(int)
@@ -377,35 +146,32 @@ class DatasetProcessor(DatasetHandler):
         if "split" not in self.columns:
             self.columns.append("split")
 
-    def apply_transformation_group(self, transform_manager: TransformManager) -> None:
-        """Apply the transformation group to the data."""
-        for (
-            column_name,
-            transform_name,
-            _params,
-        ) in self.dataset_manager.get_transform_logic()["transformations"]:
-            transformed_data, add_row = transform_manager.transform_column(
-                column_name,
-                transform_name,
-                self.data[column_name],
-            )
-            if add_row:
-                new_rows = self.data.with_columns(
-                    pl.Series(column_name, transformed_data),
-                )
-                self.data = pl.vstack(self.data, new_rows)
-            else:
-                self.data = self.data.with_columns(
-                    pl.Series(column_name, transformed_data),
-                )
-
-    def shuffle_labels(self, seed: Optional[float] = None) -> None:
+    def apply_transformations(self) -> None:
+        """Apply the transformation group to the data.
+
+        Applies all transformations defined in self.transforms to their corresponding columns.
+        Each column can have multiple transformations that are applied sequentially.
+        """
+        for column_name, transforms_list in self.transforms.items():
+            for transform in transforms_list:
+                transformed_data = transform.transform(self.data[column_name])
+
+                if transform.add_row:
+                    new_rows = self.data.with_columns(
+                        pl.Series(column_name, transformed_data),
+                    )
+                    self.data = pl.vstack(self.data, new_rows)
+                else:
+                    self.data = self.data.with_columns(
+                        pl.Series(column_name, transformed_data),
+                    )
+
+    def shuffle_labels(self, label_columns: list[str], seed: Optional[float] = None) -> None:
         """Shuffles the labels in the data."""
         # set the np seed
         np.random.seed(seed)
 
-        label_keys = self.dataset_manager.column_categories["label"]
-        for key in label_keys:
+        for key in label_columns:
             self.data = self.data.with_columns(
                 pl.Series(key, np.random.permutation(list(self.data[key]))),
             )
@@ -416,15 +182,41 @@ class DatasetLoader(DatasetHandler):
 
     def __init__(
         self,
-        data_config: yaml_data.YamlSplitTransformDict,
+        encoders: dict[str, encoders_module.AbstractEncoder],
+        input_columns: list[str],
+        label_columns: list[str],
+        meta_columns: list[str],
         csv_path: str,
-        encoder_loader: loaders.EncoderLoader,
-        split: Union[int, None] = None,
+        split: Optional[int] = None,
     ) -> None:
         """Initialize the DatasetLoader."""
-        super().__init__(data_config, csv_path)
-        self.encoder_manager = EncodeManager(encoder_loader)
+        super().__init__(csv_path)
+        self.encoders = encoders
         self.data = self.load_csv_per_split(csv_path, split) if split is not None else self.load_csv(csv_path)
+        self.input_columns = input_columns
+        self.label_columns = label_columns
+        self.meta_columns = meta_columns
+
+    def encode_dataframe(self, dataframe: pl.DataFrame) -> dict[str, torch.Tensor]:
+        """Encode the dataframe columns using the configured encoders.
+
+        Takes a polars DataFrame and encodes each column using its corresponding encoder
+        from self.encoders.
+
+        Args:
+            dataframe: Polars DataFrame containing the columns to encode
+
+        Returns:
+            Dict mapping column names to their encoded tensors. The exact shape of each
+            tensor depends on the encoder used for that column.
+
+        Example:
+            >>> df = pl.DataFrame({"dna_seq": ["ACGT", "TGCA"], "labels": [1, 2]})
+            >>> encoded = dataset_loader.encode_dataframe(df)
+            >>> print(encoded["dna_seq"].shape)
+            torch.Size([2, 4, 4])  # 2 sequences, length 4, one-hot encoded
+        """
+        return {col: self.encoders[col].encode_all(dataframe[col].to_list()) for col in dataframe.columns}
 
     def get_all_items(self) -> tuple[dict, dict, dict]:
         """Get the full dataset as three separate dictionaries for inputs, labels and metadata.
@@ -445,14 +237,9 @@ class DatasetLoader(DatasetHandler):
             >>> print(meta_dict.keys())
             dict_keys(['passenger_id'])
         """
-        input_columns, label_columns, meta_columns = (
-            self.dataset_manager.column_categories["input"],
-            self.dataset_manager.column_categories["label"],
-            self.dataset_manager.column_categories["meta"],
-        )
-        input_data = self.encoder_manager.encode_dataframe(self.data[input_columns])
-        label_data = self.encoder_manager.encode_dataframe(self.data[label_columns])
-        meta_data = {key: self.data[key].to_list() for key in meta_columns}
+        input_data = self.encode_dataframe(self.data[self.input_columns])
+        label_data = self.encode_dataframe(self.data[self.label_columns])
+        meta_data = {key: self.data[key].to_list() for key in self.meta_columns}
         return input_data, label_data, meta_data
 
     def get_all_items_and_length(self) -> tuple[tuple[dict, dict, dict], int]:
@@ -463,10 +250,6 @@ class DatasetLoader(DatasetHandler):
         """Load the part of csv file that has the specified split value.
 
         Split is a number that for 0 is train, 1 is validation, 2 is test.
-        This is accessed through the column with category `split`. Example column name could be `split:split:int`.
-
-        NOTE that the aim of having this function is that depending on the training, validation and test scenarios,
-        we are gonna load only the relevant data for it.
         """
         if "split" not in self.columns:
             raise ValueError("The category split is not present in the csv file")
@@ -489,10 +272,6 @@ class DatasetLoader(DatasetHandler):
         Args:
             idx: The index of the data to be returned, it can be a single index, a list of indexes or a slice
         """
-        input_columns = self.dataset_manager.column_categories["input"]
-        label_columns = self.dataset_manager.column_categories["label"]
-        meta_columns = self.dataset_manager.column_categories["meta"]
-
         # Handle different index types
         if isinstance(idx, slice):
             # Get the actual indices for the slice
@@ -501,36 +280,57 @@ class DatasetLoader(DatasetHandler):
             data_at_index = self.data.slice(start, stop - start)
 
             # Process DataFrame
-            input_data = self.encoder_manager.encode_dataframe(
-                data_at_index[input_columns],
+            input_data = self.encode_dataframe(
+                data_at_index[self.input_columns],
             )
-            label_data = self.encoder_manager.encode_dataframe(
-                data_at_index[label_columns],
+            label_data = self.encode_dataframe(
+                data_at_index[self.label_columns],
             )
-            meta_data = {key: data_at_index[key].to_list() for key in meta_columns}
+            meta_data = {key: data_at_index[key].to_list() for key in self.meta_columns}
 
         elif isinstance(idx, int):
             # For single row, convert to dict with column names as keys
             row_dict = dict(zip(self.data.columns, self.data.row(idx)))
 
             # Create single-row DataFrames for encoding
-            input_df = pl.DataFrame({col: [row_dict[col]] for col in input_columns})
-            label_df = pl.DataFrame({col: [row_dict[col]] for col in label_columns})
+            input_df = pl.DataFrame({col: [row_dict[col]] for col in self.input_columns})
+            label_df = pl.DataFrame({col: [row_dict[col]] for col in self.label_columns})
 
-            input_data = self.encoder_manager.encode_dataframe(input_df)
-            label_data = self.encoder_manager.encode_dataframe(label_df)
-            meta_data = {key: [row_dict[key]] for key in meta_columns}
+            input_data = self.encode_dataframe(input_df)
+            label_data = self.encode_dataframe(label_df)
+            meta_data = {key: [row_dict[key]] for key in self.meta_columns}
 
         else:  # list or other sequence
             data_at_index = self.data.select(idx)
 
             # Process DataFrame
-            input_data = self.encoder_manager.encode_dataframe(
-                data_at_index[input_columns],
+            input_data = self.encode_dataframe(
+                data_at_index[self.input_columns],
             )
-            label_data = self.encoder_manager.encode_dataframe(
-                data_at_index[label_columns],
+            label_data = self.encode_dataframe(
+                data_at_index[self.label_columns],
             )
-            meta_data = {key: data_at_index[key].to_list() for key in meta_columns}
+            meta_data = {key: data_at_index[key].to_list() for key in self.meta_columns}
 
         return input_data, label_data, meta_data
+
+
+class TorchDataset(torch.utils.data.Dataset):
+    """Class for creating a torch dataset."""
+
+    def __init__(
+        self,
+        loader: DatasetLoader,
+    ) -> None:
+        """Initialize the TorchDataset.
+
+        Args:
+            loader: A DatasetLoader instance
+        """
+        self.loader = loader
+
+    def __len__(self) -> int:
+        return len(self.loader)
+
+    def __getitem__(self, idx: int) -> tuple[dict, dict, dict]:
+        return self.loader[idx]
diff --git a/src/stimulus/data/handlertorch.py b/src/stimulus/data/handlertorch.py
deleted file mode 100644
index 3b89140..0000000
--- a/src/stimulus/data/handlertorch.py
+++ /dev/null
@@ -1,40 +0,0 @@
-"""This file provides the class API for handling the data in pytorch using the Dataset and Dataloader classes."""
-
-from typing import Optional
-
-from torch.utils.data import Dataset
-
-from stimulus.data import data_handlers, loaders
-from stimulus.utils.yaml_data import YamlSplitTransformDict
-
-
-class TorchDataset(Dataset):
-    """Class for creating a torch dataset."""
-
-    def __init__(
-        self,
-        data_config: YamlSplitTransformDict,
-        csv_path: str,
-        encoder_loader: loaders.EncoderLoader,
-        split: Optional[int] = None,
-    ) -> None:
-        """Initialize the TorchDataset.
-
-        Args:
-            data_config: A YamlSplitTransformDict holding the configuration.
-            csv_path: Path to the CSV data file
-            encoder_loader: Encoder loader instance
-            split: Optional tuple containing split information
-        """
-        self.loader = data_handlers.DatasetLoader(
-            data_config=data_config,
-            csv_path=csv_path,
-            encoder_loader=encoder_loader,
-            split=split,
-        )
-
-    def __len__(self) -> int:
-        return len(self.loader)
-
-    def __getitem__(self, idx: int) -> tuple[dict, dict, dict]:
-        return self.loader[idx]
diff --git a/src/stimulus/data/interface/__init__.py b/src/stimulus/data/interface/__init__.py
new file mode 100644
index 0000000..ba4bd57
--- /dev/null
+++ b/src/stimulus/data/interface/__init__.py
@@ -0,0 +1,1 @@
+"""Module for interfacing configs with the pipeline."""
diff --git a/src/stimulus/data/interface/data_config_parser.py b/src/stimulus/data/interface/data_config_parser.py
new file mode 100644
index 0000000..0bc58af
--- /dev/null
+++ b/src/stimulus/data/interface/data_config_parser.py
@@ -0,0 +1,351 @@
+"""Module for parsing data configs."""
+
+from typing import Any
+
+import yaml
+
+from stimulus.data.encoding import encoders as encoders_module
+from stimulus.data.interface.data_config_schema import (
+    Columns,
+    ConfigDict,
+    Split,
+    SplitConfigDict,
+    SplitTransformDict,
+    Transform,
+)
+from stimulus.data.splitting import splitters as splitters_module
+from stimulus.data.transforming import transforms as transforms_module
+
+
+def _instantiate_component(module: Any, name: str, params: dict) -> Any:
+    """Generic component instantiator.
+
+    Args:
+        module: The module to instantiate the component from.
+        name: The name of the component to instantiate.
+        params: The parameters to instantiate the component with.
+
+    Returns:
+        The instantiated component.
+    """
+    try:
+        return getattr(module, name)(**params or {})
+    except (AttributeError, TypeError) as e:
+        raise ValueError(f"Failed to create {name} from {module.__name__}") from e
+
+
+def create_encoders(column_config: list[Columns]) -> dict[str, encoders_module.AbstractEncoder]:
+    """Factory for creating encoders from config."""
+    return {
+        column.column_name: _instantiate_component(
+            module=encoders_module,
+            name=column.encoder[0].name,
+            params=column.encoder[0].params,
+        )
+        for column in column_config
+    }
+
+
+def create_transforms(transform_config: list[Transform]) -> dict[str, list[Any]]:
+    """Factory for creating transforms from config.
+
+    Args:
+        transform_config: List of Transform objects from the YAML config
+
+    Returns:
+        Dictionary mapping column names to lists of instantiated transform objects
+    """
+    transforms = {}
+    for transform in transform_config:
+        for column in transform.columns:
+            transforms[column.column_name] = [
+                _instantiate_component(
+                    module=transforms_module,
+                    name=transformation.name,
+                    params=transformation.params,
+                )
+                for transformation in column.transformations
+            ]
+    return transforms
+
+
+def create_splitter(split_config: Split) -> splitters_module.AbstractSplitter:
+    """Factory for creating splitters from config."""
+    return _instantiate_component(
+        module=splitters_module,
+        name=split_config.split_method,
+        params=split_config.params,
+    )
+
+
+def parse_split_transform_config(
+    config: SplitTransformDict,
+) -> tuple[
+    dict[str, encoders_module.AbstractEncoder],
+    list[str],
+    list[str],
+    list[str],
+]:
+    """Parse the configuration and return a dictionary of the parsed configuration.
+
+    Args:
+        config: The configuration to parse.
+
+    Returns:
+        A tuple of the parsed configuration.
+    """
+    encoders = create_encoders(config.columns)
+    input_columns = [column.column_name for column in config.columns if column.column_type == "input"]
+    label_columns = [column.column_name for column in config.columns if column.column_type == "label"]
+    meta_columns = [column.column_name for column in config.columns if column.column_type == "meta"]
+
+    return encoders, input_columns, label_columns, meta_columns
+
+
+def extract_transform_parameters_at_index(
+    transform: Transform,
+    index: int = 0,
+) -> Transform:
+    """Get a transform with parameters at the specified index.
+
+    Args:
+        transform: The original transform containing parameter lists
+        index: Index to extract parameters from (default 0)
+
+    Returns:
+        A new transform with single parameter values at the specified index
+    """
+    # Create a copy of the transform
+    new_transform = Transform(**transform.model_dump())
+
+    # Process each column and transformation
+    for column in new_transform.columns:
+        for transformation in column.transformations:
+            if transformation.params:
+                # Convert each parameter list to single value at index
+                new_params = {}
+                for param_name, param_value in transformation.params.items():
+                    if isinstance(param_value, list):
+                        new_params[param_name] = param_value[index]
+                    else:
+                        new_params[param_name] = param_value
+                transformation.params = new_params
+
+    return new_transform
+
+
+def expand_transform_parameter_combinations(
+    transform: Transform,
+) -> list[Transform]:
+    """Get all possible transforms by extracting parameters at each valid index.
+
+    For a transform with parameter lists, creates multiple new transforms, each containing
+    single parameter values from the corresponding indices of the parameter lists.
+
+    Args:
+        transform: The original transform containing parameter lists
+
+    Returns:
+        A list of transforms, each with single parameter values from sequential indices
+    """
+    # Find the length of parameter lists - we only need to check the first list we find
+    # since all lists must have the same length (enforced by pydantic validator)
+    max_length = 1
+    for column in transform.columns:
+        for transformation in column.transformations:
+            if transformation.params:
+                list_lengths = [len(v) for v in transformation.params.values() if isinstance(v, list) and len(v) > 1]
+                if list_lengths:
+                    max_length = list_lengths[0]  # All lists have same length due to validator
+                    break
+
+    # Generate a transform for each index
+    transforms = []
+    for i in range(max_length):
+        transforms.append(extract_transform_parameters_at_index(transform, i))
+
+    return transforms
+
+
+def expand_transform_list_combinations(
+    transform_list: list[Transform],
+) -> list[Transform]:
+    """Expands a list of transforms into all possible parameter combinations.
+
+    Takes a list of transforms where each transform may contain parameter lists,
+    and expands them into separate transforms with single parameter values.
+    For example, if a transform has parameters [0.1, 0.2] and [1, 2], this will
+    create two transforms: one with 0.1/1 and another with 0.2/2.
+
+    Args:
+        transform_list: A list of YamlTransform objects containing parameter lists
+            that need to be expanded into individual transforms.
+
+    Returns:
+        list[YamlTransform]: A flattened list of transforms where each transform
+            has single parameter values instead of parameter lists. The length of
+            the returned list will be the sum of the number of parameter combinations
+            for each input transform.
+    """
+    sub_transforms = []
+    for transform in transform_list:
+        sub_transforms.extend(expand_transform_parameter_combinations(transform))
+    return sub_transforms
+
+
+def generate_split_configs(config: ConfigDict) -> list[SplitConfigDict]:
+    """Generates all possible split configuration from a YAML config.
+
+    Takes a YAML configuration that may contain parameter lists and splits,
+    and generates all unique splits into separate data configurations.
+
+    For example, if the config has:
+    - Two transforms with parameters [0.1, 0.2], [0.3, 0.4]
+    - Two splits [0.7/0.3] and [0.8/0.2]
+    This will generate 2 configs, 2 for each split.
+        config_1:
+            transform: [[0.1, 0.2], [0.3, 0.4]]
+            split: [0.7, 0.3]
+
+        config_2:
+            transform: [[0.1, 0.2], [0.3, 0.4]]
+            split: [0.8, 0.2]
+
+    Args:
+        config: The source YAML configuration containing transforms with
+            parameter lists and multiple splits.
+
+    Returns:
+        list[SplitConfigDict]: A list of data configurations, where each
+            config has a list of parameters and one split configuration. The
+            length will be the product of the number of parameter combinations
+            and the number of splits.
+    """
+    if isinstance(config, dict) and not isinstance(config, ConfigDict):
+        raise TypeError("Input must be a ConfigDict object")
+
+    sub_splits = config.split
+    sub_configs = []
+    for split in sub_splits:
+        sub_configs.append(
+            SplitConfigDict(
+                global_params=config.global_params,
+                columns=config.columns,
+                transforms=config.transforms,
+                split=split,
+            ),
+        )
+    return sub_configs
+
+
+def generate_split_transform_configs(
+    config: SplitConfigDict,
+) -> list[SplitTransformDict]:
+    """Generates all the transform configuration for a given split.
+
+    Takes a YAML configuration that may contain a transform or a list of transform,
+    and generates all unique transform for a split into separate data configurations.
+
+    For example, if the config has:
+    - Two transforms with parameters [0.1, 0.2], [0.3, 0.4]
+    - A split [0.7, 0.3]
+    This will generate 2 configs, 2 for each split.
+        transform_config_1:
+            transform: [0.1, 0.2]
+            split: [0.7, 0.3]
+
+        transform_config_2:
+            transform: [0.3, 0.4]
+            split: [0.7, 0.3]
+
+    Args:
+        config: The source YAML configuration containing each
+            a split with transforms with parameters lists
+
+    Returns:
+        list[SplitTransformDict]: A list of data configurations, where each
+            config has a list of parameters and one split configuration. The
+            length will be the product of the number of parameter combinations
+            and the number of splits.
+    """
+    if isinstance(config, dict) and not isinstance(
+        config,
+        SplitConfigDict,
+    ):
+        raise TypeError("Input must be a list of SplitConfigDict")
+
+    sub_transforms = expand_transform_list_combinations(config.transforms)
+    split_transform_config: list[SplitTransformDict] = []
+    for transform in sub_transforms:
+        split_transform_config.append(
+            SplitTransformDict(
+                global_params=config.global_params,
+                columns=config.columns,
+                transforms=transform,
+                split=config.split,
+            ),
+        )
+    return split_transform_config
+
+
+def dump_yaml_list_into_files(
+    yaml_list: list[SplitConfigDict],
+    directory_path: str,
+    base_name: str,
+) -> None:
+    """Dumps YAML configurations to files with consistent formatting."""
+
+    class CleanDumper(yaml.SafeDumper):
+        """Simplified dumper maintaining key functionality."""
+
+        def ignore_aliases(self, _data: Any) -> bool:
+            return True  # Disable anchor/alias generation
+
+        def write_line_break(self, data: Any = None) -> None:
+            """Maintain root-level spacing."""
+            super().write_line_break(data)
+            if not self.indent:  # At root level
+                super().write_line_break()
+
+    # Register type handlers
+    CleanDumper.add_representer(type(None), lambda d, _: d.represent_scalar("tag:yaml.org,2002:null", ""))
+
+    CleanDumper.add_representer(
+        list,
+        lambda d, data: d.represent_sequence(
+            "tag:yaml.org,2002:seq",
+            data,
+            flow_style=isinstance(data[0], (list, dict)) if data else False,
+        ),
+    )
+
+    for i, yaml_dict in enumerate(yaml_list):
+        data = _clean_params(yaml_dict.model_dump(exclude_none=True))
+
+        with open(f"{directory_path}/{base_name}_{i}.yaml", "w") as f:
+            yaml.dump(
+                data,
+                f,
+                Dumper=CleanDumper,
+                sort_keys=False,
+                indent=2,
+                width=float("inf"),
+                default_flow_style=None,  # Let representers handle flow style
+            )
+
+
+def _clean_params(data: dict) -> dict:
+    """Recursive cleaner for empty parameters (replaces fix_params)."""
+    if isinstance(data, dict):
+        return {
+            k: (
+                # Handle special cases for encoder/transformations lists
+                [dict(e, params=e.get("params") or {}) for e in v]
+                if k in ("encoder", "transformations")
+                else _clean_params(v)
+            )
+            for k, v in data.items()
+        }
+    if isinstance(data, list):
+        return [_clean_params(item) for item in data]
+    return data
diff --git a/src/stimulus/data/interface/data_config_schema.py b/src/stimulus/data/interface/data_config_schema.py
new file mode 100644
index 0000000..9b3fcc0
--- /dev/null
+++ b/src/stimulus/data/interface/data_config_schema.py
@@ -0,0 +1,137 @@
+"""Module for defining the data config schema."""
+
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, field_validator
+
+
+class GlobalParams(BaseModel):
+    """Model for global parameters in YAML configuration."""
+
+    seed: int
+
+
+class ColumnsEncoder(BaseModel):
+    """Model for column encoder configuration."""
+
+    name: str
+    params: Optional[dict[str, Union[str, list[Any]]]]  # Allow both string and list values
+
+
+class Columns(BaseModel):
+    """Model for column configuration."""
+
+    column_name: str
+    column_type: str
+    data_type: str
+    encoder: list[ColumnsEncoder]
+
+
+class TransformColumnsTransformation(BaseModel):
+    """Model for column transformation configuration."""
+
+    name: str
+    params: Optional[dict[str, Union[list[Any], float]]]  # Allow both list and float values
+
+
+class TransformColumns(BaseModel):
+    """Model for transform columns configuration."""
+
+    column_name: str
+    transformations: list[TransformColumnsTransformation]
+
+
+class Transform(BaseModel):
+    """Model for transform configuration."""
+
+    transformation_name: str
+    columns: list[TransformColumns]
+
+    @field_validator("columns")
+    @classmethod
+    def validate_param_lists_across_columns(
+        cls,
+        columns: list[TransformColumns],
+    ) -> list[TransformColumns]:
+        """Validate that parameter lists across columns have consistent lengths.
+
+        Args:
+            columns: List of transform columns to validate
+
+        Returns:
+            The validated columns list
+        """
+        # Get all parameter list lengths across all columns and transformations
+        all_list_lengths: set[int] = set()
+
+        for column in columns:
+            for transformation in column.transformations:
+                if transformation.params and any(
+                    isinstance(param_value, list) and len(param_value) > 0
+                    for param_value in transformation.params.values()
+                ):
+                    all_list_lengths.update(
+                        len(param_value)
+                        for param_value in transformation.params.values()
+                        if isinstance(param_value, list) and len(param_value) > 0
+                    )
+
+        # Skip validation if no lists found
+        if not all_list_lengths:
+            return columns
+
+        # Check if all lists either have length 1, or all have the same length
+        all_list_lengths.discard(1)  # Remove length 1 as it's always valid
+        if len(all_list_lengths) > 1:  # Multiple different lengths found
+            raise ValueError(
+                "All parameter lists across columns must either contain one element or have the same length",
+            )
+
+        return columns
+
+
+class Split(BaseModel):
+    """Model for split configuration."""
+
+    split_method: str
+    params: dict[str, list[float]]  # More specific type for split parameters
+    split_input_columns: list[str]
+
+
+class ConfigDict(BaseModel):
+    """Model for main YAML configuration."""
+
+    global_params: GlobalParams
+    columns: list[Columns]
+    transforms: list[Transform]
+    split: list[Split]
+
+
+class SplitConfigDict(BaseModel):
+    """Model for sub-configuration generated from main config."""
+
+    global_params: GlobalParams
+    columns: list[Columns]
+    transforms: list[Transform]
+    split: Split
+
+
+class SplitTransformDict(BaseModel):
+    """Model for sub-configuration generated from main config."""
+
+    global_params: GlobalParams
+    columns: list[Columns]
+    transforms: Transform
+    split: Split
+
+
+class Schema(BaseModel):
+    """Model for validating YAML schema."""
+
+    config: ConfigDict
+
+
+class SplitSchema(BaseModel):
+    """Model for validating a Split YAML schema."""
+
+    config: ConfigDict
diff --git a/src/stimulus/data/loaders.py b/src/stimulus/data/loaders.py
deleted file mode 100644
index 096f10c..0000000
--- a/src/stimulus/data/loaders.py
+++ /dev/null
@@ -1,293 +0,0 @@
-"""Loaders serve as interfaces between the CSV master class and custom methods.
-
-Mainly, three types of custom methods are supported:
-- Encoders: methods for encoding data before it is fed into the model
-- Data transformers: methods for transforming data (i.e. augmenting, noising...)
-- Splitters: methods for splitting data into train, validation and test sets
-
-Loaders are built from an input config YAML file which format is described in the documentation, you can find an example here: tests/test_data/dna_experiment/dna_experiment_config_template.yaml
-"""
-
-import inspect
-import logging
-from typing import Any, Optional
-
-from stimulus.data.encoding import encoders
-from stimulus.data.splitters import splitters
-from stimulus.data.transform import data_transformation_generators
-from stimulus.utils import yaml_data
-
-
-class EncoderLoader:
-    """Class for loading encoders from a config file."""
-
-    def __init__(self, seed: Optional[float] = None) -> None:
-        """Initialize the encoder loader.
-
-        Args:
-            seed: Random seed for reproducibility
-        """
-        self.seed = seed
-
-    def initialize_column_encoders_from_config(
-        self,
-        column_config: yaml_data.YamlColumns,
-    ) -> None:
-        """Build the loader from a config dictionary.
-
-        Args:
-            column_config (yaml_data.YamlColumns): Configuration dictionary containing field names (column_name) and their encoder specifications.
-        """
-        for field in column_config:
-            encoder = self.get_encoder(field.encoder[0].name, field.encoder[0].params)
-            self.set_encoder_as_attribute(field.column_name, encoder)
-
-    def get_function_encode_all(self, field_name: str) -> Any:
-        """Gets the encoding function for a specific field.
-
-        Args:
-            field_name (str): The field name to get the encoder for
-
-        Returns:
-            Any: The encode_all function for the specified field
-        """
-        return getattr(self, field_name).encode_all
-
-    def get_encoder(
-        self,
-        encoder_name: str,
-        encoder_params: Optional[dict] = None,
-    ) -> Any:
-        """Gets an encoder object from the encoders module and initializes it with the given parameters.
-
-        Args:
-            encoder_name (str): The name of the encoder to get
-            encoder_params (dict): The parameters for the encoder
-
-        Returns:
-            Any: The encoder function for the specified field and parameters
-        """
-        try:
-            return getattr(encoders, encoder_name)(**encoder_params)
-        except AttributeError:
-            logging.exception(
-                f"Encoder '{encoder_name}' not found in the encoders module.",
-            )
-            logging.exception(
-                f"Available encoders: {[name for name, obj in encoders.__dict__.items() if isinstance(obj, type) and name not in ('ABC', 'Any')]}",
-            )
-            raise
-
-        except TypeError:
-            if encoder_params is None:
-                return getattr(encoders, encoder_name)()
-            logging.exception(
-                f"Encoder '{encoder_name}' has incorrect parameters: {encoder_params}",
-            )
-            logging.exception(
-                f"Expected parameters for '{encoder_name}': {inspect.signature(getattr(encoders, encoder_name))}",
-            )
-            raise
-
-    def set_encoder_as_attribute(
-        self,
-        field_name: str,
-        encoder: encoders.AbstractEncoder,
-    ) -> None:
-        """Sets the encoder as an attribute of the loader.
-
-        Args:
-            field_name (str): The name of the field to set the encoder for
-            encoder (encoders.AbstractEncoder): The encoder to set
-        """
-        setattr(self, field_name, encoder)
-
-
-class TransformLoader:
-    """Class for loading transformations from a config file."""
-
-    def __init__(self, seed: Optional[float] = None) -> None:
-        """Initialize the transform loader.
-
-        Args:
-            seed: Random seed for reproducibility
-        """
-        self.seed = seed
-
-    def get_data_transformer(
-        self,
-        transformation_name: str,
-        transformation_params: Optional[dict] = None,
-    ) -> Any:
-        """Gets a transformer object from the transformers module.
-
-        Args:
-            transformation_name (str): The name of the transformer to get
-            transformation_params (Optional[dict]): Parameters for the transformer
-
-        Returns:
-            Any: The transformer function for the specified transformation
-        """
-        try:
-            return getattr(data_transformation_generators, transformation_name)(
-                **transformation_params,
-            )
-        except AttributeError:
-            logging.exception(
-                f"Transformer '{transformation_name}' not found in the transformers module.",
-            )
-            logging.exception(
-                f"Available transformers: {[name for name, obj in data_transformation_generators.__dict__.items() if isinstance(obj, type) and name not in ('ABC', 'Any')]}",
-            )
-            raise
-
-        except TypeError:
-            if transformation_params is None:
-                return getattr(data_transformation_generators, transformation_name)()
-            logging.exception(
-                f"Transformer '{transformation_name}' has incorrect parameters: {transformation_params}",
-            )
-            logging.exception(
-                f"Expected parameters for '{transformation_name}': {inspect.signature(getattr(data_transformation_generators, transformation_name))}",
-            )
-            raise
-
-    def set_data_transformer_as_attribute(
-        self,
-        field_name: str,
-        data_transformer: Any,
-    ) -> None:
-        """Sets the data transformer as an attribute of the loader.
-
-        Args:
-            field_name (str): The name of the field to set the data transformer for
-            data_transformer (Any): The data transformer to set
-        """
-        # check if the field already exists, if it does not, initialize it to an empty dict
-        if not hasattr(self, field_name):
-            setattr(
-                self,
-                field_name,
-                {data_transformer.__class__.__name__: data_transformer},
-            )
-        else:
-            field_value = getattr(self, field_name)
-            field_value[data_transformer.__class__.__name__] = data_transformer
-
-    def initialize_column_data_transformers_from_config(
-        self,
-        transform_config: yaml_data.YamlTransform,
-    ) -> None:
-        """Build the loader from a config dictionary.
-
-        Args:
-            transform_config (yaml_data.YamlTransform): Configuration dictionary containing transforms configurations.
-
-        Example:
-            Given a YAML config like:
-            ```yaml
-            transforms:
-              transformation_name: noise
-              columns:
-                - column_name: age
-                  transformations:
-                    - name: GaussianNoise
-                      params:
-                        std: 0.1
-                - column_name: fare
-                  transformations:
-                    - name: GaussianNoise
-                      params:
-                        std: 0.1
-            ```
-
-            The loader will:
-            1. Iterate through each column (age, fare)
-            2. For each transformation in the column:
-               - Get the transformer (GaussianNoise) with its params (std=0.1)
-               - Set it as an attribute on the loader using the column name as key
-        """
-        for column in transform_config.columns:
-            col_name = column.column_name
-            for transform_spec in column.transformations:
-                transformer = self.get_data_transformer(
-                    transform_spec.name,
-                    transform_spec.params,
-                )
-                self.set_data_transformer_as_attribute(col_name, transformer)
-
-
-class SplitLoader:
-    """Class for loading splitters from a config file."""
-
-    def __init__(self, seed: Optional[float] = None) -> None:
-        """Initialize the split loader.
-
-        Args:
-            seed: Random seed for reproducibility
-        """
-        self.seed = seed
-
-    def get_function_split(self) -> Any:
-        """Gets the function for splitting the data.
-
-        Returns:
-            Any: The split function for the specified method
-
-        Raises:
-            AttributeError: If splitter hasn't been initialized using initialize_splitter_from_config()
-        """
-        if not hasattr(self, "split"):
-            # Raise a more specific error and chain it to the original AttributeError
-            raise AttributeError(
-                "Splitter not initialized. Please call initialize_splitter_from_config() or set_splitter_as_attribute() "
-                "before attempting to get split function.",
-            )
-        return self.split.get_split_indexes
-
-    def get_splitter(
-        self,
-        splitter_name: str,
-        splitter_params: Optional[dict] = None,
-    ) -> Any:
-        """Gets a splitter object from the splitters module.
-
-        Args:
-            splitter_name (str): The name of the splitter to get
-            splitter_params (Optional[dict]): Parameters for the splitter
-
-        Returns:
-            Any: The splitter function for the specified splitter
-        """
-        try:
-            return getattr(splitters, splitter_name)(**splitter_params)
-        except TypeError:
-            if splitter_params is None:
-                return getattr(splitters, splitter_name)()
-            logging.exception(
-                f"Splitter '{splitter_name}' has incorrect parameters: {splitter_params}",
-            )
-            logging.exception(
-                f"Expected parameters for '{splitter_name}': {inspect.signature(getattr(splitters, splitter_name))}",
-            )
-            raise
-
-    def set_splitter_as_attribute(self, splitter: Any) -> None:
-        """Sets the splitter as an attribute of the loader.
-
-        Args:
-            splitter (Any): The splitter to set
-        """
-        self.split = splitter
-
-    def initialize_splitter_from_config(
-        self,
-        split_config: yaml_data.YamlSplitConfigDict,
-    ) -> None:
-        """Build the loader from a config dictionary.
-
-        Args:
-            split_config (yaml_data.YamlSplitConfigDict): Configuration dictionary containing split configurations.
-        """
-        splitter = self.get_splitter(split_config.split_method, split_config.params)
-        self.set_splitter_as_attribute(splitter)
diff --git a/src/stimulus/data/splitters/__init__.py b/src/stimulus/data/splitting/__init__.py
similarity index 71%
rename from src/stimulus/data/splitters/__init__.py
rename to src/stimulus/data/splitting/__init__.py
index c790625..5f7c557 100644
--- a/src/stimulus/data/splitters/__init__.py
+++ b/src/stimulus/data/splitting/__init__.py
@@ -1,5 +1,5 @@
 """This package provides splitter classes for splitting data into train, validation, and test sets."""
 
-from stimulus.data.splitters.splitters import AbstractSplitter, RandomSplit
+from stimulus.data.splitting.splitters import AbstractSplitter, RandomSplit
 
 __all__ = ["AbstractSplitter", "RandomSplit"]
diff --git a/src/stimulus/data/splitters/splitters.py b/src/stimulus/data/splitting/splitters.py
similarity index 100%
rename from src/stimulus/data/splitters/splitters.py
rename to src/stimulus/data/splitting/splitters.py
diff --git a/src/stimulus/data/transform/__init__.py b/src/stimulus/data/transforming/__init__.py
similarity index 100%
rename from src/stimulus/data/transform/__init__.py
rename to src/stimulus/data/transforming/__init__.py
diff --git a/src/stimulus/data/transform/data_transformation_generators.py b/src/stimulus/data/transforming/transforms.py
similarity index 98%
rename from src/stimulus/data/transform/data_transformation_generators.py
rename to src/stimulus/data/transforming/transforms.py
index afc4d89..56cc14e 100644
--- a/src/stimulus/data/transform/data_transformation_generators.py
+++ b/src/stimulus/data/transforming/transforms.py
@@ -7,7 +7,7 @@ from typing import Any
 import numpy as np
 
 
-class AbstractDataTransformer(ABC):
+class AbstractTransform(ABC):
     """Abstract class for data transformers.
 
     Data transformers implement in_place or augmentation transformations.
@@ -64,7 +64,7 @@ class AbstractDataTransformer(ABC):
         raise NotImplementedError
 
 
-class AbstractNoiseGenerator(AbstractDataTransformer):
+class AbstractNoiseGenerator(AbstractTransform):
     """Abstract class for noise generators.
 
     All noise function should have the seed in it. This is because the multiprocessing of them could unset the seed.
@@ -76,7 +76,7 @@ class AbstractNoiseGenerator(AbstractDataTransformer):
         self.add_row = False
 
 
-class AbstractAugmentationGenerator(AbstractDataTransformer):
+class AbstractAugmentationGenerator(AbstractTransform):
     """Abstract class for augmentation generators.
 
     All augmentation function should have the seed in it. This is because the multiprocessing of them could unset the seed.
diff --git a/src/stimulus/learner/raytune_learner.py b/src/stimulus/learner/raytune_learner.py
index 9905f37..6682f4e 100644
--- a/src/stimulus/learner/raytune_learner.py
+++ b/src/stimulus/learner/raytune_learner.py
@@ -16,11 +16,9 @@ from safetensors.torch import save_model as safe_save_model
 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.utils.yaml_model_schema import RayTuneModel
 
 
@@ -36,10 +34,9 @@ class TuneWrapper:
     def __init__(
         self,
         model_config: RayTuneModel,
-        data_config: YamlSplitTransformDict,
         model_class: nn.Module,
-        data_path: str,
-        encoder_loader: EncoderLoader,
+        train_dataset: TorchDataset,
+        validation_dataset: TorchDataset,
         seed: int,
         ray_results_dir: Optional[str] = None,
         tune_run_name: Optional[str] = None,
@@ -85,11 +82,6 @@ class TuneWrapper:
             stop=model_config.tune.run_params.stop,
         )
 
-        # add the data path to the config
-        if not os.path.exists(data_path):
-            raise ValueError("Data path does not exist. Given path:" + data_path)
-        self.config["data_path"] = os.path.abspath(data_path)
-
         # Set up tune_run path
         if ray_results_dir is None:
             ray_results_dir = os.environ.get("HOME", "")
@@ -104,24 +96,21 @@ class TuneWrapper:
         )
         self.config["_debug"] = debug
         self.config["model"] = model_class
-        self.config["encoder_loader"] = encoder_loader
         self.config["ray_worker_seed"] = tune.randint(0, 1000)
 
         self.gpu_per_trial = model_config.tune.gpu_per_trial
         self.cpu_per_trial = model_config.tune.cpu_per_trial
 
         self.tuner = self.tuner_initialization(
-            data_config=data_config,
-            data_path=data_path,
-            encoder_loader=encoder_loader,
+            train_dataset=train_dataset,
+            validation_dataset=validation_dataset,
             autoscaler=autoscaler,
         )
 
     def tuner_initialization(
         self,
-        data_config: YamlSplitTransformDict,
-        data_path: str,
-        encoder_loader: EncoderLoader,
+        train_dataset: TorchDataset,
+        validation_dataset: TorchDataset,
         *,
         autoscaler: bool = False,
     ) -> tune.Tuner:
@@ -150,24 +139,9 @@ class TuneWrapper:
             f"PER_TRIAL resources ->  GPU: {self.gpu_per_trial} CPU: {self.cpu_per_trial}",
         )
 
-        # Pre-load and encode datasets once, then put them in Ray's object store
-
-        training = TorchDataset(
-            data_config=data_config,
-            csv_path=data_path,
-            encoder_loader=encoder_loader,
-            split=0,
-        )
-        validation = TorchDataset(
-            data_config=data_config,
-            csv_path=data_path,
-            encoder_loader=encoder_loader,
-            split=1,
-        )
-
         # log to debug the names of the columns and shapes of tensors for a batch of training
         # Log shapes of encoded tensors for first batch of training data
-        inputs, labels, meta = training[0:10]
+        inputs, labels, meta = train_dataset[0:10]
 
         logging.debug("Training data tensor shapes:")
         for field, tensor in inputs.items():
@@ -179,8 +153,8 @@ class TuneWrapper:
         for field, values in meta.items():
             logging.debug(f"Meta field '{field}' length: {len(values)}")
 
-        training_ref = ray.put(training)
-        validation_ref = ray.put(validation)
+        training_ref = ray.put(train_dataset)
+        validation_ref = ray.put(validation_dataset)
 
         self.config["_training_ref"] = training_ref
         self.config["_validation_ref"] = validation_ref
diff --git a/src/stimulus/typing/__init__.py b/src/stimulus/typing/__init__.py
index 7ae473d..76b0889 100644
--- a/src/stimulus/typing/__init__.py
+++ b/src/stimulus/typing/__init__.py
@@ -6,28 +6,39 @@ and likely not instantiated, as well as aliases for other types to use for typin
 The aliases from this module should be used for typing purposes only.
 """
 # ruff: noqa: F401
+# ruff: noqa: F811
 
-from typing import TypeAlias
+from typing import Any, TypeAlias, TypeVar
 
 # these imports mostly alias everything
 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.interface.data_config_schema import (
+    Columns,
+    ColumnsEncoder,
+    ConfigDict,
+    GlobalParams,
+    Schema,
+    Split,
+    SplitConfigDict,
+    SplitTransformDict,
+    Transform,
+    TransformColumns,
+    TransformColumnsTransformation,
+)
+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 (
     RayTuneMetrics,
     RayTuneOptimizer,
@@ -35,56 +46,44 @@ from stimulus.learner.raytune_parser import (
     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.utils.yaml_model_schema import (
     CustomTunableParameter,
     Data,
     Loss,
     Model,
+    RayConfigLoader,
     RayTuneModel,
     RunParams,
     Scheduler,
     TunableParameter,
     Tune,
     TuneParams,
-    YamlRayConfigLoader,
 )
 
-# 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
+# data/interface/data_config_schema.py
 
 YamlData: TypeAlias = (
-    YamlColumns
-    | YamlColumnsEncoder
-    | YamlConfigDict
-    | YamlGlobalParams
-    | YamlSchema
-    | YamlSplit
-    | YamlSplitConfigDict
-    | YamlTransform
-    | YamlTransformColumns
-    | YamlTransformColumnsTransformation
+    Columns
+    | ColumnsEncoder
+    | ConfigDict
+    | GlobalParams
+    | Schema
+    | Split
+    | SplitConfigDict
+    | Transform
+    | TransformColumns
+    | TransformColumnsTransformation
+    | SplitTransformDict
 )
+
+# Replace these problematic imports
+# from stimulus.learner.raytune_learner import CheckpointDict, TuneModel, TuneWrapper
+
+# Replace with type aliases if needed
+CheckpointDict = dict[str, Any]
+TuneModel = TypeVar("TuneModel")
+TuneWrapper = TypeVar("TuneWrapper")
diff --git a/src/stimulus/utils/launch_utils.py b/src/stimulus/utils/model_file_interface.py
similarity index 100%
rename from src/stimulus/utils/launch_utils.py
rename to src/stimulus/utils/model_file_interface.py
diff --git a/src/stimulus/utils/yaml_data.py b/src/stimulus/utils/yaml_data.py
deleted file mode 100644
index 1229609..0000000
--- a/src/stimulus/utils/yaml_data.py
+++ /dev/null
@@ -1,456 +0,0 @@
-"""Utility module for handling YAML configuration files and their validation."""
-
-from typing import Any, Optional, Union
-
-import yaml
-from pydantic import BaseModel, ValidationError, field_validator
-
-
-class YamlGlobalParams(BaseModel):
-    """Model for global parameters in YAML configuration."""
-
-    seed: int
-
-
-class YamlColumnsEncoder(BaseModel):
-    """Model for column encoder configuration."""
-
-    name: str
-    params: Optional[dict[str, Union[str, list[Any]]]]  # Allow both string and list values
-
-
-class YamlColumns(BaseModel):
-    """Model for column configuration."""
-
-    column_name: str
-    column_type: str
-    data_type: str
-    encoder: list[YamlColumnsEncoder]
-
-
-class YamlTransformColumnsTransformation(BaseModel):
-    """Model for column transformation configuration."""
-
-    name: str
-    params: Optional[dict[str, Union[list[Any], float]]]  # Allow both list and float values
-
-
-class YamlTransformColumns(BaseModel):
-    """Model for transform columns configuration."""
-
-    column_name: str
-    transformations: list[YamlTransformColumnsTransformation]
-
-
-class YamlTransform(BaseModel):
-    """Model for transform configuration."""
-
-    transformation_name: str
-    columns: list[YamlTransformColumns]
-
-    @field_validator("columns")
-    @classmethod
-    def validate_param_lists_across_columns(
-        cls,
-        columns: list[YamlTransformColumns],
-    ) -> list[YamlTransformColumns]:
-        """Validate that parameter lists across columns have consistent lengths.
-
-        Args:
-            columns: List of transform columns to validate
-
-        Returns:
-            The validated columns list
-        """
-        # Get all parameter list lengths across all columns and transformations
-        all_list_lengths: set[int] = set()
-
-        for column in columns:
-            for transformation in column.transformations:
-                if transformation.params and any(
-                    isinstance(param_value, list) and len(param_value) > 0
-                    for param_value in transformation.params.values()
-                ):
-                    all_list_lengths.update(
-                        len(param_value)
-                        for param_value in transformation.params.values()
-                        if isinstance(param_value, list) and len(param_value) > 0
-                    )
-
-        # Skip validation if no lists found
-        if not all_list_lengths:
-            return columns
-
-        # Check if all lists either have length 1, or all have the same length
-        all_list_lengths.discard(1)  # Remove length 1 as it's always valid
-        if len(all_list_lengths) > 1:  # Multiple different lengths found
-            raise ValueError(
-                "All parameter lists across columns must either contain one element or have the same length",
-            )
-
-        return columns
-
-
-class YamlSplit(BaseModel):
-    """Model for split configuration."""
-
-    split_method: str
-    params: dict[str, list[float]]  # More specific type for split parameters
-    split_input_columns: list[str]
-
-
-class YamlConfigDict(BaseModel):
-    """Model for main YAML configuration."""
-
-    global_params: YamlGlobalParams
-    columns: list[YamlColumns]
-    transforms: list[YamlTransform]
-    split: list[YamlSplit]
-
-
-# TODO: Rename this class to SplitConfigDict
-class YamlSplitConfigDict(BaseModel):
-    """Model for sub-configuration generated from main config."""
-
-    global_params: YamlGlobalParams
-    columns: list[YamlColumns]
-    transforms: list[YamlTransform]
-    split: YamlSplit
-
-
-class YamlSplitTransformDict(BaseModel):
-    """Model for sub-configuration generated from main config."""
-
-    global_params: YamlGlobalParams
-    columns: list[YamlColumns]
-    transforms: YamlTransform
-    split: YamlSplit
-
-
-class YamlSchema(BaseModel):
-    """Model for validating YAML schema."""
-
-    yaml_conf: YamlConfigDict
-
-
-class YamlSplitSchema(BaseModel):
-    """Model for validating a Split YAML schema."""
-
-    yaml_conf: YamlSplitConfigDict
-
-
-def extract_transform_parameters_at_index(
-    transform: YamlTransform,
-    index: int = 0,
-) -> YamlTransform:
-    """Get a transform with parameters at the specified index.
-
-    Args:
-        transform: The original transform containing parameter lists
-        index: Index to extract parameters from (default 0)
-
-    Returns:
-        A new transform with single parameter values at the specified index
-    """
-    # Create a copy of the transform
-    new_transform = YamlTransform(**transform.model_dump())
-
-    # Process each column and transformation
-    for column in new_transform.columns:
-        for transformation in column.transformations:
-            if transformation.params:
-                # Convert each parameter list to single value at index
-                new_params = {}
-                for param_name, param_value in transformation.params.items():
-                    if isinstance(param_value, list):
-                        new_params[param_name] = param_value[index]
-                    else:
-                        new_params[param_name] = param_value
-                transformation.params = new_params
-
-    return new_transform
-
-
-def expand_transform_parameter_combinations(
-    transform: YamlTransform,
-) -> list[YamlTransform]:
-    """Get all possible transforms by extracting parameters at each valid index.
-
-    For a transform with parameter lists, creates multiple new transforms, each containing
-    single parameter values from the corresponding indices of the parameter lists.
-
-    Args:
-        transform: The original transform containing parameter lists
-
-    Returns:
-        A list of transforms, each with single parameter values from sequential indices
-    """
-    # Find the length of parameter lists - we only need to check the first list we find
-    # since all lists must have the same length (enforced by pydantic validator)
-    max_length = 1
-    for column in transform.columns:
-        for transformation in column.transformations:
-            if transformation.params:
-                list_lengths = [len(v) for v in transformation.params.values() if isinstance(v, list) and len(v) > 1]
-                if list_lengths:
-                    max_length = list_lengths[0]  # All lists have same length due to validator
-                    break
-
-    # Generate a transform for each index
-    transforms = []
-    for i in range(max_length):
-        transforms.append(extract_transform_parameters_at_index(transform, i))
-
-    return transforms
-
-
-def expand_transform_list_combinations(
-    transform_list: list[YamlTransform],
-) -> list[YamlTransform]:
-    """Expands a list of transforms into all possible parameter combinations.
-
-    Takes a list of transforms where each transform may contain parameter lists,
-    and expands them into separate transforms with single parameter values.
-    For example, if a transform has parameters [0.1, 0.2] and [1, 2], this will
-    create two transforms: one with 0.1/1 and another with 0.2/2.
-
-    Args:
-        transform_list: A list of YamlTransform objects containing parameter lists
-            that need to be expanded into individual transforms.
-
-    Returns:
-        list[YamlTransform]: A flattened list of transforms where each transform
-            has single parameter values instead of parameter lists. The length of
-            the returned list will be the sum of the number of parameter combinations
-            for each input transform.
-    """
-    sub_transforms = []
-    for transform in transform_list:
-        sub_transforms.extend(expand_transform_parameter_combinations(transform))
-    return sub_transforms
-
-
-def generate_split_configs(yaml_config: YamlConfigDict) -> list[YamlSplitConfigDict]:
-    """Generates all possible split configuration from a YAML config.
-
-    Takes a YAML configuration that may contain parameter lists and splits,
-    and generates all unique splits into separate data configurations.
-
-    For example, if the config has:
-    - Two transforms with parameters [0.1, 0.2], [0.3, 0.4]
-    - Two splits [0.7/0.3] and [0.8/0.2]
-    This will generate 2 configs, 2 for each split.
-        config_1:
-            transform: [[0.1, 0.2], [0.3, 0.4]]
-            split: [0.7, 0.3]
-
-        config_2:
-            transform: [[0.1, 0.2], [0.3, 0.4]]
-            split: [0.8, 0.2]
-
-    Args:
-        yaml_config: The source YAML configuration containing transforms with
-            parameter lists and multiple splits.
-
-    Returns:
-        list[YamlSubConfigDict]: A list of data configurations, where each
-            config has a list of parameters and one split configuration. The
-            length will be the product of the number of parameter combinations
-            and the number of splits.
-    """
-    if isinstance(yaml_config, dict) and not isinstance(yaml_config, YamlConfigDict):
-        raise TypeError("Input must be a YamlConfigDict object")
-
-    sub_splits = yaml_config.split
-    sub_configs = []
-    for split in sub_splits:
-        sub_configs.append(
-            YamlSplitConfigDict(
-                global_params=yaml_config.global_params,
-                columns=yaml_config.columns,
-                transforms=yaml_config.transforms,
-                split=split,
-            ),
-        )
-    return sub_configs
-
-
-def generate_split_transform_configs(
-    yaml_config: YamlSplitConfigDict,
-) -> list[YamlSplitTransformDict]:
-    """Generates all the transform configuration for a given split.
-
-    Takes a YAML configuration that may contain a transform or a list of transform,
-    and generates all unique transform for a split into separate data configurations.
-
-    For example, if the config has:
-    - Two transforms with parameters [0.1, 0.2], [0.3, 0.4]
-    - A split [0.7, 0.3]
-    This will generate 2 configs, 2 for each split.
-        transform_config_1:
-            transform: [0.1, 0.2]
-            split: [0.7, 0.3]
-
-        transform_config_2:
-            transform: [0.3, 0.4]
-            split: [0.7, 0.3]
-
-    Args:
-        yaml_config: The source YAML configuration containing each
-            a split with transforms with parameters lists
-
-    Returns:
-        list[YamlSubConfigTransformDict]: A list of data configurations, where each
-            config has a list of parameters and one split configuration. The
-            length will be the product of the number of parameter combinations
-            and the number of splits.
-    """
-    if isinstance(yaml_config, dict) and not isinstance(
-        yaml_config,
-        YamlSplitConfigDict,
-    ):
-        raise TypeError("Input must be a list of YamlSubConfigDict")
-
-    sub_transforms = expand_transform_list_combinations(yaml_config.transforms)
-    split_transform_config: list[YamlSplitTransformDict] = []
-    for transform in sub_transforms:
-        split_transform_config.append(
-            YamlSplitTransformDict(
-                global_params=yaml_config.global_params,
-                columns=yaml_config.columns,
-                transforms=transform,
-                split=yaml_config.split,
-            ),
-        )
-    return split_transform_config
-
-
-def dump_yaml_list_into_files(
-    yaml_list: list[YamlSplitConfigDict],
-    directory_path: str,
-    base_name: str,
-) -> None:
-    """Dumps a list of YAML configurations into separate files with custom formatting."""
-    # Create a new class attribute rather than assigning to the method
-    # Remove this line since we'll add ignore_aliases to CustomDumper instead
-
-    def represent_none(dumper: yaml.Dumper, _: Any) -> yaml.Node:
-        """Custom representer to format None values as empty strings in YAML output."""
-        return dumper.represent_scalar("tag:yaml.org,2002:null", "")
-
-    def custom_representer(dumper: yaml.Dumper, data: Any) -> yaml.Node:
-        """Custom representer to handle different types of lists with appropriate formatting."""
-        if isinstance(data, list):
-            if len(data) == 0:
-                return dumper.represent_scalar("tag:yaml.org,2002:null", "")
-            if isinstance(data[0], dict):
-                return dumper.represent_sequence(
-                    "tag:yaml.org,2002:seq",
-                    data,
-                    flow_style=False,
-                )
-            if isinstance(data[0], list):
-                return dumper.represent_sequence(
-                    "tag:yaml.org,2002:seq",
-                    data,
-                    flow_style=True,
-                )
-        return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True)
-
-    class CustomDumper(yaml.Dumper):
-        """Custom YAML dumper that adds extra formatting controls."""
-
-        def ignore_aliases(self, _data: Any) -> bool:
-            """Ignore aliases in the YAML output."""
-            return True
-
-        def write_line_break(self, _data: Any = None) -> None:
-            """Add extra newline after root-level elements."""
-            super().write_line_break(_data)
-            if len(self.indents) <= 1:  # At root level
-                super().write_line_break(_data)
-
-        def increase_indent(  # type: ignore[override]
-            self,
-            *,
-            flow: bool = False,
-            indentless: bool = False,
-        ) -> None:
-            """Ensure consistent indentation by preventing indentless sequences."""
-            return super().increase_indent(
-                flow=flow,
-                indentless=indentless,
-            )  # Force indentless to False for better formatting
-
-    # Register the custom representers with our dumper
-    yaml.add_representer(type(None), represent_none, Dumper=CustomDumper)
-    yaml.add_representer(list, custom_representer, Dumper=CustomDumper)
-
-    for i, yaml_dict in enumerate(yaml_list):
-        dict_data = yaml_dict.model_dump(exclude_none=True)
-
-        def fix_params(input_dict: dict[str, Any]) -> dict[str, Any]:
-            """Recursively process dictionary to properly handle params fields."""
-            if isinstance(input_dict, dict):
-                processed_dict: dict[str, Any] = {}
-                for key, value in input_dict.items():
-                    if key == "encoder" and isinstance(value, list):
-                        processed_dict[key] = []
-                        for encoder in value:
-                            processed_encoder = dict(encoder)
-                            if "params" not in processed_encoder or not processed_encoder["params"]:
-                                processed_encoder["params"] = {}
-                            processed_dict[key].append(processed_encoder)
-                    elif key == "transformations" and isinstance(value, list):
-                        processed_dict[key] = []
-                        for transformation in value:
-                            processed_transformation = dict(transformation)
-                            if "params" not in processed_transformation or not processed_transformation["params"]:
-                                processed_transformation["params"] = {}
-                            processed_dict[key].append(processed_transformation)
-                    elif isinstance(value, dict):
-                        processed_dict[key] = fix_params(value)
-                    elif isinstance(value, list):
-                        processed_dict[key] = [
-                            fix_params(list_item) if isinstance(list_item, dict) else list_item for list_item in value
-                        ]
-                    else:
-                        processed_dict[key] = value
-                return processed_dict
-            return input_dict
-
-        dict_data = fix_params(dict_data)
-
-        with open(f"{directory_path}/{base_name}_{i}.yaml", "w") as f:
-            yaml.dump(
-                dict_data,
-                f,
-                Dumper=CustomDumper,
-                sort_keys=False,
-                default_flow_style=False,
-                indent=2,
-                width=float("inf"),  # Prevent line wrapping
-            )
-
-
-def check_yaml_schema(config_yaml: YamlConfigDict) -> str:
-    """Validate YAML configuration fields have correct types.
-
-    If the children field is specific to a parent, the children fields class is hosted in the parent fields class.
-    If any field in not the right type, the function prints an error message explaining the problem and exits the python code.
-
-    Args:
-        config_yaml: The YamlConfigDict containing the fields of the yaml configuration file
-
-    Returns:
-        str: Empty string if validation succeeds
-
-    Raises:
-        ValueError: If validation fails
-    """
-    try:
-        YamlSchema(yaml_conf=config_yaml)
-    except ValidationError as e:
-        # Use logging instead of print for error handling
-        raise ValueError("Wrong type on a field, see the pydantic report above") from e
-    return ""
diff --git a/src/stimulus/utils/yaml_model_schema.py b/src/stimulus/utils/yaml_model_schema.py
index 7fb25eb..6b554d8 100644
--- a/src/stimulus/utils/yaml_model_schema.py
+++ b/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


Test patch:
diff --git a/tests/data/interface/__init__.py b/tests/data/interface/__init__.py
new file mode 100644
index 0000000..3daf13b
--- /dev/null
+++ b/tests/data/interface/__init__.py
@@ -0,0 +1,1 @@
+"""test interface for data handling."""
diff --git a/tests/data/interface/test_data_config_parser.py b/tests/data/interface/test_data_config_parser.py
new file mode 100644
index 0000000..dda6046
--- /dev/null
+++ b/tests/data/interface/test_data_config_parser.py
@@ -0,0 +1,163 @@
+"""Test the data config parser."""
+
+import pytest
+import yaml
+
+import stimulus.data.encoding.encoders as encoders_module  # TODO: should be imported from stimulus.typing instead.
+from stimulus.data.interface.data_config_parser import (
+    create_encoders,
+    create_splitter,
+    create_transforms,
+    expand_transform_parameter_combinations,
+    generate_split_configs,
+    generate_split_transform_configs,
+    parse_split_transform_config,
+)
+from stimulus.data.interface.data_config_schema import (
+    ConfigDict,
+    SplitConfigDict,
+    SplitTransformDict,
+)
+
+
+@pytest.fixture
+def split_config_path() -> str:
+    """Get path to test config file.
+
+    Returns:
+        str: Path to test config file
+    """
+    return "tests/test_data/titanic/titanic_unique_split.yaml"
+
+
+@pytest.fixture
+def full_config_path() -> str:
+    """Get path to test config file.
+
+    Returns:
+        str: Path to test config file
+    """
+    return "tests/test_data/titanic/titanic.yaml"
+
+
+@pytest.fixture
+def transform_config_path() -> str:
+    """Get path to test config file.
+
+    Returns:
+        str: Path to test config file
+    """
+    return "tests/test_data/titanic/titanic_unique_transform.yaml"
+
+
+@pytest.fixture
+def load_transform_config(transform_config_path: str) -> dict:
+    """Load the YAML config file.
+
+    Args:
+        config_path (str): Path to the YAML config file
+
+    Returns:
+        dict: The loaded YAML config
+    """
+    with open(transform_config_path) as f:
+        return SplitTransformDict(**yaml.safe_load(f))
+
+
+@pytest.fixture
+def load_split_config(split_config_path: str) -> dict:
+    """Load the YAML config file.
+
+    Args:
+        config_path (str): Path to the YAML config file
+
+    Returns:
+        dict: The loaded YAML config
+    """
+    with open(split_config_path) as f:
+        return SplitConfigDict(**yaml.safe_load(f))
+
+
+@pytest.fixture
+def load_full_config(full_config_path: str) -> dict:
+    """Load the YAML config file.
+
+    Args:
+        config_path (str): Path to the YAML config file
+
+    Returns:
+        dict: The loaded YAML config
+    """
+    with open(full_config_path) as f:
+        return ConfigDict(**yaml.safe_load(f))
+
+
+def test_create_encoders(load_full_config: ConfigDict) -> None:
+    """Test encoder creation from config."""
+    config = load_full_config
+
+    encoders = create_encoders(config.columns)
+
+    # Test encoder types
+    assert isinstance(encoders["sex"], encoders_module.StrClassificationEncoder)
+    assert isinstance(encoders["age"], encoders_module.NumericEncoder)
+    assert len(encoders) == len(config.columns)
+
+
+def test_create_transforms(load_split_config: SplitTransformDict) -> None:
+    """Test transform pipeline creation."""
+    transforms = create_transforms(load_split_config.transforms)
+
+    # Verify transform structure
+    age_transforms = transforms["age"]
+    assert len(age_transforms) == 1
+    assert age_transforms[0].__class__.__name__ == "GaussianNoise"
+    assert "std" in age_transforms[0].__dict__
+
+
+def test_create_splitter(load_full_config: ConfigDict) -> None:
+    """Test splitter creation with parameters."""
+    splitter = create_splitter(load_full_config.split[0])
+
+    assert splitter.__class__.__name__ == "RandomSplit"
+    assert splitter.split == [0.7, 0.15, 0.15]
+
+
+def test_parse_split_transform_config(load_transform_config: SplitTransformDict) -> None:
+    """Test full config parsing workflow."""
+    encoders, input_columns, label_columns, meta_columns = parse_split_transform_config(load_transform_config)
+
+    assert len(encoders) == 9
+    assert len(input_columns) == 7
+    assert len(label_columns) == 1
+    assert len(meta_columns) == 1
+
+
+def test_expand_transform_parameter_combinations(load_split_config: SplitTransformDict) -> None:
+    """Test parameter expansion logic."""
+    expanded = expand_transform_parameter_combinations(load_split_config.transforms[0])
+    assert len(expanded) == 3  # Because we have 3 std values in the YAML
+
+    # Verify parameter unpacking
+    first_params = expanded[0].columns[0].transformations[0].params
+    assert first_params["std"] == 0.1
+
+
+def test_generate_split_configs(load_full_config: ConfigDict) -> None:
+    """Test split configuration generation."""
+    split_configs = generate_split_configs(load_full_config)
+    assert len(split_configs) == 1  # Only one split in the test YAML
+
+    cfg = split_configs[0]
+    assert isinstance(cfg, SplitConfigDict)
+    assert cfg.split.split_method == "RandomSplit"
+
+
+def test_generate_split_transform_configs(load_full_config: ConfigDict) -> None:
+    """Test combined split/transform config generation."""
+    split_configs = generate_split_configs(load_full_config)
+    combined_configs = generate_split_transform_configs(split_configs[0])
+
+    assert len(combined_configs) == 3  # 3 parameter combinations
+    assert all(isinstance(c, SplitTransformDict) for c in combined_configs)
+    assert combined_configs[0].transforms.columns[0].transformations[0].params["std"] == 0.1
diff --git a/tests/data/test_data_handlers.py b/tests/data/test_data_handlers.py
index cf6dae1..de377e7 100644
--- a/tests/data/test_data_handlers.py
+++ b/tests/data/test_data_handlers.py
@@ -1,27 +1,18 @@
 """Tests for CSV data loading and processing functionality."""
 
+from typing import Any
+
 import pytest
-import yaml
+import torch
 
-from stimulus.data import loaders
 from stimulus.data.data_handlers import (
     DatasetLoader,
-    DatasetManager,
     DatasetProcessor,
-    EncodeManager,
-    SplitManager,
-    TransformManager,
-)
-from stimulus.utils.yaml_data import (
-    YamlConfigDict,
-    YamlSplitConfigDict,
-    YamlSplitTransformDict,
-    YamlTransform,
-    YamlTransformColumns,
-    YamlTransformColumnsTransformation,
-    generate_split_configs,
-    generate_split_transform_configs,
+    TorchDataset,
 )
+from stimulus.data.encoding import encoders as encoders_module
+from stimulus.data.splitting import splitters as splitters_module
+from stimulus.data.transforming import transforms as transforms_module
 
 
 # Fixtures
@@ -36,340 +27,213 @@ def titanic_csv_path() -> str:
     return "tests/test_data/titanic/titanic_stimulus.csv"
 
 
+# Updated component fixtures
 @pytest.fixture
-def config_path() -> str:
-    """Get path to test config file.
-
-    Returns:
-        str: Path to test config file
-    """
-    return "tests/test_data/titanic/titanic.yaml"
-
-
-@pytest.fixture
-def base_config(config_path: str) -> YamlConfigDict:
-    """Load base configuration from YAML file.
-
-    Args:
-        config_path: Path to config file
-
-    Returns:
-        YamlConfigDict: Loaded configuration
-    """
-    with open(config_path) as f:
-        return YamlConfigDict(**yaml.safe_load(f))
-
-
-@pytest.fixture
-def generate_sub_configs(base_config: YamlConfigDict) -> list[YamlConfigDict]:
-    """Generate all possible configurations from base config.
-
-    Args:
-        base_config: Base configuration to generate from
-
-    Returns:
-        list[YamlConfigDict]: List of generated configurations
-    """
-    split_configs: list[YamlSplitConfigDict] = generate_split_configs(base_config)
-    split_transform_list: list[YamlSplitTransformDict] = []
-    for split in split_configs:
-        split_transform_list.extend(generate_split_transform_configs(split))
-    return split_transform_list
-
-
-@pytest.fixture
-def dump_single_split_config_to_disk() -> YamlSplitTransformDict:
-    """Get path for dumping single split config.
-
-    Returns:
-        str: Path to dump config file
-    """
-    config_dict: YamlSplitTransformDict
-    path: str = "tests/test_data/titanic/titanic_sub_config.yaml"
-    with open(path) as f:
-        config_dict = YamlSplitTransformDict(**yaml.safe_load(f))
-    return config_dict
-
-
-# Loader fixtures
-@pytest.fixture
-def encoder_loader(generate_sub_configs: list[YamlConfigDict]) -> loaders.EncoderLoader:
-    """Create encoder loader with initialized encoders.
-
-    Args:
-        generate_sub_configs: List of configurations
-
-    Returns:
-        experiments.EncoderLoader: Initialized encoder loader
-    """
-    loader = loaders.EncoderLoader()
-    loader.initialize_column_encoders_from_config(generate_sub_configs[0].columns)
-    return loader
+def dummy_encoders() -> dict[str, Any]:
+    """Create simple encoders for test columns."""
+    return {
+        "age": encoders_module.NumericEncoder(dtype=torch.float32),
+        "fare": encoders_module.NumericEncoder(dtype=torch.float32),
+        "survived": encoders_module.NumericEncoder(dtype=torch.int64),
+    }
 
 
 @pytest.fixture
-def transform_loader(
-    generate_sub_configs: list[YamlConfigDict],
-) -> loaders.TransformLoader:
-    """Create transform loader with initialized transformers.
-
-    Args:
-        generate_sub_configs: List of configurations
-
-    Returns:
-        experiments.TransformLoader: Initialized transform loader
-    """
-    loader = loaders.TransformLoader()
-    loader.initialize_column_data_transformers_from_config(
-        generate_sub_configs[0].transforms,
-    )
-    return loader
-
-
-@pytest.fixture
-def split_loader(generate_sub_configs: list[YamlConfigDict]) -> loaders.SplitLoader:
-    """Create split loader with initialized splitter.
-
-    Args:
-        generate_sub_configs: List of configurations
-
-    Returns:
-        experiments.SplitLoader: Initialized split loader
-    """
-    loader = loaders.SplitLoader()
-    loader.initialize_splitter_from_config(generate_sub_configs[0].split)
-    return loader
-
-
-# Test DatasetManager
-def test_dataset_manager_init(
-    dump_single_split_config_to_disk: YamlSplitTransformDict,
-) -> None:
-    """Test initialization of DatasetManager."""
-    manager = DatasetManager(dump_single_split_config_to_disk)
-    assert hasattr(manager, "config")
-    assert hasattr(manager, "column_categories")
-
-
-def test_dataset_manager_organize_columns(
-    dump_single_split_config_to_disk: str,
-) -> None:
-    """Test column organization by type."""
-    manager = DatasetManager(dump_single_split_config_to_disk)
-    categories = manager.categorize_columns_by_type()
-
-    assert "pclass" in categories["input"]
-    assert "sex" in categories["input"]
-    assert "age" in categories["input"]
-    assert "survived" in categories["label"]
-    assert "passenger_id" in categories["meta"]
-
-
-def test_dataset_manager_organize_transforms(
-    dump_single_split_config_to_disk: str,
-) -> None:
-    """Test transform organization."""
-    manager = DatasetManager(dump_single_split_config_to_disk)
-    categories = manager.categorize_columns_by_type()
-
-    assert len(categories) == 3
-    assert all(key in categories for key in ["input", "label", "meta"])
-
-
-def test_dataset_manager_get_transform_logic(
-    dump_single_split_config_to_disk: str,
-) -> None:
-    """Test getting transform logic from config."""
-    manager = DatasetManager(dump_single_split_config_to_disk)
-    transform_logic = manager.get_transform_logic()
-    assert transform_logic["transformation_name"] == "noise"
-    assert len(transform_logic["transformations"]) == 2
-
-
-# Test EncodeManager
-def test_encode_manager_init() -> None:
-    """Test initialization of EncodeManager."""
-    encoder_loader = loaders.EncoderLoader()
-    manager = EncodeManager(encoder_loader)
-    assert hasattr(manager, "encoder_loader")
-
-
-def test_encode_manager_initialize_encoders() -> None:
-    """Test encoder initialization."""
-    encoder_loader = loaders.EncoderLoader()
-    manager = EncodeManager(encoder_loader)
-    assert hasattr(manager, "encoder_loader")
-
-
-def test_encode_manager_encode_numeric() -> None:
-    """Test numeric encoding."""
-    encoder_loader = loaders.EncoderLoader()
-    intencoder = encoder_loader.get_encoder("NumericEncoder")
-    encoder_loader.set_encoder_as_attribute("test_col", intencoder)
-    manager = EncodeManager(encoder_loader)
-    data = [1, 2, 3]
-    encoded = manager.encode_column("test_col", data)
-    assert encoded is not None
-
-
-# Test TransformManager
-def test_transform_manager_init() -> None:
-    """Test initialization of TransformManager."""
-    transform_loader = loaders.TransformLoader()
-    manager = TransformManager(transform_loader)
-    assert hasattr(manager, "transform_loader")
-
-
-def test_transform_manager_initialize_transforms() -> None:
-    """Test transform initialization."""
-    transform_loader = loaders.TransformLoader()
-    manager = TransformManager(transform_loader)
-    assert hasattr(manager, "transform_loader")
-
-
-def test_transform_manager_transform_column() -> None:
-    """Test column transformation."""
-    transform_loader = loaders.TransformLoader()
-    dummy_config = YamlTransform(
-        transformation_name="GaussianNoise",
-        columns=[
-            YamlTransformColumns(
-                column_name="test_col",
-                transformations=[
-                    YamlTransformColumnsTransformation(
-                        name="GaussianNoise",
-                        params={"std": 0.1},
-                    ),
-                ],
-            ),
+def dummy_transforms() -> dict[str, list]:
+    """Create test transforms."""
+    return {
+        "age": [
+            transforms_module.GaussianNoise(std=0.1),
+            transforms_module.GaussianNoise(std=0.2),
+            transforms_module.GaussianNoise(std=0.3),
         ],
-    )
-    transform_loader.initialize_column_data_transformers_from_config(dummy_config)
-    manager = TransformManager(transform_loader)
-    data = [1, 2, 3]
-    transformed, added_row = manager.transform_column("test_col", "GaussianNoise", data)
-    assert len(transformed) == len(data)
-    assert added_row is False
-
-
-# Test SplitManager
-def test_split_manager_init(split_loader: loaders.SplitLoader) -> None:
-    """Test initialization of SplitManager."""
-    manager = SplitManager(split_loader)
-    assert hasattr(manager, "split_loader")
+        "fare": [transforms_module.GaussianNoise(std=0.1)],
+    }
 
 
-def test_split_manager_initialize_splits(split_loader: loaders.SplitLoader) -> None:
-    """Test split initialization."""
-    manager = SplitManager(split_loader)
-    assert hasattr(manager, "split_loader")
-
-
-def test_split_manager_apply_split(split_loader: loaders.SplitLoader) -> None:
-    """Test applying splits to data."""
-    manager = SplitManager(split_loader)
-    data = {"col": range(100)}
-    split_indices = manager.get_split_indices(data)
-    assert len(split_indices) == 3
-    assert len(split_indices[0]) == 70
-    assert len(split_indices[1]) == 15
-    assert len(split_indices[2]) == 15
+@pytest.fixture
+def dummy_splitter() -> splitters_module.AbstractSplitter:
+    """Create test splitter."""
+    return splitters_module.RandomSplit(split=[0.7, 0.2, 0.1])
 
 
 # Test DatasetProcessor
 def test_dataset_processor_init(
-    dump_single_split_config_to_disk: str,
     titanic_csv_path: str,
+    dummy_transforms: dict,
+    dummy_splitter: Any,
 ) -> None:
     """Test initialization of DatasetProcessor."""
     processor = DatasetProcessor(
-        data_config=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
+        transforms=dummy_transforms,
+        splitter=dummy_splitter,
+        split_columns=["age", "fare"],
     )
 
-    assert isinstance(processor.dataset_manager, DatasetManager)
-    assert processor.columns is not None
+    assert processor.data.shape[0] > 0
+    assert isinstance(processor.transforms, dict)
+    assert isinstance(processor.splitter, splitters_module.RandomSplit)
+    assert isinstance(processor.transforms["age"][0], transforms_module.GaussianNoise)
 
 
 def test_dataset_processor_apply_split(
-    dump_single_split_config_to_disk: str,
     titanic_csv_path: str,
-    split_loader: loaders.SplitLoader,
+    dummy_transforms: dict,
+    dummy_splitter: Any,
 ) -> None:
     """Test applying splits in DatasetProcessor."""
     processor = DatasetProcessor(
-        data_config=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
+        transforms=dummy_transforms,
+        splitter=dummy_splitter,
+        split_columns=["age", "fare"],
     )
-    processor.data = processor.load_csv(titanic_csv_path)
-    processor.add_split(split_manager=SplitManager(split_loader))
-    assert "split" in processor.columns
+
+    # Should start without split column
+    assert "split" not in processor.data.columns
+
+    processor.add_split()
+
     assert "split" in processor.data.columns
-    assert len(processor.data["split"]) == 712
+    assert set(processor.data["split"].unique().to_list()).issubset({0, 1, 2})
 
 
 def test_dataset_processor_apply_transformation_group(
-    dump_single_split_config_to_disk: str,
     titanic_csv_path: str,
-    transform_loader: loaders.TransformLoader,
+    dummy_transforms: dict,
+    dummy_splitter: Any,
 ) -> None:
     """Test applying transformation groups."""
     processor = DatasetProcessor(
-        data_config=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
+        transforms=dummy_transforms,
+        splitter=dummy_splitter,
+        split_columns=["age", "fare"],
+    )
+
+    control = DatasetProcessor(
+        csv_path=titanic_csv_path,
+        transforms={},  # No transforms
+        splitter=dummy_splitter,
+        split_columns=["age", "fare"],
     )
-    processor.data = processor.load_csv(titanic_csv_path)
 
-    processor_control = DatasetProcessor(
-        data_config=dump_single_split_config_to_disk,
+    processor.apply_transformations()
+
+    # Transformed columns should differ
+    assert processor.data["age"].to_list() != control.data["age"].to_list()
+    assert processor.data["fare"].to_list() != control.data["fare"].to_list()
+    # Untransformed columns should match
+    assert processor.data["survived"].to_list() == control.data["survived"].to_list()
+
+
+def test_dataset_processor_shuffle_labels(
+    titanic_csv_path: str,
+    dummy_transforms: dict,
+    dummy_splitter: Any,
+) -> None:
+    """Test shuffling labels."""
+    processor = DatasetProcessor(
         csv_path=titanic_csv_path,
+        transforms=dummy_transforms,
+        splitter=dummy_splitter,
+        split_columns=["age", "fare"],
     )
-    processor_control.data = processor_control.load_csv(titanic_csv_path)
 
-    processor.apply_transformation_group(
-        transform_manager=TransformManager(transform_loader),
+    control = DatasetProcessor(
+        csv_path=titanic_csv_path,
+        transforms=dummy_transforms,
+        splitter=dummy_splitter,
+        split_columns=["age", "fare"],
     )
+    processor.shuffle_labels(label_columns=["survived"])
 
-    assert processor.data["age"].to_list() != processor_control.data["age"].to_list()
-    assert processor.data["fare"].to_list() != processor_control.data["fare"].to_list()
-    assert processor.data["parch"].to_list() == processor_control.data["parch"].to_list()
-    assert processor.data["sibsp"].to_list() == processor_control.data["sibsp"].to_list()
-    assert processor.data["pclass"].to_list() == processor_control.data["pclass"].to_list()
-    assert processor.data["embarked"].to_list() == processor_control.data["embarked"].to_list()
-    assert processor.data["sex"].to_list() == processor_control.data["sex"].to_list()
+    assert processor.data["survived"].to_list() != control.data["survived"].to_list()
 
 
 # Test DatasetLoader
 def test_dataset_loader_init(
-    dump_single_split_config_to_disk: str,
     titanic_csv_path: str,
-    encoder_loader: loaders.EncoderLoader,
+    dummy_encoders: dict,
 ) -> None:
     """Test initialization of DatasetLoader."""
     loader = DatasetLoader(
-        data_config=dump_single_split_config_to_disk,
+        encoders=dummy_encoders,
+        input_columns=["age", "fare"],
+        label_columns=["survived"],
+        meta_columns=["passenger_id"],
         csv_path=titanic_csv_path,
-        encoder_loader=encoder_loader,
     )
 
-    assert isinstance(loader.dataset_manager, DatasetManager)
-    assert loader.data is not None
-    assert loader.columns is not None
-    assert hasattr(loader, "encoder_manager")
+    assert loader.data.shape[0] > 0
+    assert isinstance(loader.encoders, dict)
+    assert len(loader.input_columns) == 2
 
 
 def test_dataset_loader_get_dataset(
-    dump_single_split_config_to_disk: str,
     titanic_csv_path: str,
-    encoder_loader: loaders.EncoderLoader,
+    dummy_encoders: dict,
 ) -> None:
     """Test getting dataset from loader."""
     loader = DatasetLoader(
-        data_config=dump_single_split_config_to_disk,
+        encoders=dummy_encoders,
+        input_columns=["age", "fare"],
+        label_columns=["survived"],
+        meta_columns=["passenger_id"],
         csv_path=titanic_csv_path,
-        encoder_loader=encoder_loader,
     )
 
-    dataset = loader.get_all_items()
-    assert isinstance(dataset, tuple)
-    assert len(dataset) == 3  # input_data, label_data, meta_data
+    inputs, labels, meta = loader.get_all_items()
+    assert isinstance(inputs, dict)
+    assert isinstance(labels, dict)
+    assert isinstance(meta, dict)
+    assert len(inputs["age"]) == len(loader.data)
+
+    # Test encoder types
+    assert isinstance(loader.encoders["age"], encoders_module.NumericEncoder)
+    assert isinstance(loader.encoders["survived"], encoders_module.NumericEncoder)
+
+
+def test_torch_dataset_init(
+    titanic_csv_path: str,
+    dummy_encoders: dict,
+) -> None:
+    """Test initialization of TorchDataset."""
+    loader = DatasetLoader(
+        encoders=dummy_encoders,
+        input_columns=["age", "fare"],
+        label_columns=["survived"],
+        meta_columns=["passenger_id"],
+        csv_path=titanic_csv_path,
+    )
+    dataset = TorchDataset(loader)
+    assert len(dataset) > 0
+    assert isinstance(dataset[0], tuple)
+    assert len(dataset[0]) == 3
+    assert isinstance(dataset[0][0], dict)
+    assert isinstance(dataset[0][1], dict)
+    assert isinstance(dataset[0][2], dict)
+
+
+def test_torch_dataset_get_item(
+    titanic_csv_path: str,
+    dummy_encoders: dict,
+) -> None:
+    """Test getting item from TorchDataset."""
+    loader = DatasetLoader(
+        encoders=dummy_encoders,
+        input_columns=["age", "fare"],
+        label_columns=["survived"],
+        meta_columns=["passenger_id"],
+        csv_path=titanic_csv_path,
+    )
+    dataset = TorchDataset(loader)
+
+    inputs, labels, meta = dataset[0]
+    assert isinstance(inputs, dict)
+    assert isinstance(labels, dict)
+    assert isinstance(meta, dict)
+    assert isinstance(inputs["age"], torch.Tensor)
+    assert isinstance(inputs["fare"], torch.Tensor)
+    assert isinstance(labels["survived"], torch.Tensor)
diff --git a/tests/data/test_experiment.py b/tests/data/test_experiment.py
deleted file mode 100644
index 8a907e3..0000000
--- a/tests/data/test_experiment.py
+++ /dev/null
@@ -1,168 +0,0 @@
-"""Tests for experiment functionality and configuration."""
-
-import pytest
-import yaml
-
-from stimulus.data import loaders
-from stimulus.data.encoding.encoders import AbstractEncoder
-from stimulus.data.splitters import splitters
-from stimulus.data.transform import data_transformation_generators
-from stimulus.utils import yaml_data
-
-
-@pytest.fixture
-def dna_experiment_config_path() -> str:
-    """Fixture that provides the path to the DNA experiment config template YAML file.
-
-    Returns:
-        str: Path to the DNA experiment config template YAML file
-    """
-    return "tests/test_data/dna_experiment/dna_experiment_config_template.yaml"
-
-
-@pytest.fixture
-def dna_experiment_sub_yaml(
-    dna_experiment_config_path: str,
-) -> yaml_data.YamlConfigDict:
-    """Get a sub-configuration from the DNA experiment config.
-
-    Args:
-        dna_experiment_config_path: Path to the DNA experiment config file
-
-    Returns:
-        yaml_data.YamlConfigDict: First generated sub-configuration
-    """
-    with open(dna_experiment_config_path) as f:
-        yaml_dict = yaml.safe_load(f)
-        yaml_config = yaml_data.YamlConfigDict(**yaml_dict)
-
-    yaml_split_configs = yaml_data.generate_split_configs(yaml_config)
-    yaml_split_transform_configs = yaml_data.generate_split_transform_configs(
-        yaml_split_configs[0],
-    )
-    return yaml_split_transform_configs[0]
-
-
-@pytest.fixture
-def titanic_yaml_path() -> str:
-    """Get path to Titanic YAML config file.
-
-    Returns:
-        str: Path to Titanic config file
-    """
-    return "tests/test_data/titanic/titanic.yaml"
-
-
-@pytest.fixture
-def titanic_sub_yaml_path() -> str:
-    """Get path to Titanic sub-config YAML file.
-
-    Returns:
-        str: Path to Titanic sub-config file
-    """
-    return "tests/test_data/titanic/titanic_sub_config_0.yaml"
-
-
-@pytest.fixture
-def text_onehot_encoder_params() -> tuple[str, dict[str, str]]:
-    """Get TextOneHotEncoder name and parameters.
-
-    Returns:
-        tuple[str, dict[str, str]]: Encoder name and parameters
-    """
-    return "TextOneHotEncoder", {"alphabet": "acgt"}
-
-
-def test_get_encoder(text_onehot_encoder_params: tuple[str, dict[str, str]]) -> None:
-    """Test the get_encoder method of the AbstractExperiment class.
-
-    Args:
-        text_onehot_encoder_params: Tuple of encoder name and parameters
-    """
-    experiment = loaders.EncoderLoader()
-    encoder_name, encoder_params = text_onehot_encoder_params
-    encoder = experiment.get_encoder(encoder_name, encoder_params)
-    assert isinstance(encoder, AbstractEncoder)
-
-
-def test_set_encoder_as_attribute(
-    text_onehot_encoder_params: tuple[str, dict[str, str]],
-) -> None:
-    """Test the set_encoder_as_attribute method of the AbstractExperiment class.
-
-    Args:
-        text_onehot_encoder_params: Tuple of encoder name and parameters
-    """
-    experiment = loaders.EncoderLoader()
-    encoder_name, encoder_params = text_onehot_encoder_params
-    encoder = experiment.get_encoder(encoder_name, encoder_params)
-    experiment.set_encoder_as_attribute("ciao", encoder)
-    assert hasattr(experiment, "ciao")
-    assert experiment.ciao == encoder
-    assert experiment.get_function_encode_all("ciao") == encoder.encode_all
-
-
-def test_build_experiment_class_encoder_dict(
-    dna_experiment_sub_yaml: yaml_data.YamlConfigDict,
-) -> None:
-    """Test the build_experiment_class_encoder_dict method.
-
-    Args:
-        dna_experiment_sub_yaml: DNA experiment sub-configuration
-    """
-    experiment = loaders.EncoderLoader()
-    config = dna_experiment_sub_yaml.columns
-    experiment.initialize_column_encoders_from_config(config)
-    assert hasattr(experiment, "hello")
-    assert hasattr(experiment, "bonjour")
-    assert hasattr(experiment, "ciao")
-
-    assert experiment.hello.encode_all(["a", "c", "g", "t"]) is not None
-
-
-def test_get_data_transformer() -> None:
-    """Test the get_data_transformer method of the TransformLoader class."""
-    experiment = loaders.TransformLoader()
-    transformer = experiment.get_data_transformer("ReverseComplement")
-    assert isinstance(transformer, data_transformation_generators.ReverseComplement)
-
-
-def test_set_data_transformer_as_attribute() -> None:
-    """Test the set_data_transformer_as_attribute method."""
-    experiment = loaders.TransformLoader()
-    transformer = experiment.get_data_transformer("ReverseComplement")
-    experiment.set_data_transformer_as_attribute("col1", transformer)
-    assert hasattr(experiment, "col1")
-    assert experiment.col1["ReverseComplement"] == transformer
-
-
-def test_initialize_column_data_transformers_from_config(
-    dna_experiment_sub_yaml: yaml_data.YamlConfigDict,
-) -> None:
-    """Test initializing column data transformers from config.
-
-    Args:
-        dna_experiment_sub_yaml: DNA experiment sub-configuration
-    """
-    experiment = loaders.TransformLoader()
-    config = dna_experiment_sub_yaml.transforms
-    experiment.initialize_column_data_transformers_from_config(config)
-
-    assert hasattr(experiment, "col1")
-    column_transformers = experiment.col1
-    assert any(isinstance(t, data_transformation_generators.ReverseComplement) for t in column_transformers.values())
-
-
-def test_initialize_splitter_from_config(
-    dna_experiment_sub_yaml: yaml_data.YamlConfigDict,
-) -> None:
-    """Test initializing splitter from configuration.
-
-    Args:
-        dna_experiment_sub_yaml: DNA experiment sub-configuration
-    """
-    experiment = loaders.SplitLoader()
-    config = dna_experiment_sub_yaml.split
-    experiment.initialize_splitter_from_config(config)
-    assert hasattr(experiment, "split")
-    assert isinstance(experiment.split, splitters.RandomSplit)
diff --git a/tests/data/test_handlertorch.py b/tests/data/test_handlertorch.py
deleted file mode 100644
index dff2490..0000000
--- a/tests/data/test_handlertorch.py
+++ /dev/null
@@ -1,139 +0,0 @@
-"""Tests for PyTorch data handling functionality."""
-
-import os
-
-import pytest
-import yaml
-
-from stimulus.data import handlertorch, loaders
-from stimulus.utils import yaml_data
-
-
-@pytest.fixture
-def titanic_config_path() -> str:
-    """Get path to Titanic config file.
-
-    Returns:
-        str: Absolute path to the config file
-    """
-    return os.path.abspath("tests/test_data/titanic/titanic_sub_config.yaml")
-
-
-@pytest.fixture
-def titanic_csv_path() -> str:
-    """Get path to Titanic CSV file.
-
-    Returns:
-        str: Absolute path to the CSV file
-    """
-    return os.path.abspath("tests/test_data/titanic/titanic_stimulus.csv")
-
-
-@pytest.fixture
-def titanic_yaml_config(titanic_config_path: str) -> dict:
-    """Load Titanic YAML config.
-
-    Args:
-        titanic_config_path: Path to the config file
-
-    Returns:
-        dict: Loaded YAML configuration
-    """
-    with open(titanic_config_path) as file:
-        return yaml_data.YamlSplitTransformDict(**yaml.safe_load(file))
-
-
-@pytest.fixture
-def titanic_encoder_loader(
-    titanic_yaml_config: yaml_data.YamlSplitTransformDict,
-) -> loaders.EncoderLoader:
-    """Get Titanic encoder loader."""
-    loader = loaders.EncoderLoader()
-    loader.initialize_column_encoders_from_config(titanic_yaml_config.columns)
-    return loader
-
-
-def test_init_handlertorch(
-    titanic_config_path: str,
-    titanic_csv_path: str,
-    titanic_encoder_loader: loaders.EncoderLoader,
-) -> None:
-    """Test TorchDataset initialization."""
-    data_config: yaml_data.YamlSplitTransformDict
-    with open(titanic_config_path) as f:
-        data_config = yaml_data.YamlSplitTransformDict(**yaml.safe_load(f))
-    handlertorch.TorchDataset(
-        data_config=data_config,
-        csv_path=titanic_csv_path,
-        encoder_loader=titanic_encoder_loader,
-    )
-
-
-def test_len_handlertorch(
-    titanic_config_path: str,
-    titanic_csv_path: str,
-    titanic_encoder_loader: loaders.EncoderLoader,
-) -> None:
-    """Test length functionality of TorchDataset.
-
-    Args:
-        titanic_config_path: Path to config file
-        titanic_csv_path: Path to CSV file
-        titanic_encoder_loader: Encoder loader instance
-    """
-    data_config: yaml_data.YamlSplitTransformDict
-    with open(titanic_config_path) as f:
-        data_config = yaml_data.YamlSplitTransformDict(**yaml.safe_load(f))
-    dataset = handlertorch.TorchDataset(
-        data_config=data_config,
-        csv_path=titanic_csv_path,
-        encoder_loader=titanic_encoder_loader,
-    )
-    assert len(dataset) == 712
-
-
-def test_getitem_handlertorch_slice(
-    titanic_config_path: str,
-    titanic_csv_path: str,
-    titanic_encoder_loader: loaders.EncoderLoader,
-) -> None:
-    """Test slice indexing functionality of TorchDataset.
-
-    Args:
-        titanic_config_path: Path to config file
-        titanic_csv_path: Path to CSV file
-        titanic_encoder_loader: Encoder loader instance
-    """
-    data_config: yaml_data.YamlSplitTransformDict
-    with open(titanic_config_path) as f:
-        data_config = yaml_data.YamlSplitTransformDict(**yaml.safe_load(f))
-    dataset = handlertorch.TorchDataset(
-        data_config=data_config,
-        csv_path=titanic_csv_path,
-        encoder_loader=titanic_encoder_loader,
-    )
-    assert len(dataset[0:5]) == 3
-    assert len(dataset[0:5][0]["pclass"]) == 5
-
-
-def test_getitem_handlertorch_int(
-    titanic_config_path: str,
-    titanic_csv_path: str,
-    titanic_encoder_loader: loaders.EncoderLoader,
-) -> None:
-    """Test integer indexing functionality of TorchDataset.
-
-    Args:
-        titanic_config_path: Path to config file
-        titanic_csv_path: Path to CSV file
-        titanic_encoder_loader: Encoder loader instance
-    """
-    data_config: yaml_data.YamlSplitTransformDict
-    with open(titanic_config_path) as f:
-        data_config = yaml_data.YamlSplitTransformDict(**yaml.safe_load(f))
-    dataset = handlertorch.TorchDataset(
-        data_config=data_config,
-        csv_path=titanic_csv_path,
-        encoder_loader=titanic_encoder_loader,
-    )
-    assert len(dataset[0]) == 3
diff --git a/tests/data/transform/test_data_transformers.py b/tests/data/transform/test_data_transformers.py
index 053a171..26ab9ad 100644
--- a/tests/data/transform/test_data_transformers.py
+++ b/tests/data/transform/test_data_transformers.py
@@ -1,22 +1,17 @@
 """Test suite for the data transformation generators."""
 
-import os
 from typing import Any
 
 import numpy as np
 import pytest
 
-from src.stimulus.data.transform.data_transformation_generators import (
-    AbstractDataTransformer,
+from src.stimulus.data.transforming.transforms import (
+    AbstractTransform,
     GaussianChunk,
     GaussianNoise,
     ReverseComplement,
     UniformTextMasker,
 )
-from stimulus.utils.yaml_data import (
-    dump_yaml_list_into_files,
-    generate_split_transform_configs,
-)
 
 
 class DataTransformerTest:
@@ -33,7 +28,7 @@ class DataTransformerTest:
 
     def __init__(  # noqa: D107
         self,
-        transformer: AbstractDataTransformer,
+        transformer: AbstractTransform,
         params: dict,
         single_input: Any,
         expected_single_output: Any,
@@ -256,19 +251,3 @@ class TestReverseComplement:
         for item in transformed_data:
             assert isinstance(item, str)
         assert transformed_data == test_data.expected_multiple_outputs
-
-
-@pytest.fixture
-def titanic_config_path(base_config: dict) -> str:
-    """Ensure the config file exists and return its path."""
-    config_path = "tests/test_data/titanic/titanic_sub_config_0.yaml"
-
-    if not os.path.exists(config_path):
-        configs = generate_split_transform_configs(base_config)
-        dump_yaml_list_into_files(
-            [configs[0]],
-            "tests/test_data/titanic/",
-            "titanic_sub_config",
-        )
-
-    return os.path.abspath(config_path)
diff --git a/tests/learner/test_raytune_learner.py b/tests/learner/test_raytune_learner.py
index 03b21e9..9d80550 100644
--- a/tests/learner/test_raytune_learner.py
+++ b/tests/learner/test_raytune_learner.py
@@ -5,13 +5,13 @@ import warnings
 
 import pytest
 import ray
+import torch
 import yaml
 
-from stimulus.data.handlertorch import TorchDataset
-from stimulus.data.loaders import EncoderLoader
+from stimulus.data.data_handlers import DatasetLoader, TorchDataset
+from stimulus.data.encoding import encoders as encoders_module
 from stimulus.learner.raytune_learner import TuneWrapper
-from stimulus.utils.yaml_data import YamlSplitTransformDict
-from stimulus.utils.yaml_model_schema import Model, RayTuneModel, YamlRayConfigLoader
+from stimulus.utils.yaml_model_schema import Model, RayConfigLoader, RayTuneModel
 from tests.test_model import titanic_model
 
 
@@ -20,35 +20,95 @@ def ray_config_loader() -> RayTuneModel:
     """Load the RayTuneModel configuration."""
     with open("tests/test_model/titanic_model_cpu.yaml") as file:
         model_config = yaml.safe_load(file)
-    return YamlRayConfigLoader(Model(**model_config)).get_config()
+    return RayConfigLoader(Model(**model_config)).get_config()
 
 
 @pytest.fixture
-def encoder_loader() -> EncoderLoader:
+def get_encoders() -> dict[str, encoders_module.AbstractEncoder]:
     """Load the EncoderLoader configuration."""
-    with open("tests/test_data/titanic/titanic_sub_config.yaml") as file:
-        data_config = yaml.safe_load(file)
-    encoder_loader = EncoderLoader()
-    encoder_loader.initialize_column_encoders_from_config(
-        YamlSplitTransformDict(**data_config).columns,
-    )
-    return encoder_loader
+    return {
+        "passenger_id": encoders_module.NumericEncoder(dtype=torch.int64),
+        "survived": encoders_module.NumericEncoder(dtype=torch.int64),
+        "pclass": encoders_module.NumericEncoder(dtype=torch.int64),
+        "sex": encoders_module.StrClassificationEncoder(),
+        "age": encoders_module.NumericEncoder(dtype=torch.float32),
+        "sibsp": encoders_module.NumericEncoder(dtype=torch.int64),
+        "parch": encoders_module.NumericEncoder(dtype=torch.int64),
+        "fare": encoders_module.NumericEncoder(dtype=torch.float32),
+        "embarked": encoders_module.StrClassificationEncoder(),
+    }
+
+
+@pytest.fixture
+def get_input_columns() -> list[str]:
+    """Get the input columns."""
+    return ["embarked", "pclass", "sex", "age", "sibsp", "parch", "fare"]
+
+
+@pytest.fixture
+def get_label_columns() -> list[str]:
+    """Get the label columns."""
+    return ["survived"]
+
+
+@pytest.fixture
+def get_meta_columns() -> list[str]:
+    """Get the meta columns."""
+    return ["passenger_id"]
 
 
 @pytest.fixture
-def titanic_dataset(encoder_loader: EncoderLoader) -> TorchDataset:
+def titanic_dataset() -> TorchDataset:
     """Create a TorchDataset instance for testing."""
+    return "tests/test_data/titanic/titanic_stimulus_split.csv"
+
+
+@pytest.fixture
+def get_train_dataset(
+    titanic_dataset: str,
+    get_encoders: dict[str, encoders_module.AbstractEncoder],
+    get_input_columns: list[str],
+    get_label_columns: list[str],
+    get_meta_columns: list[str],
+) -> TorchDataset:
+    """Get the DatasetLoader."""
+    return TorchDataset(
+        loader=DatasetLoader(
+            csv_path=titanic_dataset,
+            encoders=get_encoders,
+            input_columns=get_input_columns,
+            label_columns=get_label_columns,
+            meta_columns=get_meta_columns,
+            split=0,
+        ),
+    )
+
+
+@pytest.fixture
+def get_validation_dataset(
+    titanic_dataset: str,
+    get_encoders: dict[str, encoders_module.AbstractEncoder],
+    get_input_columns: list[str],
+    get_label_columns: list[str],
+    get_meta_columns: list[str],
+) -> TorchDataset:
+    """Get the DatasetLoader."""
     return TorchDataset(
-        csv_path="tests/test_data/titanic/titanic_stimulus_split.csv",
-        config_path="tests/test_data/titanic/titanic_sub_config.yaml",
-        encoder_loader=encoder_loader,
-        split=0,
+        loader=DatasetLoader(
+            csv_path=titanic_dataset,
+            encoders=get_encoders,
+            input_columns=get_input_columns,
+            label_columns=get_label_columns,
+            meta_columns=get_meta_columns,
+            split=1,
+        ),
     )
 
 
 def test_tunewrapper_init(
     ray_config_loader: RayTuneModel,
-    encoder_loader: EncoderLoader,
+    get_train_dataset: TorchDataset,
+    get_validation_dataset: TorchDataset,
 ) -> None:
     """Test the initialization of the TuneWrapper class."""
     # Filter ResourceWarning during Ray shutdown
@@ -58,16 +118,11 @@ def test_tunewrapper_init(
     ray.init(ignore_reinit_error=True)
 
     try:
-        data_config: YamlSplitTransformDict
-        with open("tests/test_data/titanic/titanic_sub_config.yaml") as f:
-            data_config = YamlSplitTransformDict(**yaml.safe_load(f))
-
         tune_wrapper = TuneWrapper(
             model_config=ray_config_loader,
             model_class=titanic_model.ModelTitanic,
-            data_path="tests/test_data/titanic/titanic_stimulus_split.csv",
-            data_config=data_config,
-            encoder_loader=encoder_loader,
+            train_dataset=get_train_dataset,
+            validation_dataset=get_validation_dataset,
             seed=42,
             ray_results_dir=os.path.abspath("tests/test_data/titanic/ray_results"),
             tune_run_name="test_run",
@@ -88,7 +143,8 @@ def test_tunewrapper_init(
 
 def test_tune_wrapper_tune(
     ray_config_loader: RayTuneModel,
-    encoder_loader: EncoderLoader,
+    get_train_dataset: TorchDataset,
+    get_validation_dataset: TorchDataset,
 ) -> None:
     """Test the tune method of TuneWrapper class."""
     # Filter ResourceWarning during Ray shutdown
@@ -98,16 +154,11 @@ def test_tune_wrapper_tune(
     ray.init(ignore_reinit_error=True)
 
     try:
-        data_config: YamlSplitTransformDict
-        with open("tests/test_data/titanic/titanic_sub_config.yaml") as f:
-            data_config = YamlSplitTransformDict(**yaml.safe_load(f))
-
         tune_wrapper = TuneWrapper(
             model_config=ray_config_loader,
             model_class=titanic_model.ModelTitanic,
-            data_path="tests/test_data/titanic/titanic_stimulus_split.csv",
-            data_config=data_config,
-            encoder_loader=encoder_loader,
+            train_dataset=get_train_dataset,
+            validation_dataset=get_validation_dataset,
             seed=42,
             ray_results_dir=os.path.abspath("tests/test_data/titanic/ray_results"),
             tune_run_name="test_run",
diff --git a/tests/test_data/titanic/titanic_unique_split.yaml b/tests/test_data/titanic/titanic_unique_split.yaml
index 3ef7d17..9b90604 100644
--- a/tests/test_data/titanic/titanic_unique_split.yaml
+++ b/tests/test_data/titanic/titanic_unique_split.yaml
@@ -59,30 +59,18 @@ columns:
         params: {}
 
 transforms:
-  - transformation_name: noise
+  - transformation_name: "noise"
     columns:
-      - column_name: age
+      - column_name: "age"
         transformations:
           - name: GaussianNoise
             params:
-              std: 0.1
-      - column_name: fare
+              std: [0.1, 0.2, 0.3]
+      - column_name: "fare"
         transformations:
           - name: GaussianNoise
             params:
-              std: 0.1
-  - transformation_name: noise2
-    columns:
-      - column_name: age
-        transformations:
-          - name: GaussianNoise
-            params:
-              std: 0.1
-      - column_name: fare
-        transformations:
-          - name: GaussianNoise
-            params:
-              std: 0.1
+              std: [0.1, 0.2, 0.3]
 
 split:
   split_method: RandomSplit
diff --git a/tests/test_data/titanic/titanic_sub_config.yaml b/tests/test_data/titanic/titanic_unique_transform.yaml
similarity index 100%
rename from tests/test_data/titanic/titanic_sub_config.yaml
rename to tests/test_data/titanic/titanic_unique_transform.yaml
diff --git a/tests/typing/test_typing.py b/tests/typing/test_typing.py
index d2be37f..eaa60dc 100644
--- a/tests/typing/test_typing.py
+++ b/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")
@@ -27,6 +24,7 @@ def test_learner_types() -> None:
     """Test the learner types."""
     try:
         from stimulus.typing import (
+            CheckpointDict,
             PredictWrapper,
             RayTuneMetrics,
             RayTuneOptimizer,
@@ -39,24 +37,24 @@ def test_learner_types() -> None:
         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:
@@ -67,13 +65,13 @@ def test_yaml_model_schema_types() -> None:
             Data,
             Loss,
             Model,
+            RayConfigLoader,
             RayTuneModel,
             RunParams,
             Scheduler,
             TunableParameter,
             Tune,
             TuneParams,
-            YamlRayConfigLoader,
         )
     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 Data, RayTuneData
     except ImportError:
         pytest.fail("Failed to import Type Aliases")
diff --git a/tests/utils/test_data_yaml.py b/tests/utils/test_data_yaml.py
deleted file mode 100644
index 2bc23ec..0000000
--- a/tests/utils/test_data_yaml.py
+++ /dev/null
@@ -1,137 +0,0 @@
-"""Tests for YAML data handling functionality."""
-
-import pytest
-import yaml
-
-from src.stimulus.utils import yaml_data
-from src.stimulus.utils.yaml_data import (
-    YamlConfigDict,
-    YamlSplitConfigDict,
-    YamlSplitTransformDict,
-    generate_split_configs,
-    generate_split_transform_configs,
-)
-
-
-@pytest.fixture
-def titanic_csv_path() -> str:
-    """Get path to Titanic CSV file."""
-    return "tests/test_data/titanic/titanic_stimulus.csv"
-
-
-@pytest.fixture
-def load_titanic_yaml_from_file() -> YamlConfigDict:
-    """Fixture that loads a test YAML configuration file."""
-    with open("tests/test_data/titanic/titanic.yaml") as f:
-        yaml_dict = yaml.safe_load(f)
-        return YamlConfigDict(**yaml_dict)
-
-
-@pytest.fixture
-def load_split_config_yaml_from_file() -> YamlSplitConfigDict:
-    """Fixture that loads a test unique split YAML configuration file."""
-    with open("tests/test_data/titanic/titanic_unique_split.yaml") as f:
-        yaml_dict = yaml.safe_load(f)
-        return YamlSplitConfigDict(**yaml_dict)
-
-
-@pytest.fixture
-def load_yaml_from_file() -> YamlConfigDict:
-    """Fixture that loads a test YAML configuration file."""
-    with open(
-        "tests/test_data/dna_experiment/dna_experiment_config_template.yaml",
-    ) as f:
-        yaml_dict = yaml.safe_load(f)
-        return YamlConfigDict(**yaml_dict)
-
-
-@pytest.fixture
-def load_wrong_type_yaml() -> dict:
-    """Fixture that loads a YAML configuration file with wrong typing."""
-    with open("tests/test_data/yaml_files/wrong_field_type.yaml") as f:
-        return yaml.safe_load(f)
-
-
-def test_split_config_validation(load_titanic_yaml_from_file: YamlConfigDict) -> None:
-    """Test split configuration validation."""
-    split_config = generate_split_configs(load_titanic_yaml_from_file)[0]
-    YamlSplitConfigDict.model_validate(split_config)
-
-
-def test_sub_config_validation(
-    load_split_config_yaml_from_file: YamlSplitConfigDict,
-) -> None:
-    """Test sub-config validation."""
-    split_config = generate_split_transform_configs(
-        load_split_config_yaml_from_file,
-    )[0]
-    YamlSplitTransformDict.model_validate(split_config)
-
-
-def test_expand_transform_parameter_combinations(
-    load_yaml_from_file: YamlConfigDict,
-) -> None:
-    """Tests expanding transforms with parameter lists into individual transforms."""
-    # Test transform with multiple parameter lists
-    transform = load_yaml_from_file.transforms[0]
-    results = yaml_data.expand_transform_parameter_combinations(transform)
-    assert len(results) == 1  # Only one transform returned
-    assert isinstance(
-        results[0],
-        yaml_data.YamlTransform,
-    )  # Should return YamlTransform objects
-
-
-def test_expand_transform_list_combinations(
-    load_yaml_from_file: YamlConfigDict,
-) -> None:
-    """Tests expanding a list of transforms into all parameter combinations."""
-    results = yaml_data.expand_transform_list_combinations(
-        load_yaml_from_file.transforms,
-    )
-    # 4 combinations from first transform x 2 from second
-    assert len(results) == 8
-    # Each result should be a YamlTransform
-    for result in results:
-        assert isinstance(result, yaml_data.YamlTransform)
-        assert isinstance(result.transformation_name, str)
-        assert isinstance(result.columns, list)
-
-
-def test_generate_data_configs(
-    load_yaml_from_file: YamlConfigDict,
-) -> None:
-    """Tests generating all possible data configurations."""
-    split_configs = yaml_data.generate_split_configs(load_yaml_from_file)
-    configs: list[YamlSplitTransformDict] = []
-    for s_conf in split_configs:
-        configs.extend(generate_split_transform_configs(s_conf))
-
-    assert len(configs) == 16  # 8 transform combinations x 2 splits
-
-    # Check each config individually to help debug
-    for i, config in enumerate(configs):
-        assert isinstance(
-            config,
-            yaml_data.YamlSplitTransformDict,
-        ), f"Config {i} is type {type(config)}, expected YamlSubConfigDict"
-
-
-@pytest.mark.parametrize(
-    "test_input",
-    [("load_yaml_from_file", False), ("load_wrong_type_yaml", True)],
-)
-def test_check_yaml_schema(
-    request: pytest.FixtureRequest,
-    test_input: tuple[str, bool],
-) -> None:
-    """Tests the Pydantic schema validation."""
-    data = request.getfixturevalue(test_input[0])
-    if test_input[1]:
-        with pytest.raises(
-            ValueError,
-            match="Wrong type on a field, see the pydantic report above",
-        ):
-            yaml_data.check_yaml_schema(data)
-    else:
-        yaml_data.check_yaml_schema(data)
diff --git a/tests/utils/test_model_yaml.py b/tests/utils/test_model_yaml.py
index 0e08aba..92d1fcf 100644
--- a/tests/utils/test_model_yaml.py
+++ b/tests/utils/test_model_yaml.py
@@ -26,7 +26,7 @@ def test_raytune_space_selector(titanic_model_yaml_path: str) -> None:
     with open(titanic_model_yaml_path) as file:
         model_config = yaml.safe_load(file)
     model = yaml_model_schema.Model.model_validate(model_config)
-    loader = yaml_model_schema.YamlRayConfigLoader(model)
+    loader = yaml_model_schema.RayConfigLoader(model)
 
     result = loader.raytune_space_selector(tune.choice, [1, 2, 3])
     assert str(type(result).__name__) == "Categorical"
@@ -43,7 +43,7 @@ def test_convert_raytune(titanic_model_yaml_path: str) -> None:
     with open(titanic_model_yaml_path) as file:
         model_config = yaml.safe_load(file)
     model = yaml_model_schema.Model.model_validate(model_config)
-    loader = yaml_model_schema.YamlRayConfigLoader(model)
+    loader = yaml_model_schema.RayConfigLoader(model)
 
     param = yaml_model_schema.TunableParameter(space=[16, 32, 64], mode="choice")
     result = loader.convert_raytune(param)
@@ -59,7 +59,7 @@ def test_convert_config_to_ray(titanic_model_yaml_path: str) -> None:
     with open(titanic_model_yaml_path) as file:
         model_config = yaml.safe_load(file)
     model = yaml_model_schema.Model.model_validate(model_config)
-    loader = yaml_model_schema.YamlRayConfigLoader(model)
+    loader = yaml_model_schema.RayConfigLoader(model)
 
     ray_model = loader.convert_config_to_ray(model)
     assert isinstance(ray_model, yaml_model_schema.RayTuneModel)
@@ -74,7 +74,7 @@ def test_get_config(titanic_model_yaml_path: str) -> None:
     with open(titanic_model_yaml_path) as file:
         model_config = yaml.safe_load(file)
     model = yaml_model_schema.Model.model_validate(model_config)
-    loader = yaml_model_schema.YamlRayConfigLoader(model)
+    loader = yaml_model_schema.RayConfigLoader(model)
 
     config = loader.get_config()
     assert isinstance(config, yaml_model_schema.RayTuneModel)
@@ -89,7 +89,7 @@ def test_sampint() -> None:
     sample_space = [1, 5]
     n_space = [2, 4]
 
-    result = yaml_model_schema.YamlRayConfigLoader.sampint(sample_space, n_space)
+    result = yaml_model_schema.RayConfigLoader.sampint(sample_space, n_space)
 
     assert isinstance(result, list)
     assert all(isinstance(x, int) for x in result)