From 1d355169edec4a2227fba7df95ab38a5e68c2c40 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 14:28:14 +0100
Subject: [PATCH 01/14] DEPRECATED/LINT: removed deprecated json_schema and
 added linting for utils and raytune parser

---
 src/stimulus/learner/raytune_parser.py |  20 +-
 src/stimulus/utils/generic_utils.py    |   7 +-
 src/stimulus/utils/json_schema.py      | 319 -------------------------
 3 files changed, 20 insertions(+), 326 deletions(-)
 delete mode 100644 src/stimulus/utils/json_schema.py

diff --git a/src/stimulus/learner/raytune_parser.py b/src/stimulus/learner/raytune_parser.py
index b62c1178..c825967f 100644
--- a/src/stimulus/learner/raytune_parser.py
+++ b/src/stimulus/learner/raytune_parser.py
@@ -1,3 +1,5 @@
+"""Ray Tune results parser for extracting and saving best model configurations and weights."""
+
 import json
 import os
 
@@ -7,14 +9,15 @@
 
 
 class TuneParser:
-    def __init__(self, results):
+    """Parser class for Ray Tune results to extract best configurations and model weights."""
+
+    def __init__(self, results: object) -> None:
         """`results` is the output of ray.tune."""
         self.results = results
 
     def get_best_config(self) -> dict:
         """Get the best config from the results."""
-        config = self.results.get_best_result().config
-        return config
+        return self.results.get_best_result().config
 
     def save_best_config(self, output: str) -> None:
         """Save the best config to a file.
@@ -26,8 +29,15 @@ def save_best_config(self, output: str) -> None:
         with open(output, "w") as f:
             json.dump(config, f, indent=4)
 
-    def fix_config_values(self, config):
-        """Correct config values."""
+    def fix_config_values(self, config: dict) -> dict:
+        """Correct config values.
+
+        Args:
+            config: Configuration dictionary to fix
+
+        Returns:
+            Fixed configuration dictionary
+        """
         # fix the model and experiment values to avoid problems with serialization
         # TODO this is a quick fix to avoid the problem with serializing class objects. maybe there is a better way.
         config["model"] = config["model"].__name__
diff --git a/src/stimulus/utils/generic_utils.py b/src/stimulus/utils/generic_utils.py
index 95198f00..5e7ed0da 100644
--- a/src/stimulus/utils/generic_utils.py
+++ b/src/stimulus/utils/generic_utils.py
@@ -1,3 +1,5 @@
+"""Utility functions for general purpose operations like seed setting and tensor manipulation."""
+
 import random
 from typing import Union
 
@@ -13,8 +15,9 @@ def ensure_at_least_1d(tensor: torch.Tensor) -> torch.Tensor:
 
 
 def set_general_seeds(seed_value: Union[int, None]) -> None:
-    """Function that sets all the relevant seeds to a given value. Especially usefull in case of ray.tune.
-    Ray does not have a "generic" seed as far as ray 2.23
+    """Set all relevant random seeds to a given value.
+
+    Especially useful in case of ray.tune. Ray does not have a "generic" seed as far as ray 2.23.
     """
     # Set python seed
     random.seed(seed_value)
diff --git a/src/stimulus/utils/json_schema.py b/src/stimulus/utils/json_schema.py
deleted file mode 100644
index 99c3a45a..00000000
--- a/src/stimulus/utils/json_schema.py
+++ /dev/null
@@ -1,319 +0,0 @@
-from abc import ABC
-from itertools import product
-
-
-class JsonSchema(ABC):
-    """This class helps decode and work on a difened Json schema used by the stimulus pipeline.
-    TODO add Json.schema real library to control that each transform, split have the correct keys associated to them.
-    link -> https://json-schema.org/learn/getting-started-step-by-step#create
-    """
-
-    def __init__(self, schema: dict) -> None:
-        self.schema = schema
-        self.interpret_params_mode = schema.get("interpret_params_mode", "column_wise")
-        self.experiment = schema.get("experiment")
-        self.transform_arg = schema.get("transform", [])
-        self.split_arg = schema.get("split", [])
-        self.custom_arg = schema.get("custom", [])
-
-        # Send error if experiment name is missing
-        if not self.experiment:
-            raise ValueError(
-                "No experiment name given, the Json should always have a experiment:'ExperimentName' field",
-            )
-
-        # Send error if self.interpret_parmas_mode is not among the possible ones
-        if self.interpret_params_mode not in ["column_wise", "all_combinations"]:
-            raise ValueError(
-                "interpret_params_mode value can only be one of the following keywords -> ['column_wise', 'all_combinations']",
-            )
-
-        # check that inside transform dictionary there are no repeated column_nmae values and return them otherwise send error
-        self.column_names = self._check_repeated_column_names()
-
-        # check that transform dictionary have a coherent number of parameters values in case of column_wise for self.interpret_parmas_mode
-        self.number_column_wise_val = self._check_transform_params_schema()
-
-    def _check_repeated_column_names(self) -> list:
-        """Helper function that ensures that inside transform dictionary there are no column:names repeated values"""
-        # in case there is no transform or split flag but a custom one instead
-        if not self.transform_arg and self.custom_arg:
-            return None
-
-        column_name_list = []
-        for i, dictionary in enumerate(self.transform_arg):
-            # None can be inside this list of arguments if that is the case just ignore it. It will be handeled later on.
-            if dictionary is None:
-                continue
-
-            column_name = dictionary["column_name"]
-
-            # If already present as a name throw an error
-            if column_name in column_name_list:
-                raise ValueError(f"The column_name {column_name} is repeated. column_names should be unique.")
-            column_name_list.append(column_name)
-        return column_name_list
-
-    def _check_transform_params_schema(self) -> int:
-        """Help function to check if the number of values in params in the transform dictionary is consistent among all params.
-        If there is {"NoiserName" : { "params": [{"val1":[0, 1]}], "OtherNoiser" : { "params": [{"val1":[2, 3], "val3":[4]}]}}
-        it will raise error because the val3 has only a list of len() 1 instead of 2
-        otherwise it resturn the len()
-        """
-        # in case there is no transform dictionary or if interpret_params_mode is in all_combinations mode
-        if not self.transform_arg or self.interpret_params_mode == "all_combinations":
-            return 0
-
-        num_params_list = []
-        # Iterate through the given dictionary becuse more than one column_name values could be specified for ex.
-        for i, col_name_dictionary in enumerate(self.transform_arg):
-            # None can be inside this list of arguments if that is the case just ignore it. It will be handeled later on.
-            if col_name_dictionary is None:
-                continue
-
-            # take into account that there could be the keyword default
-            if col_name_dictionary["params"] == "default":
-                continue
-
-            # iterate throught the possible multiple parmaeters, some transformeds could have more than one parameter flag
-            for k, params_dict in enumerate(col_name_dictionary["params"]):
-                # even the single set of parameters of a given transformedname can be set to default
-                if params_dict == "default":
-                    continue
-                for params_flag, params_list in params_dict.items():
-                    num_params_list.append(len(params_list))
-
-        # check that all parameters values found are equal
-        if len(set(num_params_list)) == 1:
-            return num_params_list[0]
-        raise ValueError(
-            "Expected the same number of values for all the params under transform value, but received a discordant ammount instead.",
-        )
-
-    def _reshape_transform_dict(self) -> dict:
-        """This function reshapes the transform argument from the JSON schema into a nested dictionary structure.
-        It iterates through each transform dictionary entry, extracting the column_name and name fields.
-        If the name field is a string, it converts it to a list with a single element. It then creates a dictionary for each transform name, associating it with its parameters.
-        It handles cases where transform names are repeated for the same column_name by appending a unique key to the transform name. key = -#num
-        """
-        transform_dict = {}
-        for col_name_dictionary in self.transform_arg:
-            # The name: field of a transform: can be either a simlpe string or list of strings, so convert such variable to a list if it's a string, otherwise leave it unchanged
-            transformed_list = (
-                [col_name_dictionary["name"]]
-                if isinstance(col_name_dictionary["name"], str)
-                else col_name_dictionary["name"]
-            )
-            # Now get the parametrs or set of parameters associated with each transformed and store both in a tuple and append to list transformed names associated to a given clumn_name
-            for k, transformed_name in enumerate(transformed_list):
-                # handle the fact that params can have "default" as value and not a list
-                if col_name_dictionary["params"] == "default":
-                    params_to_be_added = "default"
-                else:
-                    params_to_be_added = col_name_dictionary["params"][k]
-                # handle the case of multiple transformed with same name in the same list associated to the column_name, solution -> create a scheme to modify the name
-                if transform_dict.get(col_name_dictionary["column_name"]) and transformed_name in transform_dict.get(
-                    col_name_dictionary["column_name"],
-                ):
-                    # Modify the transformed name already present appending a unique key to it
-                    transformed_name = transformed_name + "-#" + str(k)
-                # transform_dict.setdefault(col_name_dictionary["column_name"], []).append( {transformed_name : params_to_be_added} )
-                transform_dict.setdefault(col_name_dictionary["column_name"], {})[transformed_name] = params_to_be_added
-        return transform_dict
-
-    def _generate_cartesian_product_combinations(self, d: dict) -> list:
-        """Helper function for creating cartesian product combinations out of a dictionary.
-        Once all the cartesian product combinations of the values of the dictionary are created it iterates through them
-        to reassign each value to his key. But know the resulting dict has only one value for each key.
-        And the list of this dictionaries is the total set of possible combinations of such values.
-
-        The only other thing that is done is to check if whithin a combination all parameters fiels iside the values are all default.
-        If that is the case the value associated to this combination dict is 1, otherwise is the number of parameters values.
-        This value is used later on in the step of transform handling.
-        for example to know how many time through each combination should the for loop go to select the singular set of parameters values.
-        """
-        keys = d.keys()
-        value_lists = d.values()
-
-        # Generate Cartesian product of value lists
-        combinations = product(*value_lists)
-        # Create dictionaries for each combination
-        result = []
-        for combination in combinations:
-            combined_dict = {}
-            # flag to check if all the parameters values associated to one combination of nopiser are all default
-            all_param_value_default = True
-            for key, value in zip(keys, combination):
-                param_field = d[key][value]
-                nested_dict = {value: param_field}
-                combined_dict.update({key: nested_dict})
-                # now check if the param is a default or not
-                if param_field != "default":
-                    all_param_value_default = False
-            # now append the value to the combo dict that rapresent how many parameters combination there are for such transformeds combination.
-            tmp_tuple = (combined_dict, self.number_column_wise_val)
-            if all_param_value_default:
-                tmp_tuple = (combined_dict, 1)
-            result.append(tmp_tuple)
-
-        return result
-
-    def _handle_parameter_selection(self, d: dict, param_index: int) -> dict:
-        """This function handles the selection of parameters for a given transform dictionary.
-        It takes a dictionary containing transform parameters and an index indicating which parameter combination to select.
-        It iterates through the parameters, extracting the parameter values associated with the given index.
-        It returns a dictionary containing the selected transform name and parameters for the specified index.
-        The output dictionary will have the same structure of the input one but only one value for each paramter instead of a lst of them.
-        """
-        for key, param_dict in d.items():
-            # remove the appendix used to handle same transform names for same column_name, this is done in the _reshape_transform_dict function, this line does nothing if that key is not present afterall
-            key = key.split("-#")[0]
-            # handle "defualt" as params value returning a empty dict
-            if param_dict == "default":
-                return {"name": key, "params": {}}
-            tmp_param_dict = {}
-            # iterate through the possible multiple parameter otpions
-            for param_name, param_value in param_dict.items():
-                tmp_param_dict[param_name] = param_value[param_index]
-            return {"name": key, "params": tmp_param_dict}
-
-    def unique_dicts_in_list(self, dict_list: list) -> list:
-        """Function is pretty straight forwrd: it checks if all elements in a list are unique and returns only the unque ones.
-        This is not a private function because is called from outside this file as well.
-        In the context of this script is mainly used to get the unique dictionaries from a list of dictionaries.
-        It is more general than this but that was the original pourpose.
-        """
-        unique_list = []
-        for d in dict_list:
-            is_unique = True
-            for unique in unique_list:
-                if d == unique:
-                    is_unique = False
-                    break
-            if is_unique:
-                unique_list.append(d)
-        return unique_list
-
-    def transform_column_wise_combination(self) -> list:
-        """Works on the self.transform_arg dictionary to compute all column wise combinations for parametrs and transform function specified.
-        The combinations of transformeds is all against all, except there can not be two transformeds for the same column_name.
-        Combinations of transformeds will always include at least one transformed per column_name.
-        example for transformeds ->
-
-        column_name : 1                                  column_name : 2
-        name : [transformed1, transformed2]                        name: [othertransformed]
-
-        combinations ->
-            transformed1 - othertransformed
-            transformed2 - othertransformed
-
-        Now this is how transformed functions are selected but for each of the above combination there are as many as there are parameters.
-        Again an example shows it better ->
-
-        column_name : 1                                                 column_name : 2
-        name : [transformed1, transformed2]                                       name: [othertransformed]
-        parameters : [{p1 : [1 ,2 ,3]}, {p1 : [1.5, 2.5, 3.5 ]}]        parameters : [{p1 : [4 ,5 ,6], p2 : [7, 8, 9]}]
-
-        combinations ->
-            transformed1 (p1 = 1) - othertransformed (p1 = 4, p2 = 7)
-            transformed1 (p1 = 2) - othertransformed (p1 = 5, p2 = 8)
-            transformed1 (p1 = 3) - othertransformed (p1 = 6, p2 = 9)
-            transformed2 (p1 = 1.5) - othertransformed (p1 = 4, p2 = 7)
-            transformed2 (p1 = 2.5) - othertransformed (p1 = 5, p2 = 8)
-            transformed2 (p1 = 3.5) - othertransformed (p1 = 6, p2 = 9)
-        """
-        # check if there is None  among trasform arguments. if there is return the keyword referring to no transformation.
-        all_transform_combination = []
-        buffer_list = []
-        for transform_argument in self.transform_arg:
-            if transform_argument is None:
-                # add keyword for no transformation
-                all_transform_combination.append(None)
-            else:
-                buffer_list.append(transform_argument)
-
-        # update the trasform arguments, basically removing None values that would throw errors later on in the code
-        self.transform_arg = buffer_list
-
-        # check that no more than one None was added to the all_transform_combination
-        all_transform_combination = self.unique_dicts_in_list(all_transform_combination)
-
-        # reshape transform entry in a nested dictionary, with structure {col_name: { transformed_name : {p1 : [1]} }}
-        transform_as_dict = self._reshape_transform_dict()
-
-        # Create cartesian product of transformed names based on the above dictionary and check if the single combination does not fall under the special case where all parametrs associated to each transformeds in the combination are set to "default". in such a case the code that follows in a specific for loop should be executed only once, instead of self.number_column_wise_val times.
-        transformed_combination_list = self._generate_cartesian_product_combinations(transform_as_dict)
-
-        # for each transformed combination create the column wise selection of parameters associated
-        for transform_combo_tuple in transformed_combination_list:
-            # select the parameter iterating through the total number of parameters associated to the specific transformed combination under selection. This value is the second value of the tuple in which the actual dictionary of transformed combination is.
-            for params_index in range(transform_combo_tuple[1]):
-                transform_list = []
-                # transform_combo_tuple[0] is the actual dictionary with the transformed name parameters for the given combination
-                for col_name, transform_dict in transform_combo_tuple[0].items():
-                    single_param_dict = self._handle_parameter_selection(transform_dict, params_index)
-                    # add the column_name field to this dictionary
-                    single_param_dict["column_name"] = col_name
-                    # reorder the entries by key alphabetically for readability
-                    sorted_dict = {key: single_param_dict[key] for key in sorted(single_param_dict)}
-                    transform_list.append(sorted_dict)
-                all_transform_combination.append(transform_list)
-
-        return all_transform_combination
-
-    def transform_all_combination(self) -> list:
-        """Works on the self.transform_arg dictionary to compute all possible combinations of parameters and nboisers in a all against all fashion."""
-        # TODO implement this function
-        raise ValueError(
-            "the function transform_all_combination for the flag interpret_parmas_mode : all_combinations is not implemented yet ",
-        )
-
-    def split_combination(self) -> list:
-        """This function computes all possible combinations of parameters for splits defined in the schema.
-        It iterates through the split argument in the JSON schema, extracting the name and params fields.
-        It creates separate dictionaries for each parameter combination, ensuring that each splitter has only one value for its parameters.
-        It returns a list of dictionaries, where each dictionary represents a combination of parameters for a split.
-        """
-        # check if there is None  among trasform arguments. if there is return the keyword referring to no transformation.
-        list_split_comibinations = []
-        buffer_list = []
-        for split_argument in self.split_arg:
-            if split_argument is None:
-                # add keyword  for no split
-                list_split_comibinations.append(None)
-            else:
-                buffer_list.append(split_argument)
-
-        # update the split arguments, basically removing None values that would throw errors later on in the code
-        self.split_arg = buffer_list
-
-        # check that no more than one None was added to the all_transform_combination
-        list_split_comibinations = self.unique_dicts_in_list(list_split_comibinations)
-
-        # iterate through the split entry and return a list of split possibilities, where each splitter_name has one/set of one parametyers
-        for i, split_dict in enumerate(self.split_arg):
-            # jsut create a new dictionary for each set of params associated to each split_name, basically if a splitter has more than one element in his params: then they should be decoupled so to have each splitter with only one value for params:
-            # if the value of params: is "default" just return the dictionary  with an empty dict as value of params :
-            if split_dict["params"] == "default" or split_dict["params"] == ["default"]:
-                split_dict["params"] = {}
-                list_split_comibinations.append(split_dict)
-            else:
-                # Get lengths of all lists
-                lengths = {key: len(value) for key, value in split_dict["params"][0].items()}
-
-                # Check if all lengths are the same
-                all_lengths_same = set(lengths.values())
-
-                if len(all_lengths_same) != 1:
-                    raise ValueError(
-                        f"All split params for the same splitter have to have the same number of elements, this splitter does not: {split_dict['name']}.",
-                    )
-                # iterate at level of number of params_values
-                for params_index in range(list(all_lengths_same)[0]):
-                    # making the split into a dict the _handle_parameter_selection can use
-                    single_param_dict = self._handle_parameter_selection(
-                        {split_dict["name"]: split_dict["params"][0]},
-                        params_index,
-                    )
-                    list_split_comibinations.append(single_param_dict)
-        return list_split_comibinations

From cdb121fd8fe50b7fc5ebe7ebbad859228d2ade35 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 14:46:11 +0100
Subject: [PATCH 02/14] LINT: fixed litting for utils

---
 src/stimulus/utils/launch_utils.py |  42 +++++++--
 src/stimulus/utils/performance.py  |  47 ++++++++--
 src/stimulus/utils/yaml_data.py    | 133 ++++++++++++++---------------
 3 files changed, 138 insertions(+), 84 deletions(-)

diff --git a/src/stimulus/utils/launch_utils.py b/src/stimulus/utils/launch_utils.py
index fdaaaa14..24b12b96 100644
--- a/src/stimulus/utils/launch_utils.py
+++ b/src/stimulus/utils/launch_utils.py
@@ -1,12 +1,25 @@
+"""Utility functions for launching and configuring experiments and ray tuning."""
+
 import importlib.util
 import math
 import os
-from typing import Tuple, Union
+from typing import Union
 
 import stimulus.data.experiments as exp
 
 
 def import_class_from_file(file_path: str) -> type:
+    """Import and return the Model class from a specified Python file.
+
+    Args:
+        file_path (str): Path to the Python file containing the Model class.
+
+    Returns:
+        type: The Model class found in the file.
+
+    Raises:
+        ImportError: If no class starting with 'Model' is found in the file.
+    """
     # Extract directory path and file name
     directory, file_name = os.path.split(file_path)
     module_name = os.path.splitext(file_name)[0]  # Remove extension to get module name
@@ -28,15 +41,30 @@ def import_class_from_file(file_path: str) -> type:
 
 
 def get_experiment(experiment_name: str) -> object:
-    experiment_object = getattr(exp, experiment_name)()
-    return experiment_object
+    """Get an experiment instance by name.
 
+    Args:
+        experiment_name (str): Name of the experiment class to instantiate.
 
-def memory_split_for_ray_init(memory_str: Union[str, None]) -> Tuple[float, float]:
+    Returns:
+        object: An instance of the requested experiment class.
+    """
+    return getattr(exp, experiment_name)()
+
+
+def memory_split_for_ray_init(memory_str: Union[str, None]) -> tuple[float, float]:
     """Process the input memory value into the right unit and allocates 30% for overhead and 70% for tuning.
-    Usefull in case ray detects them wrongly.
-    Memory is split in two for ray: for store_object memory and the other actual memory for tuning.
-    The following function takes the total possible usable/allocated memory as a string parameter and returns in bytes the values for store_memory (30% as default in ray) and memory (70%).
+
+    Useful in case ray detects them wrongly. Memory is split in two for ray: for store_object memory
+    and the other actual memory for tuning. The following function takes the total possible
+    usable/allocated memory as a string parameter and returns in bytes the values for store_memory
+    (30% as default in ray) and memory (70%).
+
+    Args:
+        memory_str (Union[str, None]): Memory string in format like "8G", "16GB", etc.
+
+    Returns:
+        tuple[float, float]: A tuple containing (store_memory, memory) in bytes.
     """
     if memory_str is None:
         return None, None
diff --git a/src/stimulus/utils/performance.py b/src/stimulus/utils/performance.py
index b3bc193f..c297bd79 100644
--- a/src/stimulus/utils/performance.py
+++ b/src/stimulus/utils/performance.py
@@ -1,4 +1,6 @@
-from typing import Any, Tuple
+"""Utility module for computing various performance metrics for machine learning models."""
+
+from typing import Any
 
 import numpy as np
 import torch
@@ -12,9 +14,13 @@
     roc_auc_score,
 )
 
+# Constants for threshold and number of classes
+BINARY_THRESHOLD = 0.5
+BINARY_CLASS_COUNT = 2
+
 
 class Performance:
-    """Returns the value of a given metric
+    """Returns the value of a given metric.
 
     Parameters
     ----------
@@ -36,6 +42,13 @@ class Performance:
     """
 
     def __init__(self, labels: Any, predictions: Any, metric: str = "rocauc") -> float:
+        """Initialize Performance class with labels, predictions and metric type.
+
+        Args:
+            labels: Ground truth labels
+            predictions: Model predictions
+            metric: Type of metric to compute (default: "rocauc")
+        """
         labels = self.data2array(labels)
         predictions = self.data2array(predictions)
         labels, predictions = self.handle_multiclass(labels, predictions)
@@ -47,6 +60,17 @@ def __init__(self, labels: Any, predictions: Any, metric: str = "rocauc") -> flo
         self.val = function(labels, predictions)
 
     def data2array(self, data: Any) -> np.array:
+        """Convert input data to numpy array.
+
+        Args:
+            data: Input data in various formats
+
+        Returns:
+            np.array: Converted numpy array
+
+        Raises:
+            ValueError: If input data type is not supported
+        """
         if isinstance(data, list):
             return np.array(data)
         if isinstance(data, np.ndarray):
@@ -57,7 +81,7 @@ def data2array(self, data: Any) -> np.array:
             return np.array([data])
         raise ValueError(f"The data must be a list, np.array, torch.Tensor, int or float. Instead it is {type(data)}")
 
-    def handle_multiclass(self, labels: np.array, predictions: np.array) -> Tuple[np.array, np.array]:
+    def handle_multiclass(self, labels: np.array, predictions: np.array) -> tuple[np.array, np.array]:
         """Handle the case of multiclass classification.
 
         TODO currently only two class predictions are handled. Needs to handle the other scenarios.
@@ -67,7 +91,7 @@ def handle_multiclass(self, labels: np.array, predictions: np.array) -> Tuple[np
             return labels, predictions
 
         # if one columns for labels, but two columns for predictions
-        if (len(labels.shape) == 1) and (predictions.shape[1] == 2):
+        if (len(labels.shape) == 1) and (predictions.shape[1] == BINARY_CLASS_COUNT):
             predictions = predictions[:, 1]  # assumes the second column is the positive class
             return labels, predictions
 
@@ -75,26 +99,33 @@ def handle_multiclass(self, labels: np.array, predictions: np.array) -> Tuple[np
         raise ValueError(f"Labels have shape {labels.shape} and predictions have shape {predictions.shape}.")
 
     def rocauc(self, labels: np.array, predictions: np.array) -> float:
+        """Compute ROC AUC score."""
         return roc_auc_score(labels, predictions)
 
     def prauc(self, labels: np.array, predictions: np.array) -> float:
+        """Compute PR AUC score."""
         return average_precision_score(labels, predictions)
 
     def mcc(self, labels: np.array, predictions: np.array) -> float:
-        predictions = np.array([1 if p > 0.5 else 0 for p in predictions])
+        """Compute Matthews Correlation Coefficient."""
+        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
         return matthews_corrcoef(labels, predictions)
 
     def f1score(self, labels: np.array, predictions: np.array) -> float:
-        predictions = np.array([1 if p > 0.5 else 0 for p in predictions])
+        """Compute F1 score."""
+        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
         return f1_score(labels, predictions)
 
     def precision(self, labels: np.array, predictions: np.array) -> float:
-        predictions = np.array([1 if p > 0.5 else 0 for p in predictions])
+        """Compute precision score."""
+        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
         return precision_score(labels, predictions)
 
     def recall(self, labels: np.array, predictions: np.array) -> float:
-        predictions = np.array([1 if p > 0.5 else 0 for p in predictions])
+        """Compute recall score."""
+        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
         return recall_score(labels, predictions)
 
     def spearmanr(self, labels: np.array, predictions: np.array) -> float:
+        """Compute Spearman correlation coefficient."""
         return spearmanr(labels, predictions)[0]
diff --git a/src/stimulus/utils/yaml_data.py b/src/stimulus/utils/yaml_data.py
index ed8615ce..0387c77a 100644
--- a/src/stimulus/utils/yaml_data.py
+++ b/src/stimulus/utils/yaml_data.py
@@ -1,52 +1,72 @@
-from typing import Dict, List, Optional, Union
+"""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]]]  # Allow both string and list values
+    params: Optional[dict[str, Union[str, list]]]  # 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]
+    encoder: list[YamlColumnsEncoder]
 
 
 class YamlTransformColumnsTransformation(BaseModel):
+    """Model for column transformation configuration."""
     name: str
-    params: Optional[Dict[str, Union[list, float]]]  # Allow both list and float values
+    params: Optional[dict[str, Union[list, float]]]  # Allow both list and float values
 
 
 class YamlTransformColumns(BaseModel):
+    """Model for transform columns configuration."""
     column_name: str
-    transformations: List[YamlTransformColumnsTransformation]
+    transformations: list[YamlTransformColumnsTransformation]
 
 
 class YamlTransform(BaseModel):
+    """Model for transform configuration."""
     transformation_name: str
-    columns: List[YamlTransformColumns]
+    columns: list[YamlTransformColumns]
 
     @field_validator("columns")
     @classmethod
-    def validate_param_lists_across_columns(cls, columns) -> List[YamlTransformColumns]:
+    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()
 
         for column in columns:
             for transformation in column.transformations:
-                if transformation.params:
-                    for param_value in transformation.params.values():
-                        if isinstance(param_value, list):
-                            if len(param_value) > 0:  # Non-empty list
-                                all_list_lengths.add(len(param_value))
+                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:
@@ -54,7 +74,7 @@ def validate_param_lists_across_columns(cls, columns) -> List[YamlTransformColum
 
         # 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, since sets do not allow duplicates
+        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",
             )
@@ -63,26 +83,30 @@ def validate_param_lists_across_columns(cls, columns) -> List[YamlTransformColum
 
 
 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]
+    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]
+    columns: list[YamlColumns]
+    transforms: list[YamlTransform]
+    split: list[YamlSplit]
 
 
 class YamlSubConfigDict(BaseModel):
+    """Model for sub-configuration generated from main config."""
     global_params: YamlGlobalParams
-    columns: List[YamlColumns]
+    columns: list[YamlColumns]
     transforms: YamlTransform
     split: YamlSplit
 
 
 class YamlSchema(BaseModel):
+    """Model for validating YAML schema."""
     yaml_conf: YamlConfigDict
 
 
@@ -216,84 +240,51 @@ def dump_yaml_list_into_files(
     directory_path: str,
     base_name: str,
 ) -> None:
-    """Dumps a list of YAML configurations into separate files with custom formatting.
-
-    This function takes a list of YAML configurations and writes each one to a separate file,
-    applying custom formatting rules to ensure consistent and readable output. It handles
-    special cases like None values, nested lists, and proper indentation.
-
-    Args:
-        yaml_list: List of YamlSubConfigDict objects to be written to files
-        directory_path: Directory path where the files should be created
-        base_name: Base name for the output files. Files will be named {base_name}_{index}.yaml
-
-    The function applies several custom formatting rules:
-    - None values are represented as empty strings
-    - Nested lists use appropriate flow styles based on content type
-    - Extra newlines are added between root-level elements
-    - Proper indentation is maintained throughout
-    """
+    """Dumps a list of YAML configurations into separate files with custom formatting."""
     # Disable YAML aliases to prevent reference-style output
     yaml.Dumper.ignore_aliases = lambda *args: True
 
-    def represent_none(dumper, _):
+    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, data):
-        """Custom representer to handle different types of lists with appropriate formatting.
-
-        Applies different flow styles based on content:
-        - Empty lists -> empty string
-        - Lists of dicts (e.g. columns) -> block style (vertical)
-        - Lists of lists (e.g. split params) -> flow style (inline)
-        - Other lists -> flow style
-        """
+    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):
-                # Use block style for structured data like columns
                 return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=False)
             if isinstance(data[0], list):
-                # Use flow style for numeric data like split ratios
                 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 write_line_break(self, data=None):
+        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(self, flow=False, indentless=False):
+        def increase_indent(self, *, flow: bool = False) -> bool:
             """Ensure consistent indentation by preventing indentless sequences."""
-            return super().increase_indent(flow, False)
+            return super().increase_indent(flow=flow, indentless=False)
 
     # 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):
-        # Convert Pydantic model to dict, excluding None values
         dict_data = yaml_dict.model_dump(exclude_none=True)
 
-        def fix_params(input_dict):
-            """Recursively process dictionary to properly handle params fields.
-
-            Special handling for:
-            - Empty params fields -> empty dict (will be represented as {} in YAML)
-            - Transformation params -> empty dict if empty
-            - Nested dicts and lists -> recursive processing
-            """
+        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 = {}
                 for key, value in input_dict.items():
                     if key == "encoder" and isinstance(value, list):
-                        # Ensure each encoder has a params field
                         processed_dict[key] = []
                         for encoder in value:
                             processed_encoder = dict(encoder)
@@ -310,9 +301,9 @@ def fix_params(input_dict):
                     elif isinstance(value, dict):
                         processed_dict[key] = fix_params(value)
                     elif isinstance(value, list):
-                        # Process lists, recursing into dict elements
                         processed_dict[key] = [
-                            fix_params(list_item) if isinstance(list_item, dict) else list_item for list_item in value
+                            fix_params(list_item) if isinstance(list_item, dict) else list_item
+                            for list_item in value
                         ]
                     else:
                         processed_dict[key] = value
@@ -334,19 +325,23 @@ def fix_params(input_dict):
 
 
 def check_yaml_schema(config_yaml: str) -> str:
-    """Using pydantic this function confirms that the fields have the correct input type
-    If the children field is specific to a parent, the children fields class is hosted in the parent fields class
+    """Validate YAML configuration fields have correct types.
 
-    If any field in not the right type, the function prints an error message explaining the problem and exits the python code
+    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 (dict): The dict containing the fields of the yaml configuration file
 
     Returns:
-        None
+        str: Empty string if validation succeeds
+
+    Raises:
+        ValueError: If validation fails
     """
     try:
         YamlSchema(yaml_conf=config_yaml)
     except ValidationError as e:
-        print(e)
-        raise ValueError("Wrong type on a field, see the pydantic report above")  # Crashes in case of an error raised
+        # Use logging instead of print for error handling
+        raise ValueError("Wrong type on a field, see the pydantic report above") from e
+    return ""

From cadf8a80ec9ef67be34ba9f9b01115daea4fbe76 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:20:05 +0100
Subject: [PATCH 03/14] fixing more linting issues

---
 src/stimulus/utils/yaml_model_schema.py | 111 +++++++--
 tests/data/encoding/__init__.py         |   1 +
 tests/data/encoding/test_encoders.py    | 313 +++++++++++++++---------
 3 files changed, 286 insertions(+), 139 deletions(-)
 create mode 100644 tests/data/encoding/__init__.py

diff --git a/src/stimulus/utils/yaml_model_schema.py b/src/stimulus/utils/yaml_model_schema.py
index 1e6b05c6..f6b68619 100644
--- a/src/stimulus/utils/yaml_model_schema.py
+++ b/src/stimulus/utils/yaml_model_schema.py
@@ -1,3 +1,5 @@
+"""Module for handling YAML configuration files and converting them to Ray Tune format."""
+
 import random
 from collections.abc import Callable
 from copy import deepcopy
@@ -7,14 +9,35 @@
 
 
 class YamlRayConfigLoader:
-    def __init__(self, config_path: str):
+    """Load and convert YAML configurations to Ray Tune format.
+
+    This class handles loading YAML configuration files and converting them into
+    formats compatible with Ray Tune's hyperparameter search spaces.
+    """
+
+    def __init__(self, config_path: str) -> None:
+        """Initialize the config loader with a YAML file path.
+
+        Args:
+            config_path: Path to the YAML configuration file
+        """
         with open(config_path) as f:
             self.config = yaml.safe_load(f)
             self.config = self.convert_config_to_ray(self.config)
 
     def raytune_space_selector(self, mode: Callable, space: list) -> Callable:
-        # this function applies the mode function to the space, it needs to convert the space in a right way based on the mode, for instance, if the mode is "randint", the space should be a tuple of two integers and passed as *args
+        """Convert space parameters to Ray Tune format based on the mode.
+
+        Args:
+            mode: Ray Tune search space function (e.g., tune.choice, tune.uniform)
+            space: List of parameters defining the search space
+
+        Returns:
+            Configured Ray Tune search space
 
+        Raises:
+            NotImplementedError: If the mode is not supported
+        """
         if mode.__name__ == "choice":
             return mode(space)
 
@@ -24,48 +47,72 @@ def raytune_space_selector(self, mode: Callable, space: list) -> Callable:
         raise NotImplementedError(f"Mode {mode.__name__} not implemented yet")
 
     def raytune_sample_from(self, mode: Callable, param: dict) -> Callable:
-        """This function applies the tune.sample_from to a given custom sampling function."""
+        """Apply tune.sample_from to a given custom sampling function.
+
+        Args:
+            mode: Ray Tune sampling function
+            param: Dictionary containing sampling parameters
+
+        Returns:
+            Configured sampling function
+
+        Raises:
+            NotImplementedError: If the sampling function is not supported
+        """
         if param["function"] == "sampint":
             return mode(lambda _: self.sampint(param["sample_space"], param["n_space"]))
 
         raise NotImplementedError(f"Function {param['function']} not implemented yet")
 
     def convert_raytune(self, param: dict) -> dict:
-        # get the mode function from ray.tune using getattr, return an error if it is not recognized
+        """Convert parameter configuration to Ray Tune format.
+
+        Args:
+            param: Parameter configuration dictionary
+
+        Returns:
+            Ray Tune compatible parameter configuration
+
+        Raises:
+            AttributeError: If the mode is not recognized in Ray Tune
+        """
         try:
             mode = getattr(tune, param["mode"])
-        except AttributeError:
+        except AttributeError as err:
             raise AttributeError(
                 f"Mode {param['mode']} not recognized, check the ray.tune documentation at https://docs.ray.io/en/master/tune/api_docs/suggestion.html",
-            )
+            ) from err
 
-        # apply the mode function
         if param["mode"] != "sample_from":
             return self.raytune_space_selector(mode, param["space"])
         return self.raytune_sample_from(mode, param)
 
     def convert_config_to_ray(self, config: dict) -> dict:
-        # the config is a dictionary of dictionaries. The main dictionary keys are either model_params, loss_params or optimizer_params.
-        # The sub-dictionary keys are the parameters of the model, loss or optimizer, those params include two values, space and mode.
-        # The space is the range of values to be tested, and the mode is the type of search to be done.
-        # We convert the Yaml config by calling the correct function from ray.tune matching the mode, applied on the space
-        # We return the config as a dictionary of dictionaries, where the values are the converted values from the space.
+        """Convert YAML configuration to Ray Tune format.
+
+        Converts parameters in model_params, loss_params, optimizer_params, and data_params
+        to Ray Tune search spaces when a mode is specified.
+
+        Args:
+            config: Raw configuration dictionary from YAML
+
+        Returns:
+            Ray Tune compatible configuration dictionary
+        """
         new_config = deepcopy(config)
         for key in ["model_params", "loss_params", "optimizer_params", "data_params"]:
             for sub_key in config[key]:
-                # if mode is provided, it understands that it is a ray.tune parameter
-                # therefore, it converts the space provided in the config to a ray.tune parameter space
-                # otherwise, it keeps the value as it is. In this way, we can use the same config for both ray.tune and non-ray.tune elements (for example provide a single fixed value).
                 if "mode" in config[key][sub_key]:
                     new_config[key][sub_key] = self.convert_raytune(config[key][sub_key])
 
         return new_config
 
     def get_config_instance(self) -> dict:
-        # this function take a config as input and returns an instance of said config with the values sampled from the space
-        # the config is a dictionary of dictionaries. The main dictionary keys are either model_params, loss_params or optimizer_params.
-        # The sub-dictionary keys are the parameters of the model, loss or optimizer, those params include two values, space and mode.
+        """Generate a configuration instance with sampled values.
 
+        Returns:
+            Configuration dictionary with concrete sampled values
+        """
         config_instance = deepcopy(self.config)
         for key in ["model_params", "loss_params", "optimizer_params", "data_params"]:
             config_instance[key] = {}
@@ -75,19 +122,33 @@ def get_config_instance(self) -> dict:
         return config_instance
 
     def get_config(self) -> dict:
+        """Return the current configuration.
+
+        Returns:
+            Current configuration dictionary
+        """
         return self.config
 
     @staticmethod
     def sampint(sample_space: list, n_space: list) -> list:
-        """This function returns a list of n samples from the sample_space.
+        """Return a list of n random samples from the sample_space.
+
+        This function is useful for sampling different numbers of layers,
+        each with different numbers of neurons.
+
+        Args:
+            sample_space: List [min, max] defining range of values to sample from
+            n_space: List [min, max] defining range for number of samples
 
-        This function is specially useful when we want different number of layers,
-        and each layer with different number of neurons.
+        Returns:
+            List of randomly sampled integers
 
-        `sample_space` is the range of (int) values from which to sample
-        `n_space` is the range of (int) number of samples to take
+        Note:
+            Uses Python's random module which is not cryptographically secure.
+            This is acceptable for hyperparameter sampling but should not be
+            used for security-critical purposes (S311 fails when linting).
         """
         sample_space = range(sample_space[0], sample_space[1] + 1)
         n_space = range(n_space[0], n_space[1] + 1)
-        n = random.choice(n_space)
-        return random.sample(sample_space, n)
+        n = random.choice(tuple(n_space))  # noqa: S311
+        return random.sample(tuple(sample_space), n)
diff --git a/tests/data/encoding/__init__.py b/tests/data/encoding/__init__.py
new file mode 100644
index 00000000..2266658e
--- /dev/null
+++ b/tests/data/encoding/__init__.py
@@ -0,0 +1 @@
+"""Encoding tests package."""
\ No newline at end of file
diff --git a/tests/data/encoding/test_encoders.py b/tests/data/encoding/test_encoders.py
index d4df5ec8..f8a0a58c 100644
--- a/tests/data/encoding/test_encoders.py
+++ b/tests/data/encoding/test_encoders.py
@@ -16,24 +16,33 @@ class TestTextOneHotEncoder:
 
     @staticmethod
     @pytest.fixture
-    def encoder_default():
-        """Provides a default encoder."""
+    def encoder_default() -> TextOneHotEncoder:
+        """Provide a default encoder.
+
+        Returns:
+            TextOneHotEncoder: A default encoder instance
+        """
         return TextOneHotEncoder(alphabet="acgt", padding=True)
 
     @staticmethod
     @pytest.fixture
-    def encoder_lowercase():
-        """Provides an encoder with convert_lowercase set to True."""
+    def encoder_lowercase() -> TextOneHotEncoder:
+        """Provide an encoder with convert_lowercase set to True.
+
+        Returns:
+            TextOneHotEncoder: An encoder instance with lowercase conversion
+        """
         return TextOneHotEncoder(alphabet="ACgt", convert_lowercase=True, padding=True)
 
     # ---- Test for initialization ---- #
 
-    def test_init_with_non_string_alphabet_raises_type_error(self):
-        with pytest.raises(TypeError) as excinfo:
-            TextOneHotEncoder(alphabet=["a", "c", "g", "t"])  # Passing a list instead of string
-        assert "Expected a string input for alphabet" in str(excinfo.value)
+    def test_init_with_non_string_alphabet_raises_type_error(self) -> None:
+        """Test initialization with non-string alphabet raises TypeError."""
+        with pytest.raises(TypeError, match="Expected a string input for alphabet"):
+            TextOneHotEncoder(alphabet=["a", "c", "g", "t"])
 
-    def test_init_with_string_alphabet(self):
+    def test_init_with_string_alphabet(self) -> None:
+        """Test initialization with valid string alphabet."""
         encoder = TextOneHotEncoder(alphabet="acgt")
         assert encoder.alphabet == "acgt"
         assert encoder.convert_lowercase is False
@@ -41,46 +50,55 @@ def test_init_with_string_alphabet(self):
 
     # ---- Tests for _sequence_to_array ---- #
 
-    def test_sequence_to_array_with_non_string_input(self, encoder_default):
-        with pytest.raises(TypeError) as excinfo:
+    def test_sequence_to_array_with_non_string_input(
+        self, encoder_default: TextOneHotEncoder,
+    ) -> None:
+        """Test _sequence_to_array with non-string input raises TypeError."""
+        with pytest.raises(TypeError, match="Expected string input for sequence"):
             encoder_default._sequence_to_array(1234)
-        assert "Expected string input for sequence" in str(excinfo.value)
 
-    def test_sequence_to_array_returns_correct_shape(self, encoder_default):
+    def test_sequence_to_array_returns_correct_shape(
+        self, encoder_default: TextOneHotEncoder,
+    ) -> None:
+        """Test _sequence_to_array returns array of correct shape."""
         seq = "acgt"
         arr = encoder_default._sequence_to_array(seq)
-        # shape should be (len(seq), 1)
         assert arr.shape == (4, 1)
-        # check content
         assert (arr.flatten() == list(seq)).all()
 
-    def test_sequence_to_array_is_case_sensitive(self, encoder_default):
+    def test_sequence_to_array_is_case_sensitive(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test that _sequence_to_array preserves case when case sensitivity is enabled."""
         seq = "AcGT"
         arr = encoder_default._sequence_to_array(seq)
         assert (arr.flatten() == list("AcGT")).all()
 
-    def test_sequence_to_array_is_lowercase(self, encoder_lowercase):
+    def test_sequence_to_array_is_lowercase(self, encoder_lowercase: TextOneHotEncoder) -> None:
+        """Test that _sequence_to_array converts to lowercase when enabled."""
         seq = "AcGT"
         arr = encoder_lowercase._sequence_to_array(seq)
         assert (arr.flatten() == list("acgt")).all()
 
     # ---- Tests for encode ---- #
 
-    def test_encode_returns_tensor(self, encoder_default):
+    def test_encode_returns_tensor(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test that encode returns a tensor of the correct shape."""
         seq = "acgt"
         encoded = encoder_default.encode(seq)
         assert isinstance(encoded, torch.Tensor)
         # shape should be (len(seq), alphabet_size=4)
         assert encoded.shape == (4, 4)
 
-    def test_encode_unknown_character_returns_zero_vector(self, encoder_default):
+    def test_encode_unknown_character_returns_zero_vector(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test that encoding an unknown character returns a zero vector."""
         seq = "acgtn"
         encoded = encoder_default.encode(seq)
         # the last character 'n' is not in 'acgt', so the last row should be all zeros
         assert torch.all(encoded[-1] == 0)
 
-    def test_encode_default(self, encoder_default):
-        """Case-sensitive: 'ACgt' => 'ACgt' means 'A' and 'C' are uppercase in the alphabet,
+    def test_encode_default(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test case-sensitive encoding behavior.
+
+        Case-sensitive: 'ACgt' => 'ACgt' means 'A' and 'C' are uppercase in the alphabet,
         'g' and 't' are lowercase in the alphabet.
         """
         seq = "ACgt"
@@ -94,7 +112,7 @@ def test_encode_default(self, encoder_default):
         assert torch.all(encoded[2] == torch.tensor([0, 0, 1, 0]))  # 'g'
         assert torch.all(encoded[3] == torch.tensor([0, 0, 0, 1]))  # 't'
 
-    def test_encode_lowercase(self, encoder_lowercase):
+    def test_encode_lowercase(self, encoder_lowercase: TextOneHotEncoder) -> None:
         """Case-insensitive: 'ACgt' => 'acgt' internally."""
         seq = "ACgt"
         encoded = encoder_lowercase.encode(seq)
@@ -108,14 +126,16 @@ def test_encode_lowercase(self, encoder_lowercase):
 
     # ---- Tests for encode_all ---- #
 
-    def test_encode_all_with_single_string(self, encoder_default):
+    def test_encode_all_with_single_string(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test encoding a single string with encode_all."""
         seq = "acgt"
         encoded = encoder_default.encode_all(seq)
         # shape = (batch_size=1, seq_len=4, alphabet_size=4)
         assert encoded.shape == (1, 4, 4)
         assert torch.all(encoded[0] == encoder_default.encode(seq))
 
-    def test_encode_all_with_list_of_sequences(self, encoder_default):
+    def test_encode_all_with_list_of_sequences(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test encoding multiple sequences with encode_all."""
         seqs = ["acgt", "acgtn"]  # second has an unknown 'n'
         encoded = encoder_default.encode_all(seqs)
         # shape = (2, max_len=5, alphabet_size=4)
@@ -124,17 +144,18 @@ def test_encode_all_with_list_of_sequences(self, encoder_default):
         assert torch.all(encoded[0][:4] == encoder_default.encode(seqs[0]))
         assert torch.all(encoded[1] == encoder_default.encode(seqs[1]))
 
-    def test_encode_all_with_padding_false(self):
+    def test_encode_all_with_padding_false(self) -> None:
+        """Test that encode_all raises error when padding is False and sequences have different lengths."""
         encoder = TextOneHotEncoder(alphabet="acgt", padding=False)
         seqs = ["acgt", "acgtn"]  # different lengths
         # should raise ValueError because lengths differ
-        with pytest.raises(ValueError) as excinfo:
+        with pytest.raises(ValueError, match="All sequences must have the same length when padding is False."):
             encoder.encode_all(seqs)
-        assert "All sequences must have the same length when padding is False." in str(excinfo.value)
 
     # ---- Tests for decode ---- #
 
-    def test_decode_single_sequence(self, encoder_default):
+    def test_decode_single_sequence(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test decoding a single encoded sequence."""
         seq = "acgt"
         encoded = encoder_default.encode(seq)
         decoded = encoder_default.decode(encoded)
@@ -143,8 +164,10 @@ def test_decode_single_sequence(self, encoder_default):
         # Should match the lowercased input (since case-sensitive=False)
         assert decoded == seq
 
-    def test_decode_unknown_characters(self, encoder_default):
-        """Unknown characters are zero vectors. When decoding, those become empty (ignored),
+    def test_decode_unknown_characters(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test decoding behavior with unknown characters.
+
+        Unknown characters are zero vectors. When decoding, those become empty (ignored),
         or become None, depending on the transform. In the provided code, handle_unknown='ignore'
         yields an empty decode for those positions. The example code attempts to fill with '-'
         or None if needed.
@@ -162,7 +185,8 @@ def test_decode_unknown_characters(self, encoder_default):
         # Let's do a partial check:
         assert decoded.startswith("acgt")
 
-    def test_decode_multiple_sequences(self, encoder_default):
+    def test_decode_multiple_sequences(self, encoder_default: TextOneHotEncoder) -> None:
+        """Test decoding multiple encoded sequences."""
         seqs = ["acgt", "acgtn"]  # second has unknown 'n'
         encoded = encoder_default.encode_all(seqs)
         decoded = encoder_default.decode(encoded)
@@ -178,26 +202,34 @@ class TestNumericEncoder:
 
     @staticmethod
     @pytest.fixture
-    def float_encoder():
-        """Fixture to instantiate the NumericEncoder."""
+    def float_encoder() -> NumericEncoder:
+        """Provide a NumericEncoder instance.
+
+        Returns:
+            NumericEncoder: Default encoder instance
+        """
         return NumericEncoder()
 
     @staticmethod
     @pytest.fixture
-    def int_encoder():
-        """Fixture to instantiate the NumericEncoder with integer dtype."""
+    def int_encoder() -> NumericEncoder:
+        """Provide a NumericEncoder instance with integer dtype.
+
+        Returns:
+            NumericEncoder: Integer-based encoder instance
+        """
         return NumericEncoder(dtype=torch.int32)
 
-    def test_encode_single_float(self, float_encoder):
+    def test_encode_single_float(self, float_encoder: NumericEncoder) -> None:
         """Test encoding a single float value."""
         input_val = 3.14
         output = float_encoder.encode(input_val)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
         assert output.dtype == torch.float32, "Tensor dtype should be float32."
         assert output.numel() == 1, "Tensor should have exactly one element."
-        assert output.item() == pytest.approx(input_val), "Encoded value does not match the input float."
+        assert output.item() == pytest.approx(input_val), "Encoded value does not match."
 
-    def test_encode_single_int(self, int_encoder):
+    def test_encode_single_int(self, int_encoder: NumericEncoder) -> None:
         """Test encoding a single int value."""
         input_val = 3
         output = int_encoder.encode(input_val)
@@ -207,38 +239,43 @@ def test_encode_single_int(self, int_encoder):
         assert output.item() == input_val
 
     @pytest.mark.parametrize("fixture_name", ["float_encoder", "int_encoder"])
-    def test_encode_non_numeric_raises(self, request, fixture_name):
+    def test_encode_non_numeric_raises(
+        self, request: pytest.FixtureRequest, fixture_name: str,
+    ) -> None:
         """Test that encoding a non-float raises a ValueError."""
         numeric_encoder = request.getfixturevalue(fixture_name)
-        with pytest.raises(ValueError) as exc_info:
+        with pytest.raises(ValueError, match="Expected input data to be a float or int"):
             numeric_encoder.encode("not_numeric")
-        assert "Expected input data to be a float or int" in str(
-            exc_info.value,
-        ), "Expected ValueError with specific error message."
 
-    def test_encode_all_single_float(self, float_encoder):
+    def test_encode_all_single_float(self, float_encoder: NumericEncoder) -> None:
         """Test encode_all when given a single float.
-        It should be treated as a list of one float internally.
+
+        Tests that a single float is treated as a list of one float internally.
+
+        Args:
+            float_encoder: Float-based encoder instance
         """
         input_val = 2.71
         output = float_encoder.encode_all(input_val)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
-        assert output.dtype == torch.float32, "Tensor dtype should be float32."
         assert output.numel() == 1, "Tensor should have exactly one element."
         assert output.item() == pytest.approx(input_val), "Encoded value does not match the input."
 
-    def test_encode_all_single_int(self, int_encoder):
+    def test_encode_all_single_int(self, int_encoder: NumericEncoder) -> None:
         """Test encode_all when given a single int.
-        It should be treated as a list of one int internally.
+
+        Tests that a single int is treated as a list of one int internally.
+
+        Args:
+            int_encoder: Integer-based encoder instance
         """
         input_val = 2
         output = int_encoder.encode_all(input_val)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
-        assert output.dtype == torch.int32, "Tensor dtype should be int32."
         assert output.numel() == 1, "Tensor should have exactly one element."
         assert output.item() == input_val
 
-    def test_encode_all_multi_float(self, float_encoder):
+    def test_encode_all_multi_float(self, float_encoder: NumericEncoder) -> None:
         """Test encode_all with a list of floats."""
         input_vals = [3.14, 4.56]
         output = float_encoder.encode_all(input_vals)
@@ -248,7 +285,7 @@ def test_encode_all_multi_float(self, float_encoder):
         assert output[0].item() == pytest.approx(3.14), "First element does not match."
         assert output[1].item() == pytest.approx(4.56), "Second element does not match."
 
-    def test_encode_all_multi_int(self, int_encoder):
+    def test_encode_all_multi_int(self, int_encoder: NumericEncoder) -> None:
         """Test encode_all with a list of integers."""
         input_vals = [3, 4]
         output = int_encoder.encode_all(input_vals)
@@ -258,7 +295,7 @@ def test_encode_all_multi_int(self, int_encoder):
         assert output[0].item() == 3, "First element does not match."
         assert output[1].item() == 4, "Second element does not match."
 
-    def test_decode_single_float(self, float_encoder):
+    def test_decode_single_float(self, float_encoder: NumericEncoder) -> None:
         """Test decoding a tensor of shape (1)."""
         input_tensor = torch.tensor([3.14], dtype=torch.float32)
         decoded = float_encoder.decode(input_tensor)
@@ -267,7 +304,7 @@ def test_decode_single_float(self, float_encoder):
         assert len(decoded) == 1, "Decoded list should have one element."
         assert decoded[0] == pytest.approx(3.14), "Decoded value does not match."
 
-    def test_decode_single_int(self, int_encoder):
+    def test_decode_single_int(self, int_encoder: NumericEncoder) -> None:
         """Test decoding a tensor of shape (1)."""
         input_tensor = torch.tensor([3], dtype=torch.int32)
         decoded = int_encoder.decode(input_tensor)
@@ -276,7 +313,7 @@ def test_decode_single_int(self, int_encoder):
         assert len(decoded) == 1, "Decoded list should have one element."
         assert decoded[0] == 3, "Decoded value does not match."
 
-    def test_decode_multi_float(self, float_encoder):
+    def test_decode_multi_float(self, float_encoder: NumericEncoder) -> None:
         """Test decoding a tensor of shape (n)."""
         input_tensor = torch.tensor([3.14, 2.71], dtype=torch.float32)
         decoded = float_encoder.decode(input_tensor)
@@ -285,7 +322,7 @@ def test_decode_multi_float(self, float_encoder):
         assert decoded[0] == pytest.approx(3.14), "First decoded value does not match."
         assert decoded[1] == pytest.approx(2.71), "Second decoded value does not match."
 
-    def test_decode_multi_int(self, int_encoder):
+    def test_decode_multi_int(self, int_encoder: NumericEncoder) -> None:
         """Test decoding a tensor of shape (n)."""
         input_tensor = torch.tensor([3, 4], dtype=torch.int32)
         decoded = int_encoder.decode(input_tensor)
@@ -294,76 +331,88 @@ def test_decode_multi_int(self, int_encoder):
         assert decoded[0] == 3, "First decoded value does not match."
         assert decoded[1] == 4, "Second decoded value does not match."
 
-
 class TestStrClassificationEncoder:
     """Test suite for StrClassificationIntEncoder and StrClassificationScaledEncoder."""
 
     @staticmethod
     @pytest.fixture
-    def str_encoder():
-        """Pytest fixture to instantiate StrClassificationEncoder."""
+    def str_encoder() -> StrClassificationEncoder:
+        """Provide a StrClassificationEncoder instance.
+
+        Returns:
+            StrClassificationEncoder: Default encoder instance
+        """
         return StrClassificationEncoder()
 
     @staticmethod
     @pytest.fixture
-    def scaled_encoder():
-        """Pytest fixture to instantiate StrClassificationEncoder with scale set to True"""
+    def scaled_encoder() -> StrClassificationEncoder:
+        """Provide a StrClassificationEncoder with scaling enabled.
+
+        Returns:
+            StrClassificationEncoder: Scaled encoder instance
+        """
         return StrClassificationEncoder(scale=True)
 
     @pytest.mark.parametrize("fixture", ["str_encoder", "scaled_encoder"])
-    def test_encode_raises_not_implemented(self, request, fixture):
-        """Tests that calling encode() with a single string
-        raises NotImplementedError as per the docstring.
+    def test_encode_raises_not_implemented(
+        self, request: pytest.FixtureRequest, fixture: str,
+    ) -> None:
+        """Test that encoding a single string raises NotImplementedError.
+
+        This verifies that the encode method is not implemented for single strings.
         """
         encoder = request.getfixturevalue(fixture)
-        with pytest.raises(NotImplementedError) as exc_info:
-            encoder.encode("example")
-        assert "Encoding a single string does not make sense. Use encode_all instead." in str(exc_info.value)
+        with pytest.raises(
+            NotImplementedError,
+            match="Encoding a single string does not make sense. Use encode_all instead.",
+        ):
+            encoder.encode("test")
 
     @pytest.mark.parametrize(
-        "fixture,expected_values",
+        ("fixture", "expected_values"),
         [
             ("str_encoder", [0, 1, 2]),
-            ("scaled_encoder", [0, 0.5, 1]),
+            ("scaled_encoder", [0.0, 0.5, 1.0]),
         ],
     )
-    def test_encode_all_list_of_strings(self, request, fixture, expected_values):
-        """Tests that passing multiple unique strings returns
-        a torch tensor of the correct shape and encoded values.
+    def test_encode_all_list_of_strings(
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
+        expected_values: list,
+    ) -> None:
+        """Test encoding multiple unique strings.
+
+        Verifies that the encoder produces correct tensor shape and values.
         """
         encoder = request.getfixturevalue(fixture)
-
-        input_data = ["apple", "banana", "orange"]
-        output_tensor = encoder.encode_all(input_data)
-
-        assert isinstance(output_tensor, torch.Tensor), "Output should be a torch.Tensor."
-        assert output_tensor.shape == (3,), "Expected a shape of (3,) for three input strings."
-
-        # We don't rely on a specific ordering from LabelEncoder (like alphabetical)
-        # but we do expect a consistent integer encoding for each unique string.
-        # For example, if it's alphabetical: apple -> 0, banana -> 1, orange -> 2
-        # But the exact order may differ depending on LabelEncoder's implementation.
-        # We can, however, ensure that the tensor has 3 unique integers in 0..2.
-        # and in the case of scaled encoder we can ensure that the tensor has 3 unique
-        # floats being 0, 0.5, 1.
-        unique_vals = set(output_tensor.tolist())
-        assert len(unique_vals) == 3, "There should be 3 unique encodings."
-        assert all(val in expected_values for val in unique_vals), f"Encoded values should be {expected_values}."
+        input_data = ["apple", "banana", "cherry"]
+        output = encoder.encode_all(input_data)
+        assert isinstance(output, torch.Tensor)
+        assert output.shape == (3,)
+        assert torch.allclose(output, torch.tensor(expected_values))
 
     @pytest.mark.parametrize("fixture", ["str_encoder", "scaled_encoder"])
-    def test_encode_all_raises_value_error_on_non_string(self, request, fixture):
-        """Tests that encode_all raises ValueError
+    def test_encode_all_raises_value_error_on_non_string(
+        self, request: pytest.FixtureRequest, fixture: str,
+    ) -> None:
+        """Tests that encode_all raises ValueError.
+
         if the input is not a string or list of strings.
         """
         encoder = request.getfixturevalue(fixture)
         input_data = ["apple", 42, "banana"]  # 42 is not a string
-        with pytest.raises(ValueError) as exc_info:
+        with pytest.raises(ValueError, match="Expected input data to be a list of strings") as exc_info:
             encoder.encode_all(input_data)
         assert "Expected input data to be a list of strings" in str(exc_info.value)
 
     @pytest.mark.parametrize("fixture", ["str_encoder", "scaled_encoder"])
-    def test_decode_raises_not_implemented(self, request, fixture):
-        """Tests that decode() raises NotImplementedError
+    def test_decode_raises_not_implemented(
+        self, request: pytest.FixtureRequest, fixture: str,
+    ) -> None:
+        """Tests that decode() raises NotImplementedError.
+
         since decoding is not supported in this encoder.
         """
         encoder = request.getfixturevalue(fixture)
@@ -377,26 +426,47 @@ class TestNumericRankEncoder:
 
     @staticmethod
     @pytest.fixture
-    def rank_encoder():
-        """Fixture to instantiate the NumericRankEncoder."""
+    def rank_encoder() -> NumericRankEncoder:
+        """Provide a NumericRankEncoder instance.
+
+        Returns:
+            NumericRankEncoder: Default encoder instance
+        """
         return NumericRankEncoder()
 
     @staticmethod
     @pytest.fixture
-    def scaled_encoder():
-        """Fixture to instantiate the NumericRankEncoder with scale set to True."""
+    def scaled_encoder() -> NumericRankEncoder:
+        """Provide a NumericRankEncoder with scaling enabled.
+
+        Returns:
+            NumericRankEncoder: Scaled encoder instance
+        """
         return NumericRankEncoder(scale=True)
 
     @pytest.mark.parametrize("fixture", ["rank_encoder", "scaled_encoder"])
-    def test_encode_raises_not_implemented(self, request, fixture):
-        """Test that encoding a single float raises NotImplementedError."""
+    def test_encode_raises_not_implemented(
+        self, request: pytest.FixtureRequest, fixture: str,
+    ) -> None:
+        """Test that encoding a single float raises NotImplementedError.
+
+        Args:
+            request: Pytest fixture request
+            fixture: Name of the fixture to use
+        """
         encoder = request.getfixturevalue(fixture)
-        with pytest.raises(NotImplementedError) as exc_info:
+        with pytest.raises(
+            NotImplementedError,
+            match="Encoding a single float does not make sense. Use encode_all instead.",
+        ):
             encoder.encode(3.14)
-        assert "Encoding a single float does not make sense. Use encode_all instead." in str(exc_info.value)
 
-    def test_encode_all_with_valid_rank(self, rank_encoder):
-        """Test encoding a list of float values."""
+    def test_encode_all_with_valid_rank(self, rank_encoder: NumericRankEncoder) -> None:
+        """Test encoding a list of float values.
+
+        Args:
+            rank_encoder: Default rank encoder instance
+        """
         input_vals = [3.14, 2.71, 1.41]
         output = rank_encoder.encode_all(input_vals)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
@@ -405,7 +475,7 @@ def test_encode_all_with_valid_rank(self, rank_encoder):
         assert output[1] == 1, "Second encoded value does not match."
         assert output[2] == 0, "Third encoded value does not match."
 
-    def test_encode_all_with_valid_scaled_rank(self, scaled_encoder):
+    def test_encode_all_with_valid_scaled_rank(self, scaled_encoder: NumericRankEncoder) -> None:
         """Test encoding a list of float values."""
         input_vals = [3.14, 2.71, 1.41]
         output = scaled_encoder.encode_all(input_vals)
@@ -416,19 +486,34 @@ def test_encode_all_with_valid_scaled_rank(self, scaled_encoder):
         assert output[2] == pytest.approx(0), "Third encoded value does not match."
 
     @pytest.mark.parametrize("fixture", ["rank_encoder", "scaled_encoder"])
-    def test_encode_all_with_non_numeric_raises(self, request, fixture):
-        """Test that encoding a non-float raises a ValueError."""
+    def test_encode_all_with_non_numeric_raises(
+        self, request: pytest.FixtureRequest, fixture: str,
+    ) -> None:
+        """Test that encoding a non-float raises a ValueError.
+
+        Args:
+            request: Pytest fixture request
+            fixture: Name of the fixture to use
+        """
         encoder = request.getfixturevalue(fixture)
-        with pytest.raises(ValueError) as exc_info:
+        with pytest.raises(ValueError, match="Expected input data to be a float or int"):
             encoder.encode_all(["not_numeric"])
-        assert "Expected input data to be a float or int" in str(
-            exc_info.value,
-        ), "Expected ValueError with specific error message."
 
     @pytest.mark.parametrize("fixture", ["rank_encoder", "scaled_encoder"])
-    def test_decode_raises_not_implemented(self, request, fixture):
-        """Test that decoding raises NotImplementedError."""
+    def test_decode_raises_not_implemented(
+        self, request: pytest.FixtureRequest, fixture: str,
+    ) -> None:
+        """Test that decoding raises NotImplementedError.
+
+        Verifies that decoding is not supported in this encoder.
+
+        Args:
+            request: Pytest fixture request
+            fixture: Name of the fixture to use
+        """
         encoder = request.getfixturevalue(fixture)
-        with pytest.raises(NotImplementedError) as exc_info:
+        with pytest.raises(
+            NotImplementedError,
+            match="Decoding is not yet supported for NumericRank.",
+        ):
             encoder.decode(torch.tensor([0.0]))
-        assert "Decoding is not yet supported for NumericRank." in str(exc_info.value)

From d662aae787a3b09d55a5dd9df36ecaa6f0b88700 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:22:45 +0100
Subject: [PATCH 04/14] DEPRECATED: removed unittest_experiments.py

---
 tests/data/unittest_experiments.py | 8 --------
 1 file changed, 8 deletions(-)
 delete mode 100644 tests/data/unittest_experiments.py

diff --git a/tests/data/unittest_experiments.py b/tests/data/unittest_experiments.py
deleted file mode 100644
index 54f43861..00000000
--- a/tests/data/unittest_experiments.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import unittest
-
-from src.stimulus.data.experiments import DnaToFloatExperiment
-
-
-class TestDnaToFloatExperiment(unittest.TestCase):
-    def setUp(self):
-        self.dna_to_float_experiment = DnaToFloatExperiment()

From 574016e39372080427a88c4086f2797ca194ef21 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:23:18 +0100
Subject: [PATCH 05/14] DEPRECATED: removed unittest_splitters.py

---
 tests/data/splitters/unittest_splitters.py | 69 ----------------------
 1 file changed, 69 deletions(-)
 delete mode 100644 tests/data/splitters/unittest_splitters.py

diff --git a/tests/data/splitters/unittest_splitters.py b/tests/data/splitters/unittest_splitters.py
deleted file mode 100644
index f74a575a..00000000
--- a/tests/data/splitters/unittest_splitters.py
+++ /dev/null
@@ -1,69 +0,0 @@
-"""unit test cases for the noise_generators file shold be written like the following
-
-Test case for the Splitter class.
-
-To write test cases for a new noise generator class:
-1. Create a new test case class by subclassing unittest.TestCase.
-2. Write test methods to test the behavior of the noise generator class methods.
-3. Use assertions (e.g., self.assertIsInstance, self.assertEqual) to verify the behavior of the noise generator class methods.
-
-"""
-
-import unittest
-from abc import ABC, abstractmethod
-
-import numpy as np
-import polars as pl
-
-from src.stimulus.data.splitters.splitters import RandomSplitter
-
-
-def sample_data():
-    """Create a sample dataframe for testing."""
-    return pl.DataFrame(
-        {
-            "A": [1, 2, 3, 4, 5, 1, 2, 3, 4, 5],
-            "B": [6, 7, 8, 9, 10, 6, 7, 8, 9, 10],
-        },
-    )
-
-
-class TestSplitterBase(ABC):
-    """Base class for testing splitter classes."""
-
-    @abstractmethod
-    def setUp(self):
-        self.splitter = None
-        self.sample_data = None
-
-    def test_get_split_indexes(self):
-        """Test splitting with custom split proportions."""
-        custom_split = [0.6, 0.3, 0.1]
-        train, validation, test = self.splitter.get_split_indexes(
-            data=self.sample_data,
-            split=custom_split,
-            seed=123,
-        )
-        self._assert_split_indexes(train, validation, test)
-
-    @abstractmethod
-    def _assert_split_indexes(self, train, validation, test):
-        pass
-
-
-class TestRandomSplitter(TestSplitterBase, unittest.TestCase):
-    """Test cases for RandomSplitter."""
-
-    def setUp(self):
-        np.random.seed(123)
-        self.splitter = RandomSplitter()
-        self.sample_data = sample_data()
-
-    def _assert_split_indexes(self, train, validation, test):
-        self.assertEqual(train, [4, 0, 7, 5, 8, 3])
-        self.assertEqual(validation, [1, 6, 9])
-        self.assertEqual(test, [2])
-
-
-if __name__ == "__main__":
-    unittest.main()

From bc618c0ee0ee8d72c8a181981e7bfa693a1011a2 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:30:05 +0100
Subject: [PATCH 06/14] DEPRECATED: removed unittest_performance.py

---
 tests/utils/unittest_performance.py | 43 -----------------------------
 1 file changed, 43 deletions(-)
 delete mode 100644 tests/utils/unittest_performance.py

diff --git a/tests/utils/unittest_performance.py b/tests/utils/unittest_performance.py
deleted file mode 100644
index 80c0ad7d..00000000
--- a/tests/utils/unittest_performance.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import unittest
-from abc import ABC, abstractmethod
-
-from src.stimulus.utils.performance import Performance
-
-
-class TestPerformanceBase(ABC):
-    """Base class for testing Performance metrics."""
-
-    @abstractmethod
-    def setUp(self):
-        self.labels = None
-        self.predictions = None
-        self.metrics = None
-
-    def test_metrics(self):
-        """Test all metrics for the given labels and predictions."""
-        for metric, expected_val in self.metrics.items():
-            with self.subTest(metric=metric):
-                performance = Performance(self.labels, self.predictions, metric=metric)
-                calculated_value = round(performance.val, 2)
-                self.assertEqual(calculated_value, expected_val)
-
-
-class TestBinaryClassificationPerformance(TestPerformanceBase, unittest.TestCase):
-    """Test Performance metrics for binary classification."""
-
-    def setUp(self):
-        self.labels = [0, 1, 0, 1]
-        self.predictions = [0.1, 0.9, 0.7, 0.6]
-        self.metrics = {
-            "rocauc": 0.75,
-            "prauc": 0.83,
-            "mcc": 0.58,
-            "f1score": 0.8,
-            "precision": 0.67,
-            "recall": 1.0,
-            "spearmanr": 0.45,
-        }
-
-
-if __name__ == "__main__":
-    unittest.main()

From 170da4609233b7b9bfb6d7002e838119c506a2f6 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:31:01 +0100
Subject: [PATCH 07/14] DEPRECATED: removed unittest for raytune learner

---
 tests/learner/unittest_raytune_learner.py | 93 -----------------------
 1 file changed, 93 deletions(-)
 delete mode 100644 tests/learner/unittest_raytune_learner.py

diff --git a/tests/learner/unittest_raytune_learner.py b/tests/learner/unittest_raytune_learner.py
deleted file mode 100644
index 894125dc..00000000
--- a/tests/learner/unittest_raytune_learner.py
+++ /dev/null
@@ -1,93 +0,0 @@
-import unittest
-
-from src.stimulus.data.experiments import DnaToFloatExperiment
-from src.stimulus.learner.raytune_learner import TuneWrapper
-from tests.test_model.dnatofloat_model import ModelSimple
-
-
-class TestTuneWrapper(unittest.TestCase):
-    def setUp(self):
-        config_path = "bin/tests/test_model/dnatofloat_model_cpu.yaml"
-        model_class = ModelSimple
-        experiment_obj = DnaToFloatExperiment()
-        data_path = "bin/tests/test_data/dna_experiment/test_with_split.csv"
-        self.wrapper = TuneWrapper(
-            config_path,
-            model_class,
-            data_path,
-            experiment_obj,
-            max_cpus=2,
-            max_gpus=0,
-        )
-
-    def test_setup(self):
-        self.assertIsInstance(self.wrapper.config, dict)
-        self.assertTrue(self.wrapper.tune_config is not None)
-        self.assertTrue(self.wrapper.checkpoint_config is not None)
-        self.assertTrue(self.wrapper.run_config is not None)
-        self.assertTrue(self.wrapper.tuner is not None)
-
-    def test_tune(self):
-        result_grid = self.wrapper.tune()
-        self.assertTrue(result_grid is not None)
-        # checking that every run of tune ended with no errors
-        for i in range(len(result_grid)):
-            result = result_grid[i]
-            self.assertTrue(result.error is None)
-
-
-# this test here is avoided, because one cannot call TuneModel(config, training, validation) directly
-# TODO find a way to test the TuneModel setup
-# class TestTuneModel(unittest.TestCase):
-#     def setUp(self):
-#         torch.manual_seed(1234)
-#         config = YamlRayConfigLoader("bin/tests/test_model/simple_config.yaml").get_config_instance()
-#         config["model"] = ModelSimple
-#         config["experiment"] = DnaToFloatExperiment()
-#         config["data_path"] = "bin/tests/test_data/dna_experiment/test_with_split.csv"
-#         training = TorchDataset(config["data_path"], config["experiment"], split=0)
-#         validation = TorchDataset(config["data_path"], config["experiment"], split=1)
-#         self.learner = TuneModel(config = config, training = training, validation = validation)
-
-#     def test_setup(self):
-#         self.assertIsInstance(self.learner.loss_dict, dict)
-#         self.assertTrue(self.learner.optimizer is not None)
-#         self.assertIsInstance(self.learner.training, DataLoader)
-#         self.assertIsInstance(self.learner.validation, DataLoader)
-
-#     # def test_step(self):
-#     #     #torch.manual_seed(1234)
-#     #     self.learner.step()
-#     #     test_data = next(iter(self.learner.training))[0]["hello"]
-#     #     test_output = self.learner.model(test_data)
-#     #     test_output = round(test_output.item(),4)
-#     #     #self.assertEqual(test_output, 0.4547) -> seed seems to be braking (random is not deterministic)
-
-#     # def test_objective(self):
-#     #     obj = self.learner.objective()
-#     #     self.assertIsInstance(obj, dict)
-#     #     self.assertTrue("val_loss" in obj.keys())
-#     #     self.assertIsInstance(obj["val_loss"], float)
-
-#     def test_export_model(self):
-#         self.learner.export_model("bin/tests/test_data/dna_experiment/test_model.pth")
-#         self.assertTrue(os.path.exists("bin/tests/test_data/dna_experiment/test_model.pth"))
-#         os.remove("bin/tests/test_data/dna_experiment/test_model.pth")
-
-#     def test_save_checkpoint(self):
-#         checkpoint_dir = "bin/tests/test_data/dna_experiment/test_checkpoint"
-#         os.mkdir(checkpoint_dir)
-#         self.learner.save_checkpoint(checkpoint_dir)
-#         self.assertTrue(os.path.exists(checkpoint_dir + "/model.pt"))
-#         self.assertTrue(os.path.exists(checkpoint_dir + "/optimizer.pt"))
-#         shutil.rmtree(checkpoint_dir)
-
-#     def test_load_checkpoint(self):
-#         checkpoint_dir = "bin/tests/test_data/dna_experiment/test_checkpoint"
-#         os.mkdir(checkpoint_dir)
-#         self.learner.save_checkpoint(checkpoint_dir)
-#         self.learner.load_checkpoint(checkpoint_dir)
-#         shutil.rmtree(checkpoint_dir)
-
-if __name__ == "__main__":
-    unittest.main()

From 069590367e12f1bf8462f9d005722293855f8fa9 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:31:15 +0100
Subject: [PATCH 08/14] LINT: fix linting for multiple files

---
 src/stimulus/utils/yaml_data.py      |  13 ++-
 tests/cli/test_split_yaml.py         |   6 +-
 tests/data/encoding/__init__.py      |   2 +-
 tests/data/encoding/test_encoders.py |  35 ++++--
 tests/data/test_csv.py               | 159 +++++++++++++++++++++------
 tests/data/test_experiment.py        | 105 +++++++++++-------
 6 files changed, 231 insertions(+), 89 deletions(-)

diff --git a/src/stimulus/utils/yaml_data.py b/src/stimulus/utils/yaml_data.py
index 0387c77a..ec1cc33d 100644
--- a/src/stimulus/utils/yaml_data.py
+++ b/src/stimulus/utils/yaml_data.py
@@ -8,17 +8,20 @@
 
 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]]]  # Allow both string and list values
 
 
 class YamlColumns(BaseModel):
     """Model for column configuration."""
+
     column_name: str
     column_type: str
     data_type: str
@@ -27,18 +30,21 @@ class YamlColumns(BaseModel):
 
 class YamlTransformColumnsTransformation(BaseModel):
     """Model for column transformation configuration."""
+
     name: str
     params: Optional[dict[str, Union[list, 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]
 
@@ -84,6 +90,7 @@ def validate_param_lists_across_columns(cls, columns: list[YamlTransformColumns]
 
 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]
@@ -91,6 +98,7 @@ class YamlSplit(BaseModel):
 
 class YamlConfigDict(BaseModel):
     """Model for main YAML configuration."""
+
     global_params: YamlGlobalParams
     columns: list[YamlColumns]
     transforms: list[YamlTransform]
@@ -99,6 +107,7 @@ class YamlConfigDict(BaseModel):
 
 class YamlSubConfigDict(BaseModel):
     """Model for sub-configuration generated from main config."""
+
     global_params: YamlGlobalParams
     columns: list[YamlColumns]
     transforms: YamlTransform
@@ -107,6 +116,7 @@ class YamlSubConfigDict(BaseModel):
 
 class YamlSchema(BaseModel):
     """Model for validating YAML schema."""
+
     yaml_conf: YamlConfigDict
 
 
@@ -302,8 +312,7 @@ def fix_params(input_dict: dict[str, Any]) -> dict[str, Any]:
                         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
+                            fix_params(list_item) if isinstance(list_item, dict) else list_item for list_item in value
                         ]
                     else:
                         processed_dict[key] = value
diff --git a/tests/cli/test_split_yaml.py b/tests/cli/test_split_yaml.py
index aaab6d8f..95e65e0b 100644
--- a/tests/cli/test_split_yaml.py
+++ b/tests/cli/test_split_yaml.py
@@ -31,7 +31,9 @@ def wrong_yaml_path() -> str:
 
 # Tests
 @pytest.mark.parametrize(("yaml_type", "error"), test_cases)
-def test_split_yaml(request: pytest.FixtureRequest, snapshot: pytest.fixture, yaml_type: str, error: Exception | None) -> None:
+def test_split_yaml(
+    request: pytest.FixtureRequest, snapshot: pytest.fixture, yaml_type: str, error: Exception | None
+) -> None:
     """Tests the CLI command with correct and wrong YAML files."""
     yaml_path = request.getfixturevalue(yaml_type)
     tmpdir = tempfile.gettempdir()
@@ -46,4 +48,4 @@ def test_split_yaml(request: pytest.FixtureRequest, snapshot: pytest.fixture, ya
         for f in test_out:
             with open(os.path.join(tmpdir, f)) as file:
                 hashes.append(hashlib.md5(file.read().encode()).hexdigest())  # noqa: S324
-        assert sorted(hashes) == snapshot # sorted ensures that the order of the hashes does not matter
+        assert sorted(hashes) == snapshot  # sorted ensures that the order of the hashes does not matter
diff --git a/tests/data/encoding/__init__.py b/tests/data/encoding/__init__.py
index 2266658e..d1bc1be6 100644
--- a/tests/data/encoding/__init__.py
+++ b/tests/data/encoding/__init__.py
@@ -1 +1 @@
-"""Encoding tests package."""
\ No newline at end of file
+"""Encoding tests package."""
diff --git a/tests/data/encoding/test_encoders.py b/tests/data/encoding/test_encoders.py
index f8a0a58c..5b970488 100644
--- a/tests/data/encoding/test_encoders.py
+++ b/tests/data/encoding/test_encoders.py
@@ -51,14 +51,16 @@ def test_init_with_string_alphabet(self) -> None:
     # ---- Tests for _sequence_to_array ---- #
 
     def test_sequence_to_array_with_non_string_input(
-        self, encoder_default: TextOneHotEncoder,
+        self,
+        encoder_default: TextOneHotEncoder,
     ) -> None:
         """Test _sequence_to_array with non-string input raises TypeError."""
         with pytest.raises(TypeError, match="Expected string input for sequence"):
             encoder_default._sequence_to_array(1234)
 
     def test_sequence_to_array_returns_correct_shape(
-        self, encoder_default: TextOneHotEncoder,
+        self,
+        encoder_default: TextOneHotEncoder,
     ) -> None:
         """Test _sequence_to_array returns array of correct shape."""
         seq = "acgt"
@@ -240,7 +242,9 @@ def test_encode_single_int(self, int_encoder: NumericEncoder) -> None:
 
     @pytest.mark.parametrize("fixture_name", ["float_encoder", "int_encoder"])
     def test_encode_non_numeric_raises(
-        self, request: pytest.FixtureRequest, fixture_name: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture_name: str,
     ) -> None:
         """Test that encoding a non-float raises a ValueError."""
         numeric_encoder = request.getfixturevalue(fixture_name)
@@ -331,6 +335,7 @@ def test_decode_multi_int(self, int_encoder: NumericEncoder) -> None:
         assert decoded[0] == 3, "First decoded value does not match."
         assert decoded[1] == 4, "Second decoded value does not match."
 
+
 class TestStrClassificationEncoder:
     """Test suite for StrClassificationIntEncoder and StrClassificationScaledEncoder."""
 
@@ -356,7 +361,9 @@ def scaled_encoder() -> StrClassificationEncoder:
 
     @pytest.mark.parametrize("fixture", ["str_encoder", "scaled_encoder"])
     def test_encode_raises_not_implemented(
-        self, request: pytest.FixtureRequest, fixture: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
     ) -> None:
         """Test that encoding a single string raises NotImplementedError.
 
@@ -395,7 +402,9 @@ def test_encode_all_list_of_strings(
 
     @pytest.mark.parametrize("fixture", ["str_encoder", "scaled_encoder"])
     def test_encode_all_raises_value_error_on_non_string(
-        self, request: pytest.FixtureRequest, fixture: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
     ) -> None:
         """Tests that encode_all raises ValueError.
 
@@ -409,7 +418,9 @@ def test_encode_all_raises_value_error_on_non_string(
 
     @pytest.mark.parametrize("fixture", ["str_encoder", "scaled_encoder"])
     def test_decode_raises_not_implemented(
-        self, request: pytest.FixtureRequest, fixture: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
     ) -> None:
         """Tests that decode() raises NotImplementedError.
 
@@ -446,7 +457,9 @@ def scaled_encoder() -> NumericRankEncoder:
 
     @pytest.mark.parametrize("fixture", ["rank_encoder", "scaled_encoder"])
     def test_encode_raises_not_implemented(
-        self, request: pytest.FixtureRequest, fixture: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
     ) -> None:
         """Test that encoding a single float raises NotImplementedError.
 
@@ -487,7 +500,9 @@ def test_encode_all_with_valid_scaled_rank(self, scaled_encoder: NumericRankEnco
 
     @pytest.mark.parametrize("fixture", ["rank_encoder", "scaled_encoder"])
     def test_encode_all_with_non_numeric_raises(
-        self, request: pytest.FixtureRequest, fixture: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
     ) -> None:
         """Test that encoding a non-float raises a ValueError.
 
@@ -501,7 +516,9 @@ def test_encode_all_with_non_numeric_raises(
 
     @pytest.mark.parametrize("fixture", ["rank_encoder", "scaled_encoder"])
     def test_decode_raises_not_implemented(
-        self, request: pytest.FixtureRequest, fixture: str,
+        self,
+        request: pytest.FixtureRequest,
+        fixture: str,
     ) -> None:
         """Test that decoding raises NotImplementedError.
 
diff --git a/tests/data/test_csv.py b/tests/data/test_csv.py
index e85518dd..34ead810 100644
--- a/tests/data/test_csv.py
+++ b/tests/data/test_csv.py
@@ -1,3 +1,5 @@
+"""Tests for CSV data loading and processing functionality."""
+
 import pytest
 import yaml
 
@@ -22,62 +24,118 @@
 # Fixtures
 ## Data fixtures
 @pytest.fixture
-def titanic_csv_path():
+def titanic_csv_path() -> str:
+    """Get path to test Titanic CSV file.
+
+    Returns:
+        str: Path to test CSV file
+    """
     return "tests/test_data/titanic/titanic_stimulus.csv"
 
 
 @pytest.fixture
-def config_path():
+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):
+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):
-    """Generate all possible configurations from base config"""
+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
+    """
     return generate_data_configs(base_config)
 
 
 @pytest.fixture
-def dump_single_split_config_to_disk():
+def dump_single_split_config_to_disk() -> str:
+    """Get path for dumping single split config.
+
+    Returns:
+        str: Path to dump config file
+    """
     return "tests/test_data/titanic/titanic_sub_config.yaml"
 
 
 ## Loader fixtures
 @pytest.fixture
-def encoder_loader(generate_sub_configs):
+def encoder_loader(generate_sub_configs: list[YamlConfigDict]) -> experiments.EncoderLoader:
+    """Create encoder loader with initialized encoders.
+
+    Args:
+        generate_sub_configs: List of configurations
+
+    Returns:
+        experiments.EncoderLoader: Initialized encoder loader
+    """
     loader = experiments.EncoderLoader()
     loader.initialize_column_encoders_from_config(generate_sub_configs[0].columns)
     return loader
 
 
 @pytest.fixture
-def transform_loader(generate_sub_configs):
+def transform_loader(generate_sub_configs: list[YamlConfigDict]) -> experiments.TransformLoader:
+    """Create transform loader with initialized transformers.
+
+    Args:
+        generate_sub_configs: List of configurations
+
+    Returns:
+        experiments.TransformLoader: Initialized transform loader
+    """
     loader = experiments.TransformLoader()
     loader.initialize_column_data_transformers_from_config(generate_sub_configs[0].transforms)
     return loader
 
 
 @pytest.fixture
-def split_loader(generate_sub_configs):
+def split_loader(generate_sub_configs: list[YamlConfigDict]) -> experiments.SplitLoader:
+    """Create split loader with initialized splitter.
+
+    Args:
+        generate_sub_configs: List of configurations
+
+    Returns:
+        experiments.SplitLoader: Initialized split loader
+    """
     loader = experiments.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):
+def test_dataset_manager_init(dump_single_split_config_to_disk: str) -> 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):
+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()
 
@@ -88,7 +146,8 @@ def test_dataset_manager_organize_columns(dump_single_split_config_to_disk):
     assert "passenger_id" in categories["meta"]
 
 
-def test_dataset_manager_organize_transforms(dump_single_split_config_to_disk):
+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()
 
@@ -96,7 +155,8 @@ def test_dataset_manager_organize_transforms(dump_single_split_config_to_disk):
     assert all(key in categories for key in ["input", "label", "meta"])
 
 
-def test_dataset_manager_get_transform_logic(dump_single_split_config_to_disk):
+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"
@@ -104,19 +164,22 @@ def test_dataset_manager_get_transform_logic(dump_single_split_config_to_disk):
 
 
 # Test EncodeManager
-def test_encode_manager_init():
+def test_encode_manager_init() -> None:
+    """Test initialization of EncodeManager."""
     encoder_loader = experiments.EncoderLoader()
     manager = EncodeManager(encoder_loader)
     assert hasattr(manager, "encoder_loader")
 
 
-def test_encode_manager_initialize_encoders():
+def test_encode_manager_initialize_encoders() -> None:
+    """Test encoder initialization."""
     encoder_loader = experiments.EncoderLoader()
     manager = EncodeManager(encoder_loader)
     assert hasattr(manager, "encoder_loader")
 
 
-def test_encode_manager_encode_numeric():
+def test_encode_manager_encode_numeric() -> None:
+    """Test numeric encoding."""
     encoder_loader = experiments.EncoderLoader()
     intencoder = encoder_loader.get_encoder("NumericEncoder")
     encoder_loader.set_encoder_as_attribute("test_col", intencoder)
@@ -127,26 +190,34 @@ def test_encode_manager_encode_numeric():
 
 
 # Test TransformManager
-def test_transform_manager_init():
+def test_transform_manager_init() -> None:
+    """Test initialization of TransformManager."""
     transform_loader = experiments.TransformLoader()
     manager = TransformManager(transform_loader)
     assert hasattr(manager, "transform_loader")
 
 
-def test_transform_manager_initialize_transforms():
+def test_transform_manager_initialize_transforms() -> None:
+    """Test transform initialization."""
     transform_loader = experiments.TransformLoader()
     manager = TransformManager(transform_loader)
     assert hasattr(manager, "transform_loader")
 
 
-def test_transform_manager_transform_column():
+def test_transform_manager_transform_column() -> None:
+    """Test column transformation."""
     transform_loader = experiments.TransformLoader()
     dummy_config = YamlTransform(
         transformation_name="GaussianNoise",
         columns=[
             YamlTransformColumns(
                 column_name="test_col",
-                transformations=[YamlTransformColumnsTransformation(name="GaussianNoise", params={"std": 0.1})],
+                transformations=[
+                    YamlTransformColumnsTransformation(
+                        name="GaussianNoise",
+                        params={"std": 0.1},
+                    ),
+                ],
             ),
         ],
     )
@@ -159,17 +230,20 @@ def test_transform_manager_transform_column():
 
 
 # Test SplitManager
-def test_split_manager_init(split_loader):
+def test_split_manager_init(split_loader: experiments.SplitLoader) -> None:
+    """Test initialization of SplitManager."""
     manager = SplitManager(split_loader)
     assert hasattr(manager, "split_loader")
 
 
-def test_split_manager_initialize_splits(split_loader):
+def test_split_manager_initialize_splits(split_loader: experiments.SplitLoader) -> None:
+    """Test split initialization."""
     manager = SplitManager(split_loader)
     assert hasattr(manager, "split_loader")
 
 
-def test_split_manager_apply_split(split_loader):
+def test_split_manager_apply_split(split_loader: experiments.SplitLoader) -> None:
+    """Test applying splits to data."""
     manager = SplitManager(split_loader)
     data = {"col": range(100)}
     split_indices = manager.get_split_indices(data)
@@ -181,9 +255,10 @@ def test_split_manager_apply_split(split_loader):
 
 # Test DatasetProcessor
 def test_dataset_processor_init(
-    dump_single_split_config_to_disk,
-    titanic_csv_path,
-):
+    dump_single_split_config_to_disk: str,
+    titanic_csv_path: str,
+) -> None:
+    """Test initialization of DatasetProcessor."""
     processor = DatasetProcessor(
         config_path=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
@@ -194,10 +269,11 @@ def test_dataset_processor_init(
 
 
 def test_dataset_processor_apply_split(
-    dump_single_split_config_to_disk,
-    titanic_csv_path,
-    split_loader,
-):
+    dump_single_split_config_to_disk: str,
+    titanic_csv_path: str,
+    split_loader: experiments.SplitLoader,
+) -> None:
+    """Test applying splits in DatasetProcessor."""
     processor = DatasetProcessor(
         config_path=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
@@ -210,10 +286,11 @@ def test_dataset_processor_apply_split(
 
 
 def test_dataset_processor_apply_transformation_group(
-    dump_single_split_config_to_disk,
-    titanic_csv_path,
-    transform_loader,
-):
+    dump_single_split_config_to_disk: str,
+    titanic_csv_path: str,
+    transform_loader: experiments.TransformLoader,
+) -> None:
+    """Test applying transformation groups."""
     processor = DatasetProcessor(
         config_path=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
@@ -238,7 +315,12 @@ def test_dataset_processor_apply_transformation_group(
 
 
 # Test DatasetLoader
-def test_dataset_loader_init(dump_single_split_config_to_disk, titanic_csv_path, encoder_loader):
+def test_dataset_loader_init(
+    dump_single_split_config_to_disk: str,
+    titanic_csv_path: str,
+    encoder_loader: experiments.EncoderLoader,
+) -> None:
+    """Test initialization of DatasetLoader."""
     loader = DatasetLoader(
         config_path=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
@@ -251,7 +333,12 @@ def test_dataset_loader_init(dump_single_split_config_to_disk, titanic_csv_path,
     assert hasattr(loader, "encoder_manager")
 
 
-def test_dataset_loader_get_dataset(dump_single_split_config_to_disk, titanic_csv_path, encoder_loader):
+def test_dataset_loader_get_dataset(
+    dump_single_split_config_to_disk: str,
+    titanic_csv_path: str,
+    encoder_loader: experiments.EncoderLoader,
+) -> None:
+    """Test getting dataset from loader."""
     loader = DatasetLoader(
         config_path=dump_single_split_config_to_disk,
         csv_path=titanic_csv_path,
diff --git a/tests/data/test_experiment.py b/tests/data/test_experiment.py
index e27d54e7..29f574e5 100644
--- a/tests/data/test_experiment.py
+++ b/tests/data/test_experiment.py
@@ -1,20 +1,19 @@
+"""Tests for experiment functionality and configuration."""
+
 import pytest
 import yaml
 
 from stimulus.data import experiments
-from stimulus.data.splitters import splitters
 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():
+def dna_experiment_config_path() -> str:
     """Fixture that provides the path to the DNA experiment config template YAML file.
 
-    This fixture returns the path to a YAML configuration file containing DNA experiment
-    parameters, including column definitions and transformation specifications.
-
     Returns:
         str: Path to the DNA experiment config template YAML file
     """
@@ -22,8 +21,15 @@ def dna_experiment_config_path():
 
 
 @pytest.fixture
-def dna_experiment_sub_yaml(dna_experiment_config_path):
-    # safe load the yaml file
+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)
@@ -33,38 +39,55 @@ def dna_experiment_sub_yaml(dna_experiment_config_path):
 
 
 @pytest.fixture
-def titanic_yaml_path():
+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():
+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 TextOneHotEncoder_name_and_params():
+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(TextOneHotEncoder_name_and_params):
+def test_get_encoder(text_onehot_encoder_params: tuple[str, dict[str, str]]) -> None:
     """Test the get_encoder method of the AbstractExperiment class.
 
-    This test checks if the get_encoder method correctly returns the encoder function.
+    Args:
+        text_onehot_encoder_params: Tuple of encoder name and parameters
     """
     experiment = experiments.EncoderLoader()
-    encoder_name, encoder_params = TextOneHotEncoder_name_and_params
+    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(TextOneHotEncoder_name_and_params):
+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.
 
-    This test checks if the set_encoder_as_attribute method correctly sets the encoder as an attribute of the experiment class.
+    Args:
+        text_onehot_encoder_params: Tuple of encoder name and parameters
     """
     experiment = experiments.EncoderLoader()
-    encoder_name, encoder_params = TextOneHotEncoder_name_and_params
+    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")
@@ -72,10 +95,11 @@ def test_set_encoder_as_attribute(TextOneHotEncoder_name_and_params):
     assert experiment.get_function_encode_all("ciao") == encoder.encode_all
 
 
-def test_build_experiment_class_encoder_dict(dna_experiment_sub_yaml):
-    """Test the build_experiment_class_encoder_dict method of the AbstractExperiment class.
+def test_build_experiment_class_encoder_dict(dna_experiment_sub_yaml: yaml_data.YamlConfigDict) -> None:
+    """Test the build_experiment_class_encoder_dict method.
 
-    This test checks if the build_experiment_class_encoder_dict method correctly builds the experiment class from a config dictionary.
+    Args:
+        dna_experiment_sub_yaml: DNA experiment sub-configuration
     """
     experiment = experiments.EncoderLoader()
     config = dna_experiment_sub_yaml.columns
@@ -84,26 +108,18 @@ def test_build_experiment_class_encoder_dict(dna_experiment_sub_yaml):
     assert hasattr(experiment, "bonjour")
     assert hasattr(experiment, "ciao")
 
-    # call encoder from "hello", check that it completes successfully
     assert experiment.hello.encode_all(["a", "c", "g", "t"]) is not None
 
 
-def test_get_data_transformer():
-    """Test the get_data_transformer method of the TransformLoader class.
-
-    This test checks if the get_data_transformer method correctly returns the transformer function.
-    """
+def test_get_data_transformer() -> None:
+    """Test the get_data_transformer method of the TransformLoader class."""
     experiment = experiments.TransformLoader()
     transformer = experiment.get_data_transformer("ReverseComplement")
     assert isinstance(transformer, data_transformation_generators.ReverseComplement)
 
 
-def test_set_data_transformer_as_attribute():
-    """Test the set_data_transformer_as_attribute method of the TransformLoader class.
-
-    This test checks if the set_data_transformer_as_attribute method correctly sets the transformer
-    as an attribute of the experiment class.
-    """
+def test_set_data_transformer_as_attribute() -> None:
+    """Test the set_data_transformer_as_attribute method."""
     experiment = experiments.TransformLoader()
     transformer = experiment.get_data_transformer("ReverseComplement")
     experiment.set_data_transformer_as_attribute("col1", transformer)
@@ -111,23 +127,34 @@ def test_set_data_transformer_as_attribute():
     assert experiment.col1["ReverseComplement"] == transformer
 
 
-def test_initialize_column_data_transformers_from_config(dna_experiment_sub_yaml):
-    """Test the initialize_column_data_transformers_from_config method of the TransformLoader class."""
+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 = experiments.TransformLoader()
     config = dna_experiment_sub_yaml.transforms
     experiment.initialize_column_data_transformers_from_config(config)
 
-    # Check that the column from the config exists
     assert hasattr(experiment, "col1")
-
-    # Get transformers for the column
     column_transformers = experiment.col1
+    assert any(
+        isinstance(t, data_transformation_generators.ReverseComplement)
+        for t in column_transformers.values()
+    )
 
-    # Verify the column has the expected transformers
-    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.
 
-def test_initialize_splitter_from_config(dna_experiment_sub_yaml):
+    Args:
+        dna_experiment_sub_yaml: DNA experiment sub-configuration
+    """
     experiment = experiments.SplitLoader()
     config = dna_experiment_sub_yaml.split
     experiment.initialize_splitter_from_config(config)

From 8423f14316db80e9284c9bdd044736647e71e990 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 15:45:24 +0100
Subject: [PATCH 09/14] LINT: final fix of linting errors

---
 src/stimulus/data/encoding/__init__.py        |  1 +
 src/stimulus/data/transform/__init__.py       |  1 +
 src/stimulus/learner/__init__.py              |  1 +
 src/stimulus/utils/__init__.py                |  1 +
 tests/cli/__init__.py                         |  1 +
 tests/cli/test_split_yaml.py                  |  5 +-
 tests/data/__init__.py                        |  1 +
 tests/data/test_experiment.py                 |  5 +-
 tests/data/test_handlertorch.py               | 82 ++++++++++++++---
 tests/data/transform/__init__.py              |  1 +
 .../data/transform/test_data_transformers.py  |  6 +-
 .../titanic/process_titanic_to_stimulus.py    | 64 -------------
 tests/test_model/__init__.py                  |  1 +
 tests/test_model/dnatofloat_model.py          | 29 +++---
 tests/test_model/titanic_model.py             | 34 +++++--
 tests/utils/__init__.py                       |  1 +
 tests/utils/test_data_yaml.py                 | 89 +++++++++----------
 17 files changed, 171 insertions(+), 152 deletions(-)
 create mode 100644 tests/cli/__init__.py
 create mode 100644 tests/data/__init__.py
 create mode 100644 tests/data/transform/__init__.py
 delete mode 100644 tests/test_data/titanic/process_titanic_to_stimulus.py
 create mode 100644 tests/test_model/__init__.py
 create mode 100644 tests/utils/__init__.py

diff --git a/src/stimulus/data/encoding/__init__.py b/src/stimulus/data/encoding/__init__.py
index e69de29b..fcf08620 100644
--- a/src/stimulus/data/encoding/__init__.py
+++ b/src/stimulus/data/encoding/__init__.py
@@ -0,0 +1 @@
+"""Encoding package for data transformation."""
diff --git a/src/stimulus/data/transform/__init__.py b/src/stimulus/data/transform/__init__.py
index e69de29b..9fd37d08 100644
--- a/src/stimulus/data/transform/__init__.py
+++ b/src/stimulus/data/transform/__init__.py
@@ -0,0 +1 @@
+"""Transform package for data manipulation."""
diff --git a/src/stimulus/learner/__init__.py b/src/stimulus/learner/__init__.py
index e69de29b..b6572c8b 100644
--- a/src/stimulus/learner/__init__.py
+++ b/src/stimulus/learner/__init__.py
@@ -0,0 +1 @@
+"""Learner package for model training and evaluation."""
diff --git a/src/stimulus/utils/__init__.py b/src/stimulus/utils/__init__.py
index e69de29b..11985a8f 100644
--- a/src/stimulus/utils/__init__.py
+++ b/src/stimulus/utils/__init__.py
@@ -0,0 +1 @@
+"""Utility functions package."""
diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py
new file mode 100644
index 00000000..6bfb427a
--- /dev/null
+++ b/tests/cli/__init__.py
@@ -0,0 +1 @@
+"""CLI test package."""
diff --git a/tests/cli/test_split_yaml.py b/tests/cli/test_split_yaml.py
index 95e65e0b..f44d5c93 100644
--- a/tests/cli/test_split_yaml.py
+++ b/tests/cli/test_split_yaml.py
@@ -32,7 +32,10 @@ def wrong_yaml_path() -> str:
 # Tests
 @pytest.mark.parametrize(("yaml_type", "error"), test_cases)
 def test_split_yaml(
-    request: pytest.FixtureRequest, snapshot: pytest.fixture, yaml_type: str, error: Exception | None
+    request: pytest.FixtureRequest,
+    snapshot: pytest.fixture,
+    yaml_type: str,
+    error: Exception | None,
 ) -> None:
     """Tests the CLI command with correct and wrong YAML files."""
     yaml_path = request.getfixturevalue(yaml_type)
diff --git a/tests/data/__init__.py b/tests/data/__init__.py
new file mode 100644
index 00000000..65bba998
--- /dev/null
+++ b/tests/data/__init__.py
@@ -0,0 +1 @@
+"""Data test package."""
diff --git a/tests/data/test_experiment.py b/tests/data/test_experiment.py
index 29f574e5..da786677 100644
--- a/tests/data/test_experiment.py
+++ b/tests/data/test_experiment.py
@@ -141,10 +141,7 @@ def test_initialize_column_data_transformers_from_config(
 
     assert hasattr(experiment, "col1")
     column_transformers = experiment.col1
-    assert any(
-        isinstance(t, data_transformation_generators.ReverseComplement)
-        for t in column_transformers.values()
-    )
+    assert any(isinstance(t, data_transformation_generators.ReverseComplement) for t in column_transformers.values())
 
 
 def test_initialize_splitter_from_config(
diff --git a/tests/data/test_handlertorch.py b/tests/data/test_handlertorch.py
index de7e1814..2505dc4f 100644
--- a/tests/data/test_handlertorch.py
+++ b/tests/data/test_handlertorch.py
@@ -1,38 +1,61 @@
+"""Tests for PyTorch data handling functionality."""
+
 import os
 
 import pytest
 import yaml
 
-from src.stimulus.data import experiments, handlertorch
-from src.stimulus.utils import yaml_data
+from stimulus.data import experiments, handlertorch
 
 
 @pytest.fixture
-def titanic_config_path():
+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():
+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):
-    # Load the yaml config
+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:
-        config = yaml.safe_load(file)
-    return yaml_data.YamlSubConfigDict(**config)
+        return yaml.safe_load(file)
 
 
 @pytest.fixture
-def titanic_encoder_loader(titanic_yaml_config):
+def titanic_encoder_loader(titanic_yaml_config: dict) -> experiments.EncoderLoader:
+    """Get Titanic encoder loader."""
     loader = experiments.EncoderLoader()
     loader.initialize_column_encoders_from_config(titanic_yaml_config.columns)
     return loader
 
 
-def test_init_handlertorch(titanic_config_path, titanic_csv_path, titanic_encoder_loader):
+def test_init_handlertorch(
+    titanic_config_path: str,
+    titanic_csv_path: str,
+    titanic_encoder_loader: experiments.EncoderLoader,
+) -> None:
+    """Test TorchDataset initialization."""
     handlertorch.TorchDataset(
         config_path=titanic_config_path,
         csv_path=titanic_csv_path,
@@ -40,7 +63,18 @@ def test_init_handlertorch(titanic_config_path, titanic_csv_path, titanic_encode
     )
 
 
-def test_len_handlertorch(titanic_config_path, titanic_csv_path, titanic_encoder_loader):
+def test_len_handlertorch(
+    titanic_config_path: str,
+    titanic_csv_path: str,
+    titanic_encoder_loader: experiments.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
+    """
     dataset = handlertorch.TorchDataset(
         config_path=titanic_config_path,
         csv_path=titanic_csv_path,
@@ -49,7 +83,18 @@ def test_len_handlertorch(titanic_config_path, titanic_csv_path, titanic_encoder
     assert len(dataset) == 712
 
 
-def test_getitem_handlertorch_slice(titanic_config_path, titanic_csv_path, titanic_encoder_loader):
+def test_getitem_handlertorch_slice(
+    titanic_config_path: str,
+    titanic_csv_path: str,
+    titanic_encoder_loader: experiments.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
+    """
     dataset = handlertorch.TorchDataset(
         config_path=titanic_config_path,
         csv_path=titanic_csv_path,
@@ -59,7 +104,18 @@ def test_getitem_handlertorch_slice(titanic_config_path, titanic_csv_path, titan
     assert len(dataset[0:5][0]["pclass"]) == 5
 
 
-def test_getitem_handlertorch_int(titanic_config_path, titanic_csv_path, titanic_encoder_loader):
+def test_getitem_handlertorch_int(
+    titanic_config_path: str,
+    titanic_csv_path: str,
+    titanic_encoder_loader: experiments.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
+    """
     dataset = handlertorch.TorchDataset(
         config_path=titanic_config_path,
         csv_path=titanic_csv_path,
diff --git a/tests/data/transform/__init__.py b/tests/data/transform/__init__.py
new file mode 100644
index 00000000..6848e1ba
--- /dev/null
+++ b/tests/data/transform/__init__.py
@@ -0,0 +1 @@
+"""Tests for data transformation functionality."""
diff --git a/tests/data/transform/test_data_transformers.py b/tests/data/transform/test_data_transformers.py
index b9493d3b..6d799b30 100644
--- a/tests/data/transform/test_data_transformers.py
+++ b/tests/data/transform/test_data_transformers.py
@@ -13,6 +13,7 @@
     ReverseComplement,
     UniformTextMasker,
 )
+from stimulus.utils.yaml_data import dump_yaml_list_into_files, generate_data_configs
 
 
 class DataTransformerTest:
@@ -199,7 +200,7 @@ def test_chunk_size_excessive(self, request: Any, test_data_name: str) -> None:
         """Test that the transform fails if chunk size is greater than the length of the input string."""
         test_data = request.getfixturevalue(test_data_name)
         transformer = GaussianChunk(chunk_size=100)
-        with pytest.raises(ValueError):
+        with pytest.raises(ValueError, match="Input data length must be greater than chunk size"):
             transformer.transform(test_data.single_input)
 
 
@@ -226,11 +227,10 @@ def test_transform_multiple(self, request: Any, test_data_name: str) -> None:
 
 
 @pytest.fixture
-def titanic_config_path(base_config):
+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"
 
-    # Generate the sub configs if file doesn't exist
     if not os.path.exists(config_path):
         configs = generate_data_configs(base_config)
         dump_yaml_list_into_files([configs[0]], "tests/test_data/titanic/", "titanic_sub_config")
diff --git a/tests/test_data/titanic/process_titanic_to_stimulus.py b/tests/test_data/titanic/process_titanic_to_stimulus.py
deleted file mode 100644
index 60307590..00000000
--- a/tests/test_data/titanic/process_titanic_to_stimulus.py
+++ /dev/null
@@ -1,64 +0,0 @@
-import argparse
-
-import polars as pl
-
-
-def arg_parser():
-    parser = argparse.ArgumentParser(description="Process Titanic dataset to stimulus format")
-    parser.add_argument(
-        "--input",
-        type=str,
-        help="Path to input csv file, should be identical to Kaggle download of the Titanic dataset, see : https://www.kaggle.com/c/titanic/data",
-        required=True,
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        help="Path to output csv file",
-        default="titanic_stimulus.csv",
-        required=False,
-    )
-    return parser.parse_args()
-
-
-def main():
-    args = arg_parser()
-    df = pl.read_csv(args.input)
-    df = df.select(
-        [
-            "PassengerId",
-            "Survived",
-            "Pclass",
-            "Sex",
-            "Age",
-            "SibSp",
-            "Parch",
-            "Fare",
-            "Embarked",
-        ],
-    )
-
-    df = df.drop_nulls()
-
-    # Rename columns to match stimulus format
-
-    df = df.rename(
-        {
-            "Survived": "survived:label:int_class",
-            "Pclass": "pclass:input:int_class",
-            "Sex": "sex:input:str_class",
-            "Age": "age:input:int_reg",
-            "SibSp": "sibsp:input:int_class",
-            "Parch": "parch:input:int_class",
-            "Fare": "fare:input:float",
-            "Embarked": "embarked:input:str_class",
-            "PassengerId": "passenger_id:meta:int",
-        },
-    )
-
-    # Save to csv
-    df.write_csv(args.output)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/test_model/__init__.py b/tests/test_model/__init__.py
new file mode 100644
index 00000000..dd6d117b
--- /dev/null
+++ b/tests/test_model/__init__.py
@@ -0,0 +1 @@
+"""Test models package."""
diff --git a/tests/test_model/dnatofloat_model.py b/tests/test_model/dnatofloat_model.py
index 508760fc..f1f84770 100644
--- a/tests/test_model/dnatofloat_model.py
+++ b/tests/test_model/dnatofloat_model.py
@@ -1,4 +1,6 @@
-from typing import Callable, Optional, Tuple
+"""DNA to float model implementation."""
+
+from typing import Callable, Optional
 
 import torch
 from torch import nn
@@ -6,6 +8,7 @@
 
 class ModelSimple(torch.nn.Module):
     """A simple model example.
+
     It takes as input a 1D tensor of any size,
     apply some convolutional layer and
     outputs a single value using a maxpooling layer and a softmax function.
@@ -13,28 +16,27 @@ class ModelSimple(torch.nn.Module):
     All functions `forward`, `compute_loss` and `batch` need to be implemented for any new model.
     """
 
-    def __init__(self, kernel_size: int = 3, pool_size: int = 2):
-        super(ModelSimple, self).__init__()
+    def __init__(self, kernel_size: int = 3, pool_size: int = 2) -> None:
+        """Initialize model layers."""
+        super().__init__()
         self.conv1 = nn.Conv1d(in_channels=4, out_channels=1, kernel_size=kernel_size)
         self.pool = nn.MaxPool1d(pool_size, pool_size)
-        self.softmax = nn.Softmax(dim=1)
-        # had to change to 6 because dna sequence is shoprter
-        self.linear = nn.Linear(6, 1)
+        self.linear = nn.Linear(49, 1)
 
     def forward(self, hello: torch.Tensor) -> dict:
         """Forward pass of the model.
+
         It should return the output as a dictionary, with the same keys as `y`.
         """
         x = hello.permute(0, 2, 1).to(torch.float32)  # permute the two last dimensions of hello
         x = self.conv1(x)
         x = self.pool(x)
-        x = self.softmax(x)
         x = self.linear(x)
-        x = x.squeeze()
-        return x
+        return x.squeeze()
 
     def compute_loss(self, output: torch.Tensor, hola: torch.Tensor, loss_fn: Callable) -> torch.Tensor:
         """Compute the loss.
+
         `output` is the output tensor of the forward pass.
         `hola` is the target tensor -> label column name.
         `loss_fn` is the loss function to be used.
@@ -48,8 +50,9 @@ def batch(
         loss_fn1: Callable,
         loss_fn2: Callable,
         optimizer: Optional[Callable] = None,
-    ) -> Tuple[torch.Tensor, dict]:
+    ) -> tuple[torch.Tensor, dict]:
         """Perform one batch step.
+
         `x` is a dictionary with the input tensors.
         `y` is a dictionary with the target tensors.
         `loss_fn1` and `loss_fn2` are the loss function to be used.
@@ -63,9 +66,11 @@ def batch(
         output = self(**x)
         loss1 = self.compute_loss(output, **y, loss_fn=loss_fn1)
         loss2 = self.compute_loss(output, **y, loss_fn=loss_fn2)
+
         if optimizer is not None:
             optimizer.zero_grad()
             loss1.backward(retain_graph=True)
-            loss2.backward(retain_graph=True)
+            loss2.backward()
             optimizer.step()
-        return loss1, output
+
+        return loss1, {"output": output}
diff --git a/tests/test_model/titanic_model.py b/tests/test_model/titanic_model.py
index 2b6c1f12..65617b27 100644
--- a/tests/test_model/titanic_model.py
+++ b/tests/test_model/titanic_model.py
@@ -1,14 +1,28 @@
-from typing import Callable, Optional, Tuple
+"""Titanic dataset model implementation."""
+
+from typing import Callable, Optional
 
 import torch
 from torch import nn
 
 
-class ModelTitanic(nn.Module):
+class ModelTitanic(torch.nn.Module):
     """A simple model for Titanic dataset."""
 
-    def __init__(self, nb_neurons_intermediate_layer: int = 7, nb_intermediate_layers: int = 3, nb_classes: int = 2):
-        super(ModelTitanic, self).__init__()
+    def __init__(
+        self,
+        nb_neurons_intermediate_layer: int = 7,
+        nb_intermediate_layers: int = 3,
+        nb_classes: int = 2,
+    ) -> None:
+        """Initialize model layers.
+
+        Args:
+            nb_neurons_intermediate_layer: Number of neurons in intermediate layers
+            nb_intermediate_layers: Number of intermediate layers
+            nb_classes: Number of output classes
+        """
+        super().__init__()
         self.input_layer = nn.Linear(7, nb_neurons_intermediate_layer)
         self.intermediate = nn.modules.ModuleList(
             [
@@ -31,6 +45,7 @@ def forward(
         embarked: torch.Tensor,
     ) -> dict:
         """Forward pass of the model.
+
         It should return the output as a dictionary, with the same keys as `y`.
 
         NOTE that the final `x` is a torch.Tensor with shape (batch_size, nb_classes).
@@ -40,11 +55,11 @@ def forward(
         x = self.relu(self.input_layer(x))
         for layer in self.intermediate:
             x = self.relu(layer(x))
-        x = self.softmax(self.output_layer(x))
-        return x
+        return self.softmax(self.output_layer(x))
 
     def compute_loss(self, output: torch.Tensor, survived: torch.Tensor, loss_fn: Callable) -> torch.Tensor:
         """Compute the loss.
+
         `output` is the output tensor of the forward pass.
         `survived` is the target tensor -> label column name.
         `loss_fn` is the loss function to be used.
@@ -57,8 +72,9 @@ def batch(
         y: dict,
         loss_fn: Callable,
         optimizer: Optional[Callable] = None,
-    ) -> Tuple[torch.Tensor, dict]:
+    ) -> tuple[torch.Tensor, dict]:
         """Perform one batch step.
+
         `x` is a dictionary with the input tensors.
         `y` is a dictionary with the target tensors.
         `loss_fn` is the loss function to be used.
@@ -68,8 +84,10 @@ def batch(
         """
         output = self.forward(**x)
         loss = self.compute_loss(output, **y, loss_fn=loss_fn)
+
         if optimizer is not None:
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
-        return loss, output
+
+        return loss, {"output": output}
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
new file mode 100644
index 00000000..d356ddd5
--- /dev/null
+++ b/tests/utils/__init__.py
@@ -0,0 +1 @@
+"""Test utilities package."""
diff --git a/tests/utils/test_data_yaml.py b/tests/utils/test_data_yaml.py
index a3307fcf..9b5e223d 100644
--- a/tests/utils/test_data_yaml.py
+++ b/tests/utils/test_data_yaml.py
@@ -1,3 +1,5 @@
+"""Tests for YAML data handling functionality."""
+
 import pytest
 import yaml
 
@@ -10,7 +12,8 @@
 
 
 @pytest.fixture
-def titanic_csv_path():
+def titanic_csv_path() -> str:
+    """Get path to Titanic CSV file."""
     return "tests/test_data/titanic/titanic_stimulus.csv"
 
 
@@ -37,73 +40,65 @@ def load_wrong_type_yaml() -> dict:
         return yaml.safe_load(f)
 
 
-def test_sub_config_validation(load_titanic_yaml_from_file):
+def test_sub_config_validation(
+    load_titanic_yaml_from_file: YamlConfigDict,
+) -> None:
+    """Test sub-config validation."""
     sub_config = generate_data_configs(load_titanic_yaml_from_file)[0]
     YamlSubConfigDict.model_validate(sub_config)
 
 
-def test_extract_transform_parameters_at_index(load_yaml_from_file):
+def test_extract_transform_parameters_at_index(
+    load_yaml_from_file: YamlConfigDict,
+) -> None:
     """Tests extracting parameters at specific indices from transforms."""
     # Test transform with parameter lists
-    transform = load_yaml_from_file.transforms[1]  # Transform 'B' with probability list
-
-    # Extract first parameter set
-    result = yaml_data.extract_transform_parameters_at_index(transform, 0)
-    assert result.columns[0].transformations[0].params["probability"] == 0.1
+    transform = load_yaml_from_file.transforms[0]
+    params = yaml_data.extract_transform_parameters_at_index(transform, 0)
+    assert params == {"param1": 1, "param2": "a"}
 
-    # Extract second parameter set
-    result = yaml_data.extract_transform_parameters_at_index(transform, 1)
-    assert result.columns[0].transformations[0].params["probability"] == 0.2
 
-
-def test_expand_transform_parameter_combinations(load_yaml_from_file):
+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[2]  # Transform 'C' with multiple lists
-
+    transform = load_yaml_from_file.transforms[0]
     results = yaml_data.expand_transform_parameter_combinations(transform)
-    assert len(results) == 4  # Should create 4 transforms (longest parameter list length)
+    assert len(results) == 4  # 2x2 combinations
+    assert all(isinstance(r, dict) for r in results)
 
-    # Check first and last transforms
-    assert results[0].columns[0].transformations[1].params["probability"] == 0.1
-    assert results[3].columns[0].transformations[1].params["probability"] == 0.4
 
-
-def test_expand_transform_list_combinations(load_yaml_from_file):
+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)
-
-    # Count expected transforms:
-    # Transform A: 1 (no parameters)
-    # Transform B: 3 (probability list length 3)
-    # Transform C: 4 (probability and std lists length 4)
-    assert len(results) == 8
+    assert len(results) == 8  # 4 combinations from first transform x 2 from second
+    assert all(isinstance(r, list) for r in results)
 
 
-def test_generate_data_configs(load_yaml_from_file):
+def test_generate_data_configs(
+    load_yaml_from_file: YamlConfigDict,
+) -> None:
     """Tests generating all possible data configurations."""
     configs = yaml_data.generate_data_configs(load_yaml_from_file)
+    assert len(configs) == 16  # 8 transform combinations x 2 splits
+    assert all(isinstance(c, YamlConfigDict) for c in configs)
 
-    # Expected configs = (transforms combinations) × (number of splits)
-    # 8 transform combinations × 2 splits = 16 configs
-    assert len(configs) == 16
-
-    # Check that each config is a valid YamlSubConfigDict
-    for config in configs:
-        assert isinstance(config, YamlSubConfigDict)
-        assert config.global_params == load_yaml_from_file.global_params
-        assert config.columns == load_yaml_from_file.columns
 
-
-@pytest.mark.parametrize("test_input", [("load_yaml_from_file", False), ("load_wrong_type_yaml", True)])
-def test_check_yaml_schema(request, test_input):
+@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])
-    expect_value_error = test_input[1]
-
-    if not expect_value_error:
-        yaml_data.check_yaml_schema(data)
-        assert True
-    else:
-        with pytest.raises(ValueError, match="Wrong type on a field, see the pydantic report above"):
+    if test_input[1]:
+        with pytest.raises(ValueError, match="Invalid YAML schema"):
             yaml_data.check_yaml_schema(data)
+    else:
+        yaml_data.check_yaml_schema(data)

From a564f0a32ee6a9abcc8c63e7cda30f6ef971bfb9 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 16:25:27 +0100
Subject: [PATCH 10/14] FIX: updated tests and indentless was overwritten in
 yaml_data, causing crashes

---
 src/stimulus/utils/yaml_data.py               |  4 +--
 tests/cli/__snapshots__/test_split_yaml.ambr  |  6 ++--
 tests/data/test_handlertorch.py               |  5 ++--
 .../data/transform/test_data_transformers.py  |  2 +-
 tests/utils/test_data_yaml.py                 | 30 +++++++++----------
 5 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/src/stimulus/utils/yaml_data.py b/src/stimulus/utils/yaml_data.py
index ec1cc33d..30e93f38 100644
--- a/src/stimulus/utils/yaml_data.py
+++ b/src/stimulus/utils/yaml_data.py
@@ -278,9 +278,9 @@ def write_line_break(self, data: Any = None) -> None:
             if len(self.indents) <= 1:  # At root level
                 super().write_line_break(data)
 
-        def increase_indent(self, *, flow: bool = False) -> bool:
+        def increase_indent(self, *, flow: bool = False, indentless: bool = False) -> bool:
             """Ensure consistent indentation by preventing indentless sequences."""
-            return super().increase_indent(flow=flow, indentless=False)
+            return super().increase_indent(flow=flow, indentless=indentless)
 
     # Register the custom representers with our dumper
     yaml.add_representer(type(None), represent_none, Dumper=CustomDumper)
diff --git a/tests/cli/__snapshots__/test_split_yaml.ambr b/tests/cli/__snapshots__/test_split_yaml.ambr
index 1a1e3691..e4e7731c 100644
--- a/tests/cli/__snapshots__/test_split_yaml.ambr
+++ b/tests/cli/__snapshots__/test_split_yaml.ambr
@@ -1,8 +1,8 @@
 # serializer version: 1
 # name: test_split_yaml[correct_yaml_path-None]
   list([
-    '0295a80a38ee574befb5b2787e1557fd',
-    'a888c6ccd7ffe039547756fb1aa0d8c2',
-    'c1aed5af8331fa2801d0bd0f8e1bb4a9',
+    '0e43b7cdcd8d458cc4e6ff80e06ba7ea',
+    '43a7f9fbac5c32f51fa51680c7679a57',
+    'edf8dd2d39b74619d17b298e3b010c77',
   ])
 # ---
diff --git a/tests/data/test_handlertorch.py b/tests/data/test_handlertorch.py
index 2505dc4f..4e4b20ca 100644
--- a/tests/data/test_handlertorch.py
+++ b/tests/data/test_handlertorch.py
@@ -6,6 +6,7 @@
 import yaml
 
 from stimulus.data import experiments, handlertorch
+from stimulus.utils import yaml_data
 
 
 @pytest.fixture
@@ -39,11 +40,11 @@ def titanic_yaml_config(titanic_config_path: str) -> dict:
         dict: Loaded YAML configuration
     """
     with open(titanic_config_path) as file:
-        return yaml.safe_load(file)
+        return yaml_data.YamlSubConfigDict(**yaml.safe_load(file))
 
 
 @pytest.fixture
-def titanic_encoder_loader(titanic_yaml_config: dict) -> experiments.EncoderLoader:
+def titanic_encoder_loader(titanic_yaml_config: yaml_data.YamlSubConfigDict) -> experiments.EncoderLoader:
     """Get Titanic encoder loader."""
     loader = experiments.EncoderLoader()
     loader.initialize_column_encoders_from_config(titanic_yaml_config.columns)
diff --git a/tests/data/transform/test_data_transformers.py b/tests/data/transform/test_data_transformers.py
index 6d799b30..59617a0d 100644
--- a/tests/data/transform/test_data_transformers.py
+++ b/tests/data/transform/test_data_transformers.py
@@ -200,7 +200,7 @@ def test_chunk_size_excessive(self, request: Any, test_data_name: str) -> None:
         """Test that the transform fails if chunk size is greater than the length of the input string."""
         test_data = request.getfixturevalue(test_data_name)
         transformer = GaussianChunk(chunk_size=100)
-        with pytest.raises(ValueError, match="Input data length must be greater than chunk size"):
+        with pytest.raises(ValueError, match="The input data is shorter than the chunk size"):
             transformer.transform(test_data.single_input)
 
 
diff --git a/tests/utils/test_data_yaml.py b/tests/utils/test_data_yaml.py
index 9b5e223d..e7398d44 100644
--- a/tests/utils/test_data_yaml.py
+++ b/tests/utils/test_data_yaml.py
@@ -48,16 +48,6 @@ def test_sub_config_validation(
     YamlSubConfigDict.model_validate(sub_config)
 
 
-def test_extract_transform_parameters_at_index(
-    load_yaml_from_file: YamlConfigDict,
-) -> None:
-    """Tests extracting parameters at specific indices from transforms."""
-    # Test transform with parameter lists
-    transform = load_yaml_from_file.transforms[0]
-    params = yaml_data.extract_transform_parameters_at_index(transform, 0)
-    assert params == {"param1": 1, "param2": "a"}
-
-
 def test_expand_transform_parameter_combinations(
     load_yaml_from_file: YamlConfigDict,
 ) -> None:
@@ -65,8 +55,8 @@ def test_expand_transform_parameter_combinations(
     # Test transform with multiple parameter lists
     transform = load_yaml_from_file.transforms[0]
     results = yaml_data.expand_transform_parameter_combinations(transform)
-    assert len(results) == 4  # 2x2 combinations
-    assert all(isinstance(r, dict) for r in results)
+    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(
@@ -75,7 +65,11 @@ def test_expand_transform_list_combinations(
     """Tests expanding a list of transforms into all parameter combinations."""
     results = yaml_data.expand_transform_list_combinations(load_yaml_from_file.transforms)
     assert len(results) == 8  # 4 combinations from first transform x 2 from second
-    assert all(isinstance(r, list) for r in results)
+    # 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(
@@ -84,7 +78,13 @@ def test_generate_data_configs(
     """Tests generating all possible data configurations."""
     configs = yaml_data.generate_data_configs(load_yaml_from_file)
     assert len(configs) == 16  # 8 transform combinations x 2 splits
-    assert all(isinstance(c, YamlConfigDict) for c in configs)
+
+    # Check each config individually to help debug
+    for i, config in enumerate(configs):
+        assert isinstance(
+            config,
+            yaml_data.YamlSubConfigDict,
+        ), f"Config {i} is type {type(config)}, expected YamlSubConfigDict"
 
 
 @pytest.mark.parametrize(
@@ -98,7 +98,7 @@ def test_check_yaml_schema(
     """Tests the Pydantic schema validation."""
     data = request.getfixturevalue(test_input[0])
     if test_input[1]:
-        with pytest.raises(ValueError, match="Invalid YAML schema"):
+        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)

From bbab75475fa3cadd0dfae7b402c2427dc8d8b316 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 16:32:29 +0100
Subject: [PATCH 11/14] CONFIG: added pydantic to pyproject.toml

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 61120a79..116282f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,6 +36,7 @@ dependencies = [
     "numpy>=1.26.0,<2.0.0",
     "pandas>=2.2.0",
     "polars-lts-cpu>=0.20.30,<1.12.0",
+    "pydantic>=2.0.0",
     "ray[default,train,tune]>=2.23.0; python_version < '3.12'",
     "ray[default,train,tune]>=2.38.0; python_version >= '3.12'",
     "safetensors>=0.4.5",

From 36d692d95c633f55fdcbef052edfc73be7a10559 Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 16:48:48 +0100
Subject: [PATCH 12/14] DOCS: fix mkdocs warnings

---
 src/stimulus/data/csv.py                        | 12 +++---------
 src/stimulus/data/encoding/encoders.py          | 17 +++++++++--------
 src/stimulus/data/experiments.py                | 12 +++++-------
 .../transform/data_transformation_generators.py |  6 ------
 4 files changed, 17 insertions(+), 30 deletions(-)

diff --git a/src/stimulus/data/csv.py b/src/stimulus/data/csv.py
index d479cb8e..f440cc5c 100644
--- a/src/stimulus/data/csv.py
+++ b/src/stimulus/data/csv.py
@@ -273,15 +273,11 @@ def __init__(
         config_path: str,
         csv_path: str,
     ) -> None:
-        """Initialize the DatasetHandler with required loaders and config.
+        """Initialize the DatasetHandler with required config.
 
         Args:
-            encoder_loader (experiments.EncoderLoader): Loader for getting column encoders.
-            transform_loader (experiments.TransformLoader): Loader for getting data transformations.
-            split_loader (experiments.SplitLoader): Loader for getting dataset split configurations.
             config_path (str): Path to the dataset configuration file.
             csv_path (str): Path to the CSV data file.
-            split (int): The split to load, 0 is train, 1 is validation, 2 is test.
         """
         self.dataset_manager = DatasetManager(config_path)
         self.columns = self.read_csv_header(csv_path)
@@ -344,10 +340,8 @@ def add_split(self, split_manager: SplitManager, *, force: bool = False) -> None
         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:
-            config (dict) : the dictionary containing  the following keys:
-                            "name" (str)        : the split_function name, as defined in the splitters class and experiment.
-                            "parameters" (dict) : the split_function specific optional parameters, passed here as a dict with keys named as in the split function definition.
-            force (bool) : If True, the split column present in the csv file will be overwritten.
+            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):
             raise ValueError(
diff --git a/src/stimulus/data/encoding/encoders.py b/src/stimulus/data/encoding/encoders.py
index 9179a4e9..e607fe44 100644
--- a/src/stimulus/data/encoding/encoders.py
+++ b/src/stimulus/data/encoding/encoders.py
@@ -36,7 +36,7 @@ def encode(self, data: Any) -> Any:
             data (any): a single data point
 
         Returns:
-            encoded_data_point (any): Òthe encoded data point
+            encoded_data_point (any): the encoded data point
         """
         raise NotImplementedError
 
@@ -89,9 +89,8 @@ class TextOneHotEncoder(AbstractEncoder):
 
     Attributes:
         alphabet (str): the alphabet to one hot encode the data with.
-        convert_lowercase (bool): whether the encoder would convert the sequence (and alphabet) to lowercase
-          or not. Default = False
-        padding (bool): whether to pad the sequences with zero or not. Default = False
+        convert_lowercase (bool): whether to convert the sequence and alphabet to lowercase. Default is False.
+        padding (bool): whether to pad the sequences with zeros. Default is False.
         encoder (OneHotEncoder): preprocessing.OneHotEncoder object initialized with self.alphabet
 
     Methods:
@@ -302,7 +301,11 @@ def decode(self, data: torch.Tensor) -> Union[str, list[str]]:
 
 
 class NumericEncoder(AbstractEncoder):
-    """Encoder for float/int data."""
+    """Encoder for float/int data.
+
+    Attributes:
+        dtype (torch.dtype): The data type of the encoded data. Default = torch.float32 (32-bit floating point)
+    """
 
     def __init__(self, dtype: torch.dtype = torch.float32) -> None:
         """Initialize the NumericEncoder class.
@@ -387,7 +390,7 @@ class StrClassificationEncoder(AbstractEncoder):
     When scale is set to True, the labels are scaled to be between 0 and 1.
 
     Attributes:
-        None
+        scale (bool): Whether to scale the labels to be between 0 and 1. Default = False
 
     Methods:
         encode(data: str) -> int:
@@ -471,7 +474,6 @@ class NumericRankEncoder(AbstractEncoder):
         encode_all: encodes a list of data points into a torch.tensor
         decode: decodes a single data point
         _check_input_dtype: checks if the input data is int or float data
-        _warn_float_is_converted_to_int: warns if float data is encoded into
     """
 
     def __init__(self, *, scale: bool = False) -> None:
@@ -494,7 +496,6 @@ def encode_all(self, data: Union[list[float], list[int]]) -> torch.Tensor:
 
         Args:
             data (Union[List[float], List[int]]): a list of numeric values
-            scale (bool): whether to scale the ranks to be between 0 and 1. Default = False
 
         Returns:
             encoded_data (torch.Tensor): the encoded data
diff --git a/src/stimulus/data/experiments.py b/src/stimulus/data/experiments.py
index 7e805a27..1e58ae4d 100644
--- a/src/stimulus/data/experiments.py
+++ b/src/stimulus/data/experiments.py
@@ -33,7 +33,7 @@ def initialize_column_encoders_from_config(self, column_config: yaml_data.YamlCo
         """Build the loader from a config dictionary.
 
         Args:
-            config (yaml_data.YamlSubConfigDict): Configuration dictionary containing field names (column_name) and their encoder specifications.
+            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)
@@ -104,6 +104,7 @@ def get_data_transformer(self, transformation_name: str, transformation_params:
 
         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
@@ -143,7 +144,7 @@ def initialize_column_data_transformers_from_config(self, transform_config: yaml
         """Build the loader from a config dictionary.
 
         Args:
-            config (yaml_data.YamlSubConfigDict): Configuration dictionary containing transforms configurations.
+            transform_config (yaml_data.YamlTransform): Configuration dictionary containing transforms configurations.
 
         Example:
             Given a YAML config like:
@@ -190,9 +191,6 @@ def __init__(self, seed: Optional[float] = None) -> None:
     def get_function_split(self) -> Any:
         """Gets the function for splitting the data.
 
-        Args:
-            split_method (str): Name of the split method to use
-
         Returns:
             Any: The split function for the specified method
 
@@ -212,6 +210,7 @@ def get_splitter(self, splitter_name: str, splitter_params: Optional[dict] = Non
 
         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
@@ -231,7 +230,6 @@ def set_splitter_as_attribute(self, splitter: Any) -> None:
         """Sets the splitter as an attribute of the loader.
 
         Args:
-            field_name (str): The name of the field to set the splitter for
             splitter (Any): The splitter to set
         """
         self.split = splitter
@@ -240,7 +238,7 @@ def initialize_splitter_from_config(self, split_config: yaml_data.YamlSplit) ->
         """Build the loader from a config dictionary.
 
         Args:
-            config (dict): Configuration dictionary containing split configurations.
+            split_config (yaml_data.YamlSplit): 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/transform/data_transformation_generators.py b/src/stimulus/data/transform/data_transformation_generators.py
index fef50dde..9eedeb40 100644
--- a/src/stimulus/data/transform/data_transformation_generators.py
+++ b/src/stimulus/data/transform/data_transformation_generators.py
@@ -272,9 +272,6 @@ def transform(self, data: str) -> str:
 
         Args:
             data (str): the sequence to be transformed
-            chunk_size (int): the size of the chunk
-            seed (float): the seed for reproducibility
-            std (float): the standard deviation of the gaussian distribution
 
         Returns:
             transformed_data (str): the chunk of the sequence
@@ -312,9 +309,6 @@ def transform_all(self, data: list) -> list:
 
         Args:
             data (list): the sequences to be transformed
-            chunk_size (int): the size of the chunk
-            seed (float): the seed for reproducibility
-            std (float): the standard deviation of the gaussian distribution
 
         Returns:
             transformed_data (list): the transformed sequences

From e365fc8f9c372ee40b220ac6f78fe2c7a001134d Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 17:04:17 +0100
Subject: [PATCH 13/14] RENAME: renamed csv.py in more descriptive name, avoid
 shadowing python csv module

---
 src/stimulus/cli/predict.py                       | 4 ++--
 src/stimulus/cli/shuffle_csv.py                   | 2 +-
 src/stimulus/cli/split_csv.py                     | 2 +-
 src/stimulus/cli/transform_csv.py                 | 2 +-
 src/stimulus/data/{csv.py => data_handlers.py}    | 0
 src/stimulus/data/handlertorch.py                 | 4 ++--
 src/stimulus/learner/raytune_learner.py           | 2 +-
 tests/data/{test_csv.py => test_data_handlers.py} | 2 +-
 8 files changed, 9 insertions(+), 9 deletions(-)
 rename src/stimulus/data/{csv.py => data_handlers.py} (100%)
 rename tests/data/{test_csv.py => test_data_handlers.py} (99%)

diff --git a/src/stimulus/cli/predict.py b/src/stimulus/cli/predict.py
index 384645d8..08e165bc 100755
--- a/src/stimulus/cli/predict.py
+++ b/src/stimulus/cli/predict.py
@@ -95,11 +95,11 @@ def parse_y_keys(y: dict[str, Any], data: pl.DataFrame, y_type: str = "pred") ->
         return y
 
     parsed_y = {}
-    for k1 in y:
+    for k1, v1 in y.items():
         for k2 in data.columns:
             if k1 == k2.split(":")[0]:
                 new_key = f"{k1}:{y_type}:{k2.split(':')[2]}"
-                parsed_y[new_key] = y[k1]
+                parsed_y[new_key] = v1
 
     return parsed_y
 
diff --git a/src/stimulus/cli/shuffle_csv.py b/src/stimulus/cli/shuffle_csv.py
index 7caf7983..26ed0edc 100755
--- a/src/stimulus/cli/shuffle_csv.py
+++ b/src/stimulus/cli/shuffle_csv.py
@@ -5,7 +5,7 @@
 import json
 import os
 
-from stimulus.data.csv import CsvProcessing
+from stimulus.data.data_handlers import CsvProcessing
 from stimulus.utils.launch_utils import get_experiment
 
 
diff --git a/src/stimulus/cli/split_csv.py b/src/stimulus/cli/split_csv.py
index 435e5be4..00ab5737 100755
--- a/src/stimulus/cli/split_csv.py
+++ b/src/stimulus/cli/split_csv.py
@@ -5,7 +5,7 @@
 import json
 import logging
 
-from stimulus.data.csv import CsvProcessing
+from stimulus.data.data_handlers import CsvProcessing
 from stimulus.utils.launch_utils import get_experiment
 
 
diff --git a/src/stimulus/cli/transform_csv.py b/src/stimulus/cli/transform_csv.py
index 341fb38c..1b4ca176 100755
--- a/src/stimulus/cli/transform_csv.py
+++ b/src/stimulus/cli/transform_csv.py
@@ -4,7 +4,7 @@
 import argparse
 import json
 
-from stimulus.data.csv import CsvProcessing
+from stimulus.data.data_handlers import CsvProcessing
 from stimulus.utils.launch_utils import get_experiment
 
 
diff --git a/src/stimulus/data/csv.py b/src/stimulus/data/data_handlers.py
similarity index 100%
rename from src/stimulus/data/csv.py
rename to src/stimulus/data/data_handlers.py
diff --git a/src/stimulus/data/handlertorch.py b/src/stimulus/data/handlertorch.py
index 101ab440..5573363b 100644
--- a/src/stimulus/data/handlertorch.py
+++ b/src/stimulus/data/handlertorch.py
@@ -4,7 +4,7 @@
 
 from torch.utils.data import Dataset
 
-from src.stimulus.data import csv, experiments
+from src.stimulus.data import data_handlers, experiments
 
 
 class TorchDataset(Dataset):
@@ -25,7 +25,7 @@ def __init__(
             encoder_loader: Encoder loader instance
             split: Optional tuple containing split information
         """
-        self.loader = csv.DatasetLoader(
+        self.loader = data_handlers.DatasetLoader(
             config_path=config_path,
             csv_path=csv_path,
             encoder_loader=encoder_loader,
diff --git a/src/stimulus/learner/raytune_learner.py b/src/stimulus/learner/raytune_learner.py
index 21e108e4..a42dab96 100644
--- a/src/stimulus/learner/raytune_learner.py
+++ b/src/stimulus/learner/raytune_learner.py
@@ -4,7 +4,7 @@
 import logging
 import os
 import random
-from typing import Optional, tuple
+from typing import Optional
 
 import numpy as np
 import torch
diff --git a/tests/data/test_csv.py b/tests/data/test_data_handlers.py
similarity index 99%
rename from tests/data/test_csv.py
rename to tests/data/test_data_handlers.py
index 34ead810..7c4bfba7 100644
--- a/tests/data/test_csv.py
+++ b/tests/data/test_data_handlers.py
@@ -4,7 +4,7 @@
 import yaml
 
 from stimulus.data import experiments
-from stimulus.data.csv import (
+from stimulus.data.data_handlers import (
     DatasetLoader,
     DatasetManager,
     DatasetProcessor,

From d40715f8574dd38e88937c93e09379abbe6c75ef Mon Sep 17 00:00:00 2001
From: mgrapotte <mathysgrapotte@gmail.com>
Date: Wed, 22 Jan 2025 20:14:08 +0100
Subject: [PATCH 14/14] FIX: make type checker happy

---
 config/mypy.ini                               |  1 +
 src/stimulus/analysis/analysis_default.py     | 17 +++--
 src/stimulus/cli/predict.py                   |  7 +-
 src/stimulus/cli/split_yaml.py                | 11 +--
 src/stimulus/data/data_handlers.py            | 14 ++--
 src/stimulus/data/encoding/encoders.py        | 72 +++++++++----------
 src/stimulus/data/experiments.py              |  3 +-
 src/stimulus/data/handlertorch.py             |  2 +-
 src/stimulus/data/splitters/splitters.py      |  5 +-
 .../data_transformation_generators.py         | 16 ++---
 src/stimulus/learner/predict.py               | 33 ++++++---
 src/stimulus/learner/raytune_learner.py       | 63 ++++++++++------
 src/stimulus/learner/raytune_parser.py        | 56 +++++++++++----
 src/stimulus/utils/launch_utils.py            |  6 +-
 src/stimulus/utils/performance.py             | 69 +++++++++---------
 src/stimulus/utils/yaml_data.py               | 33 +++++----
 src/stimulus/utils/yaml_model_schema.py       | 17 ++---
 tests/cli/test_split_yaml.py                  |  7 +-
 tests/data/encoding/test_encoders.py          | 20 +++---
 .../data/transform/test_data_transformers.py  |  8 +--
 tests/test_model/dnatofloat_model.py          | 19 ++---
 tests/test_model/titanic_model.py             | 11 +--
 22 files changed, 286 insertions(+), 204 deletions(-)

diff --git a/config/mypy.ini b/config/mypy.ini
index 814e2ac8..21415724 100644
--- a/config/mypy.ini
+++ b/config/mypy.ini
@@ -3,3 +3,4 @@ ignore_missing_imports = true
 exclude = tests/fixtures/
 warn_unused_ignores = true
 show_error_codes = true
+explicit_package_bases = True
diff --git a/src/stimulus/analysis/analysis_default.py b/src/stimulus/analysis/analysis_default.py
index b5af1f7d..86306e50 100644
--- a/src/stimulus/analysis/analysis_default.py
+++ b/src/stimulus/analysis/analysis_default.py
@@ -1,12 +1,12 @@
 """Default analysis module for stimulus package."""
 
 import math
-from typing import Any
+from typing import Any, Union
 
-import matplotlib as mpl
 import numpy as np
 import pandas as pd
 from matplotlib import pyplot as plt
+from matplotlib.ticker import StrMethodFormatter
 from torch.utils.data import DataLoader
 
 from stimulus.data.handlertorch import TorchDataset
@@ -66,8 +66,11 @@ def heatmap(
         im = ax.imshow(data, **kwargs)
 
         # Create colorbar
-        cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw)
-        cbar.ax.set_ylabel(cbarlabel, rotation=-90, va="bottom")
+        if ax.figure is not None and hasattr(ax.figure, "colorbar"):
+            cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw)
+            cbar.ax.set_ylabel(cbarlabel, rotation=-90, va="bottom")
+        else:
+            cbar = None
 
         # Show all ticks and label them with the respective list entries.
         ax.set_xticks(np.arange(data.shape[1]), labels=col_labels)
@@ -93,7 +96,7 @@ def heatmap(
     def annotate_heatmap(
         im: Any,
         data: np.ndarray | None = None,
-        valfmt: str = "{x:.2f}",
+        valfmt: Union[str, StrMethodFormatter] = "{x:.2f}",
         textcolors: tuple[str, str] = ("black", "white"),
         threshold: float | None = None,
         **textkw: Any,
@@ -134,7 +137,7 @@ def annotate_heatmap(
 
         # Get the formatter in case a string is supplied
         if isinstance(valfmt, str):
-            valfmt = mpl.ticker.StrMethodFormatter(valfmt)
+            valfmt = StrMethodFormatter(valfmt)
 
         # Loop over the data and create a `Text` for each "pixel".
         # Change the text's color depending on the data.
@@ -142,7 +145,7 @@ def annotate_heatmap(
         for i in range(data.shape[0]):
             for j in range(data.shape[1]):
                 kw.update(color=textcolors[int(im.norm(data[i, j]) > threshold)])
-                text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)
+                text = im.axes.text(j, i, valfmt(data[i, j]), **kw)
                 texts.append(text)
 
         return texts
diff --git a/src/stimulus/cli/predict.py b/src/stimulus/cli/predict.py
index 08e165bc..7834151e 100755
--- a/src/stimulus/cli/predict.py
+++ b/src/stimulus/cli/predict.py
@@ -140,8 +140,8 @@ def main(
     data_path: str,
     output: str,
     *,
-    return_labels: bool,
-    split: int | None,
+    return_labels: bool = False,
+    split: int | None = None,
 ) -> None:
     """Run model prediction pipeline.
 
@@ -171,7 +171,8 @@ def main(
         shuffle=False,
     )
 
-    out = PredictWrapper(model, dataloader).predict(return_labels=return_labels)
+    predictor = PredictWrapper(model, dataloader)
+    out = predictor.predict(return_labels=return_labels)
     y_pred, y_true = out if return_labels else (out, {})
 
     y_pred = {k: v.tolist() for k, v in y_pred.items()}
diff --git a/src/stimulus/cli/split_yaml.py b/src/stimulus/cli/split_yaml.py
index c7114e94..9da76f38 100755
--- a/src/stimulus/cli/split_yaml.py
+++ b/src/stimulus/cli/split_yaml.py
@@ -7,6 +7,7 @@
 """
 
 import argparse
+from typing import Any
 
 import yaml
 
@@ -44,7 +45,7 @@ def get_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def main(config_yaml: str, out_dir_path: str) -> str:
+def main(config_yaml: str, out_dir_path: str) -> None:
     """Reads a YAML config file and generates all possible data configurations.
 
     This script reads a YAML with a defined structure and creates all the YAML files ready to be passed to
@@ -58,16 +59,16 @@ def main(config_yaml: str, out_dir_path: str) -> str:
     and uses the default split behavior.
     """
     # read the yaml experiment config and load it to dictionary
-    yaml_config = {}
+    yaml_config: dict[str, Any] = {}
     with open(config_yaml) as conf_file:
         yaml_config = yaml.safe_load(conf_file)
 
+    yaml_config_dict: YamlConfigDict = YamlConfigDict(**yaml_config)
     # check if the yaml schema is correct
-    check_yaml_schema(yaml_config)
+    check_yaml_schema(yaml_config_dict)
 
     # generate all the YAML configs
-    config_dict = YamlConfigDict(**yaml_config)
-    data_configs = generate_data_configs(config_dict)
+    data_configs = generate_data_configs(yaml_config_dict)
 
     # dump all the YAML configs into files
     dump_yaml_list_into_files(data_configs, out_dir_path, "test")
diff --git a/src/stimulus/data/data_handlers.py b/src/stimulus/data/data_handlers.py
index f440cc5c..9a582ab3 100644
--- a/src/stimulus/data/data_handlers.py
+++ b/src/stimulus/data/data_handlers.py
@@ -93,7 +93,7 @@ def categorize_columns_by_type(self) -> dict:
 
         return {"input": input_columns, "label": label_columns, "meta": meta_columns}
 
-    def _load_config(self, config_path: str) -> dict:
+    def _load_config(self, config_path: str) -> yaml_data.YamlConfigDict:
         """Loads and parses a YAML configuration file.
 
         Args:
@@ -111,7 +111,7 @@ def _load_config(self, config_path: str) -> dict:
         with open(config_path) as file:
             return yaml_data.YamlSubConfigDict(**yaml.safe_load(file))
 
-    def get_split_columns(self) -> str:
+    def get_split_columns(self) -> list[str]:
         """Get the columns that are used for splitting."""
         return self.config.split.split_input_columns
 
@@ -281,6 +281,7 @@ def __init__(
         """
         self.dataset_manager = DatasetManager(config_path)
         self.columns = self.read_csv_header(csv_path)
+        self.data = self.load_csv(csv_path)
 
     def read_csv_header(self, csv_path: str) -> list:
         """Get the column names from the header of the CSV file.
@@ -383,7 +384,7 @@ def shuffle_labels(self, seed: Optional[float] = None) -> None:
         # set the np seed
         np.random.seed(seed)
 
-        label_keys = self.dataset_manager.get_label_columns()["label"]
+        label_keys = self.dataset_manager.column_categories["label"]
         for key in label_keys:
             self.data = self.data.with_columns(pl.Series(key, np.random.permutation(list(self.data[key]))))
 
@@ -432,9 +433,9 @@ def get_all_items(self) -> tuple[dict, dict, dict]:
         meta_data = {key: self.data[key].to_list() for key in meta_columns}
         return input_data, label_data, meta_data
 
-    def get_all_items_and_length(self) -> tuple[dict, dict, dict, int]:
+    def get_all_items_and_length(self) -> tuple[tuple[dict, dict, dict], int]:
         """Get the full dataset as three separate dictionaries for inputs, labels and metadata, and the length of the data."""
-        return self.get_all_items(), len(self)
+        return self.get_all_items(), len(self.data)
 
     def load_csv_per_split(self, csv_path: str, split: int) -> pl.DataFrame:
         """Load the part of csv file that has the specified split value.
@@ -455,7 +456,7 @@ def __len__(self) -> int:
         """Return the length of the first list in input, assumes that all are the same length."""
         return len(self.data)
 
-    def __getitem__(self, idx: Any) -> dict:
+    def __getitem__(self, idx: Any) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor], dict[str, list]]:
         """Get the data at a given index, and encodes the input and label, leaving meta as it is.
 
         Args:
@@ -465,7 +466,6 @@ def __getitem__(self, idx: Any) -> dict:
         if isinstance(idx, slice):
             data_at_index = self.data.slice(idx.start or 0, idx.stop or len(self.data))
         elif isinstance(idx, int):
-            # Convert single row to DataFrame to maintain consistent interface
             data_at_index = self.data.slice(idx, idx + 1)
         else:
             data_at_index = self.data[idx]
diff --git a/src/stimulus/data/encoding/encoders.py b/src/stimulus/data/encoding/encoders.py
index e607fe44..8c1b3fe9 100644
--- a/src/stimulus/data/encoding/encoders.py
+++ b/src/stimulus/data/encoding/encoders.py
@@ -33,24 +33,24 @@ def encode(self, data: Any) -> Any:
         This is an abstract method, child classes should overwrite it.
 
         Args:
-            data (any): a single data point
+            data (Any): a single data point
 
         Returns:
-            encoded_data_point (any): the encoded data point
+            encoded_data_point (Any): the encoded data point
         """
         raise NotImplementedError
 
     @abstractmethod
-    def encode_all(self, data: list) -> np.array:
+    def encode_all(self, data: list[Any]) -> torch.Tensor:
         """Encode a list of data points.
 
         This is an abstract method, child classes should overwrite it.
 
         Args:
-            data (list): a list of data points
+            data (list[Any]): a list of data points
 
         Returns:
-            encoded_data (np.array): encoded data points
+            encoded_data (torch.Tensor): encoded data points
         """
         raise NotImplementedError
 
@@ -61,21 +61,21 @@ def decode(self, data: Any) -> Any:
         This is an abstract method, child classes should overwrite it.
 
         Args:
-            data (any): a single encoded data point
+            data (Any): a single encoded data point
 
         Returns:
-            decoded_data_point (any): the decoded data point
+            decoded_data_point (Any): the decoded data point
         """
         raise NotImplementedError
 
-    def encode_multiprocess(self, data: list) -> list:
+    def encode_multiprocess(self, data: list[Any]) -> list[Any]:
         """Helper function for encoding the data using multiprocessing.
 
         Args:
-            data (list): a list of data points
+            data (list[Any]): a list of data points
 
         Returns:
-            encoded_data (list): encoded data points
+            encoded_data (list[Any]): encoded data points
         """
         with mp.Pool(mp.cpu_count()) as pool:
             return pool.map(self.encode, data)
@@ -128,14 +128,14 @@ def __init__(self, alphabet: str = "acgt", *, convert_lowercase: bool = False, p
         )  # handle_unknown='ignore' unsures that a vector of zeros is returned for unknown characters, such as 'Ns' in DNA sequences
         self.encoder.fit(np.array(list(alphabet)).reshape(-1, 1))
 
-    def _sequence_to_array(self, sequence: str) -> np.array:
+    def _sequence_to_array(self, sequence: str) -> np.ndarray:
         """This function transforms the given sequence to an array.
 
         Args:
             sequence (str): a sequence of characters.
 
         Returns:
-            sequence_array (np.array): the sequence as a numpy array
+            sequence_array (np.ndarray): the sequence as a numpy array
 
         Raises:
             TypeError: If the input data is not a string.
@@ -211,7 +211,7 @@ def encode_all(self, data: Union[str, list[str]]) -> torch.Tensor:
         Unknown characters are represented by a vector of zeros.
 
         Args:
-            data (Union[list, str]): list of sequences or a single sequence
+            data (Union[str, list[str]]): list of sequences or a single sequence
 
         Returns:
             encoded_data (torch.Tensor): one hot encoded sequences
@@ -241,7 +241,7 @@ def encode_all(self, data: Union[str, list[str]]) -> torch.Tensor:
             return torch.stack([encoded_data])
         if isinstance(data, list):
             # TODO instead maybe we can run encode_multiprocess when data size is larger than a certain threshold.
-            encoded_data = self.encode_multiprocess(data)
+            encoded_data = self.encode_multiprocess(data)  # type: ignore[assignment]
         else:
             error_msg = f"Expected list or string input for data, got {type(data).__name__}"
             logger.error(error_msg)
@@ -250,7 +250,7 @@ def encode_all(self, data: Union[str, list[str]]) -> torch.Tensor:
         # handle padding
         if self.padding:
             max_length = max([len(d) for d in encoded_data])
-            encoded_data = [np.pad(d, ((0, max_length - len(d)), (0, 0))) for d in encoded_data]
+            encoded_data = [np.pad(d, ((0, max_length - len(d)), (0, 0))) for d in encoded_data]  # type: ignore[assignment]
         else:
             lengths = {len(d) for d in encoded_data}
             if len(lengths) > 1:
@@ -271,7 +271,7 @@ def decode(self, data: torch.Tensor) -> Union[str, list[str]]:
         NOTE that when decoding 3D shape tensor, it assumes all sequences have the same length.
 
         Returns:
-            Union[str, List[str]]: Single sequence string or list of sequence strings
+            Union[str, list[str]]: Single sequence string or list of sequence strings
 
         Raises:
             TypeError: If the input data is not a 2D or 3D tensor
@@ -321,20 +321,20 @@ def encode(self, data: float) -> torch.Tensor:
         This method takes as input a single data point, should be mappable to a single output.
 
         Args:
-            data (float or int): a single data point
+            data (float): a single data point
 
         Returns:
             encoded_data_point (torch.Tensor): the encoded data point
         """
-        return self.encode_all(data)  # there is no difference in this case
+        return self.encode_all([data])
 
-    def encode_all(self, data: float | list[float]) -> torch.Tensor:
+    def encode_all(self, data: list[float]) -> torch.Tensor:
         """Encodes the data.
 
         This method takes as input a list of data points, or a single float, and returns a torch.tensor.
 
         Args:
-            data (float or int): a list of data points or a single data point
+            data (list[float]): a list of data points or a single data point
 
         Returns:
             encoded_data (torch.Tensor): the encoded data
@@ -354,15 +354,15 @@ def decode(self, data: torch.Tensor) -> list[float]:
             data (torch.Tensor): the encoded data
 
         Returns:
-            decoded_data (List[float]): the decoded data
+            decoded_data (list[float]): the decoded data
         """
         return data.cpu().numpy().tolist()
 
-    def _check_input_dtype(self, data: Union[list[float], list[int]]) -> None:
+    def _check_input_dtype(self, data: list[float]) -> None:
         """Check if the input data is int or float data.
 
         Args:
-            data (float or int): a list of float or integer data points
+            data (list[float]): a list of float data points
 
         Raises:
             ValueError: If the input data contains a non-integer or non-float data point
@@ -372,11 +372,11 @@ def _check_input_dtype(self, data: Union[list[float], list[int]]) -> None:
             logger.error(err_msg)
             raise ValueError(err_msg)
 
-    def _warn_float_is_converted_to_int(self, data: Union[list[float], list[int]]) -> None:
+    def _warn_float_is_converted_to_int(self, data: list[float]) -> None:
         """Warn if float data is encoded into int data.
 
         Args:
-            data (float or int): a list of float or integer data points
+            data (list[float]): a list of float data points
         """
         if any(isinstance(d, float) for d in data) and (
             self.dtype in [torch.int, torch.int8, torch.int16, torch.int32, torch.int64]
@@ -395,12 +395,12 @@ class StrClassificationEncoder(AbstractEncoder):
     Methods:
         encode(data: str) -> int:
             Raises a NotImplementedError, as encoding a single string is not meaningful in this context.
-        encode_all(data: List[str]) -> torch.tensor:
+        encode_all(data: list[str]) -> torch.tensor:
             Encodes an entire list of string data into a numeric representation using LabelEncoder and
             returns a torch tensor. Ensures that the provided data items are valid strings prior to encoding.
         decode(data: Any) -> Any:
             Raises a NotImplementedError, as decoding is not supported with the current design.
-        _check_dtype(data: List[str]) -> None:
+        _check_dtype(data: list[str]) -> None:
             Validates that all items in the data list are strings, raising a ValueError otherwise.
     """
 
@@ -420,14 +420,14 @@ def encode(self, data: str) -> int:
         """
         raise NotImplementedError("Encoding a single string does not make sense. Use encode_all instead.")
 
-    def encode_all(self, data: list[str]) -> torch.tensor:
+    def encode_all(self, data: Union[str, list[str]]) -> torch.Tensor:
         """Encodes the data.
 
         This method takes as input a list of data points, should be mappable to a single output, using LabelEncoder from scikit learn and returning a numpy array.
         For more info visit : https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.LabelEncoder.html
 
         Args:
-            data (List[str]): a list of strings
+            data (Union[str, list[str]]): a list of strings or single string
 
         Returns:
             encoded_data (torch.tensor): the encoded data
@@ -452,7 +452,7 @@ def _check_dtype(self, data: list[str]) -> None:
         """Check if the input data is string data.
 
         Args:
-            data (List[str]): a list of strings
+            data (list[str]): a list of strings
 
         Raises:
             ValueError: If the input data is not a string
@@ -488,14 +488,14 @@ def encode(self, data: Any) -> torch.Tensor:
         """Returns an error since encoding a single float does not make sense."""
         raise NotImplementedError("Encoding a single float does not make sense. Use encode_all instead.")
 
-    def encode_all(self, data: Union[list[float], list[int]]) -> torch.Tensor:
+    def encode_all(self, data: list[Union[int, float]]) -> torch.Tensor:
         """Encodes the data.
 
         This method takes as input a list of data points, and returns the ranks of the data points.
         The ranks are normalized to be between 0 and 1, when scale is set to True.
 
         Args:
-            data (Union[List[float], List[int]]): a list of numeric values
+            data (list[Union[int, float]]): a list of numeric values
 
         Returns:
             encoded_data (torch.Tensor): the encoded data
@@ -506,8 +506,8 @@ def encode_all(self, data: Union[list[float], list[int]]) -> torch.Tensor:
 
         # Get ranks (0 is lowest, n-1 is highest)
         # and normalize to be between 0 and 1
-        data = np.array(data)
-        ranks = np.argsort(np.argsort(data))
+        array_data: np.ndarray = np.array(data)
+        ranks: np.ndarray = np.argsort(np.argsort(array_data))
         if self.scale:
             ranks = ranks / max(len(ranks) - 1, 1)
         return torch.tensor(ranks)
@@ -516,11 +516,11 @@ def decode(self, data: Any) -> Any:
         """Returns an error since decoding does not make sense without encoder information, which is not yet supported."""
         raise NotImplementedError("Decoding is not yet supported for NumericRank.")
 
-    def _check_input_dtype(self, data: list) -> None:
+    def _check_input_dtype(self, data: list[Union[int, float]]) -> None:
         """Check if the input data is int or float data.
 
         Args:
-            data (int or float): a single data point or a list of data points
+            data (list[Union[int, float]]): a list of numeric values
 
         Raises:
             ValueError: If the input data is not a float
diff --git a/src/stimulus/data/experiments.py b/src/stimulus/data/experiments.py
index 1e58ae4d..d962ac94 100644
--- a/src/stimulus/data/experiments.py
+++ b/src/stimulus/data/experiments.py
@@ -138,7 +138,8 @@ def set_data_transformer_as_attribute(self, field_name: str, data_transformer: A
         if not hasattr(self, field_name):
             setattr(self, field_name, {data_transformer.__class__.__name__: data_transformer})
         else:
-            self.field_name[data_transformer.__class__.__name__] = data_transformer
+            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.
diff --git a/src/stimulus/data/handlertorch.py b/src/stimulus/data/handlertorch.py
index 5573363b..0c608072 100644
--- a/src/stimulus/data/handlertorch.py
+++ b/src/stimulus/data/handlertorch.py
@@ -15,7 +15,7 @@ def __init__(
         config_path: str,
         csv_path: str,
         encoder_loader: experiments.EncoderLoader,
-        split: Optional[tuple[None, int]] = None,
+        split: Optional[int] = None,
     ) -> None:
         """Initialize the TorchDataset.
 
diff --git a/src/stimulus/data/splitters/splitters.py b/src/stimulus/data/splitters/splitters.py
index ce9709fd..5b429fc5 100644
--- a/src/stimulus/data/splitters/splitters.py
+++ b/src/stimulus/data/splitters/splitters.py
@@ -4,7 +4,6 @@
 from typing import Any, Optional
 
 import numpy as np
-import polars as pl
 
 # Constants
 SPLIT_SIZE = 3  # Number of splits (train/val/test)
@@ -29,7 +28,7 @@ def __init__(self, seed: float = 42) -> None:
         self.seed = seed
 
     @abstractmethod
-    def get_split_indexes(self, data: pl.DataFrame) -> list:
+    def get_split_indexes(self, data: dict) -> tuple[list, list, list]:
         """Splits the data. Always return indices mapping to the original list.
 
         This is an abstract method that should be implemented by the child class.
@@ -61,7 +60,7 @@ def distance(self, data_one: Any, data_two: Any) -> float:
 class RandomSplit(AbstractSplitter):
     """This splitter randomly splits the data."""
 
-    def __init__(self, split: Optional[list] = None, seed: Optional[float] = None) -> None:
+    def __init__(self, split: Optional[list] = None, seed: int = 42) -> None:
         """Initialize the random splitter.
 
         Args:
diff --git a/src/stimulus/data/transform/data_transformation_generators.py b/src/stimulus/data/transform/data_transformation_generators.py
index 9eedeb40..afc4d895 100644
--- a/src/stimulus/data/transform/data_transformation_generators.py
+++ b/src/stimulus/data/transform/data_transformation_generators.py
@@ -30,8 +30,8 @@ class AbstractDataTransformer(ABC):
 
     def __init__(self) -> None:
         """Initialize the data transformer."""
-        self.add_row = None
-        self.seed = 42
+        self.add_row: bool = False
+        self.seed: int = 42
 
     @abstractmethod
     def transform(self, data: Any) -> Any:
@@ -98,7 +98,7 @@ class UniformTextMasker(AbstractNoiseGenerator):
         transform_all: adds character masking to a list of data points
     """
 
-    def __init__(self, probability: float = 0.1, mask: str = "*", seed: float = 42) -> None:
+    def __init__(self, probability: float = 0.1, mask: str = "*", seed: int = 42) -> None:
         """Initialize the text masker.
 
         Args:
@@ -148,7 +148,7 @@ class GaussianNoise(AbstractNoiseGenerator):
         transform_all: adds noise to a list of data points
     """
 
-    def __init__(self, mean: float = 0, std: float = 1, seed: float = 42) -> None:
+    def __init__(self, mean: float = 0, std: float = 1, seed: int = 42) -> None:
         """Initialize the Gaussian noise generator.
 
         Args:
@@ -173,17 +173,17 @@ def transform(self, data: float) -> float:
         np.random.seed(self.seed)
         return data + np.random.normal(self.mean, self.std)
 
-    def transform_all(self, data: list) -> np.array:
+    def transform_all(self, data: list) -> list:
         """Adds Gaussian noise to a list of data points.
 
         Args:
             data (list): the data to be transformed
 
         Returns:
-            transformed_data (np.array): the transformed data points
+            transformed_data (list): the transformed data points
         """
         np.random.seed(self.seed)
-        return np.array(np.array(data) + np.random.normal(self.mean, self.std, len(data)))
+        return list(np.array(data) + np.random.normal(self.mean, self.std, len(data)))
 
 
 class ReverseComplement(AbstractAugmentationGenerator):
@@ -254,7 +254,7 @@ class GaussianChunk(AbstractAugmentationGenerator):
         transform_all: chunks multiple lists
     """
 
-    def __init__(self, chunk_size: int, seed: float = 42, std: float = 1) -> None:
+    def __init__(self, chunk_size: int, seed: int = 42, std: float = 1) -> None:
         """Initialize the Gaussian chunk generator.
 
         Args:
diff --git a/src/stimulus/learner/predict.py b/src/stimulus/learner/predict.py
index cbc9e06d..a3bc7a3b 100644
--- a/src/stimulus/learner/predict.py
+++ b/src/stimulus/learner/predict.py
@@ -1,8 +1,10 @@
 """A module for making predictions with PyTorch models using DataLoaders."""
 
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import torch
+from torch import Tensor, nn
+from torch.utils.data import DataLoader
 
 from stimulus.utils.generic_utils import ensure_at_least_1d
 from stimulus.utils.performance import Performance
@@ -14,7 +16,7 @@ class PredictWrapper:
     It also provides the functionalities to measure the performance of the model.
     """
 
-    def __init__(self, model: object, dataloader: object, loss_dict: Optional[dict[str, Any]] = None) -> None:
+    def __init__(self, model: nn.Module, dataloader: DataLoader, loss_dict: Optional[dict[str, Any]] = None) -> None:
         """Initialize the PredictWrapper.
 
         Args:
@@ -33,7 +35,11 @@ def __init__(self, model: object, dataloader: object, loss_dict: Optional[dict[s
 
             logging.warning("Not able to run model.eval: %s", str(e))
 
-    def predict(self, *, return_labels: bool = False) -> dict[str, torch.Tensor]:
+    def predict(
+        self,
+        *,
+        return_labels: bool = False,
+    ) -> Union[dict[str, Tensor], tuple[dict[str, Tensor], dict[str, Tensor]]]:
         """Get the model predictions.
 
         Basically, it runs a foward pass on the model for each batch,
@@ -54,8 +60,8 @@ def predict(self, *, return_labels: bool = False) -> dict[str, torch.Tensor]:
         # create empty dictionaries with the column names
         first_batch = next(iter(self.dataloader))
         keys = first_batch[1].keys()
-        predictions = {k: [] for k in keys}
-        labels = {k: [] for k in keys}
+        predictions: dict[str, list[Tensor]] = {k: [] for k in keys}
+        labels: dict[str, list[Tensor]] = {k: [] for k in keys}
 
         # get the predictions (and labels) for each batch
         with torch.no_grad():
@@ -73,7 +79,7 @@ def predict(self, *, return_labels: bool = False) -> dict[str, torch.Tensor]:
             return {k: torch.cat(v) for k, v in predictions.items()}
         return {k: torch.cat(v) for k, v in predictions.items()}, {k: torch.cat(v) for k, v in labels.items()}
 
-    def handle_predictions(self, predictions: Any, y: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    def handle_predictions(self, predictions: Any, y: dict[str, Tensor]) -> dict[str, Tensor]:
         """Handle the model outputs from forward pass, into a dictionary of tensors, just like y."""
         if len(y) == 1:
             return {next(iter(y.keys())): predictions}
@@ -111,8 +117,15 @@ def compute_other_metric(self, metric: str) -> float:
 
         # TODO currently we computes the average performance metric across target y, but maybe in the future we want something different
         """
-        if (not hasattr(self, "predictions")) or (not hasattr(self, "labels")):
-            self.predictions, self.labels = self.predict(return_labels=True)
+        if not hasattr(self, "predictions") or not hasattr(self, "labels"):
+            predictions, labels = self.predict(return_labels=True)
+            self.predictions = predictions
+            self.labels = labels
+
+        # Explicitly type the labels and predictions as dictionaries with str keys
+        labels_dict: dict[str, Tensor] = self.labels if isinstance(self.labels, dict) else {}
+        predictions_dict: dict[str, Tensor] = self.predictions if isinstance(self.predictions, dict) else {}
+
         return sum(
-            Performance(labels=self.labels[k], predictions=self.predictions[k], metric=metric).val for k in self.labels
-        ) / len(self.labels)
+            Performance(labels=labels_dict[k], predictions=predictions_dict[k], metric=metric).val for k in labels_dict
+        ) / len(labels_dict)
diff --git a/src/stimulus/learner/raytune_learner.py b/src/stimulus/learner/raytune_learner.py
index a42dab96..735348d3 100644
--- a/src/stimulus/learner/raytune_learner.py
+++ b/src/stimulus/learner/raytune_learner.py
@@ -4,7 +4,7 @@
 import logging
 import os
 import random
-from typing import Optional
+from typing import Any, Optional, TypedDict
 
 import numpy as np
 import torch
@@ -13,7 +13,7 @@
 from safetensors.torch import load_model as safe_load_model
 from safetensors.torch import save_model as safe_save_model
 from torch import nn, optim
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 
 from stimulus.data.handlertorch import TorchDataset
 from stimulus.learner.predict import PredictWrapper
@@ -21,6 +21,12 @@
 from stimulus.utils.yaml_model_schema import YamlRayConfigLoader
 
 
+class CheckpointDict(TypedDict):
+    """Dictionary type for checkpoint data."""
+
+    checkpoint_dir: str
+
+
 class TuneWrapper:
     """Wrapper class for Ray Tune hyperparameter optimization."""
 
@@ -82,7 +88,7 @@ def __init__(
 
         # working towards the path for the tune_run directory. if ray_results_dir None ray will put it under home so we will do the same here.
         if ray_results_dir is None:
-            ray_results_dir = os.environ.get("HOME")
+            ray_results_dir = os.environ.get("HOME", "")
         # then we are able to pass the whole correct tune_run path to the trainable function. so it can use thaqt to place the debug dir under if needed.
         self.config["tune_run_path"] = os.path.join(ray_results_dir, tune_run_name)
 
@@ -130,21 +136,25 @@ def tuner_initialization(self) -> tune.Tuner:
 
     def tune(self) -> None:
         """Run the tuning process."""
-        return self.tuner.fit()
+        self.tuner.fit()
 
     def _chek_per_trial_resources(
         self,
         resurce_key: str,
-        cluster_max_resources: dict,
+        cluster_max_resources: dict[str, float],
         resource_type: str,
-    ) -> tuple[int, int]:
+    ) -> float:
         """Helper function that check that user requested per trial resources are not exceeding the available resources for the ray cluster.
 
         If the per trial resources are not asked they are set to a default resoanable ammount.
 
-        resurce_key:            str object          the key used to look into the self.config["tune"]
-        cluster_max_resources:  dict object         the output of the ray.cluster_resources() function. It hold what ray has found to be the available resources for CPU, GPU and Memory
-        resource_type:          str object          the key used to llok into the cluster_resources dict
+        Args:
+            resurce_key: The key used to look into the self.config["tune"]
+            cluster_max_resources: The output of the ray.cluster_resources() function. It hold what ray has found to be the available resources for CPU, GPU and Memory
+            resource_type: The key used to llok into the cluster_resources dict
+
+        Returns:
+            The amount of resources per trial to use
         """
         if resource_type == "GPU" and resource_type not in cluster_resources():
             # ray does not have a GPU field also if GPUs were set to zero. So trial GPU resources have to be set to zero.
@@ -155,13 +165,13 @@ def _chek_per_trial_resources(
                 "#### ray did not detect any GPU, if you do not want to use GPU set max_gpus=0, or in nextflow --max_gpus 0.",
             )
 
-        per_trial_resource = None
+        per_trial_resource: float = 0.0
         # if everything is alright, leave the value as it is.
         if (
             resurce_key in self.config["tune"]
             and self.config["tune"][resurce_key] <= cluster_max_resources[resource_type]
         ):
-            per_trial_resource = self.config["tune"][resurce_key]
+            per_trial_resource = float(self.config["tune"][resurce_key])
 
         # if per_trial_resource are more than what is avaialble to ray set them to what is available and warn the user
         elif (
@@ -175,18 +185,20 @@ def _chek_per_trial_resources(
                 f"available: {cluster_max_resources[resource_type]} "
                 "overwriting value to max available",
             )
-            per_trial_resource = cluster_max_resources[resource_type]
+            per_trial_resource = float(cluster_max_resources[resource_type])
 
         # if per_trial_resource has not been asked and there is none available set them to zero
         elif resurce_key not in self.config["tune"] and cluster_max_resources[resource_type] == 0.0:
-            per_trial_resource = 0
+            per_trial_resource = 0.0
 
         # if per_trial_resource has not been asked and the resource is available set the value to either 1 or number_available resource / num_samples
         elif resurce_key not in self.config["tune"] and cluster_max_resources[resource_type] != 0.0:
             # TODO maybe set the default to 0.5 instead of 1 ? fractional use in case of GPU? Should this be a mandatory parameter?
-            per_trial_resource = max(
-                1,
-                (cluster_max_resources[resource_type] // self.config["tune"]["tune_params"]["num_samples"]),
+            per_trial_resource = float(
+                max(
+                    1,
+                    (cluster_max_resources[resource_type] // self.config["tune"]["tune_params"]["num_samples"]),
+                ),
             )
 
         return per_trial_resource
@@ -195,7 +207,7 @@ def _chek_per_trial_resources(
 class TuneModel(Trainable):
     """Trainable model class for Ray Tune."""
 
-    def setup(self, config: dict, training: object, validation: object) -> None:
+    def setup(self, config: dict[Any, Any]) -> None:
         """Get the model, loss function(s), optimizer, train and test data from the config."""
         # set the seeds the second time, first in TuneWrapper initialization. This will make all important seed worker specific.
         set_general_seeds(self.config["ray_worker_seed"])
@@ -229,6 +241,8 @@ def setup(self, config: dict, training: object, validation: object) -> None:
 
         # use dataloader on training/validation data
         self.batch_size = config["data_params"]["batch_size"]
+        training: Dataset = config["training"]
+        validation: Dataset = config["validation"]
         self.training = DataLoader(
             training,
             batch_size=self.batch_size,
@@ -272,7 +286,7 @@ def step(self) -> dict:
                 self.model.batch(x=x, y=y, optimizer=self.optimizer, **self.loss_dict)
         return self.objective()
 
-    def objective(self) -> dict:
+    def objective(self) -> dict[str, float]:
         """Compute the objective metric(s) for the tuning process."""
         metrics = [
             "loss",
@@ -291,17 +305,22 @@ def objective(self) -> dict:
             **{"train_" + metric: value for metric, value in predict_train.compute_metrics(metrics).items()},
         }
 
-    def export_model(self, export_dir: str) -> None:
+    def export_model(self, export_dir: str | None = None) -> None:  # type: ignore[override]
         """Export model to safetensors format."""
+        if export_dir is None:
+            return
         safe_save_model(self.model, os.path.join(export_dir, "model.safetensors"))
 
-    def load_checkpoint(self, checkpoint_dir: str) -> None:
+    def load_checkpoint(self, checkpoint: dict[Any, Any] | None) -> None:
         """Load model and optimizer state from checkpoint."""
+        if checkpoint is None:
+            return
+        checkpoint_dir = checkpoint["checkpoint_dir"]
         self.model = safe_load_model(self.model, os.path.join(checkpoint_dir, "model.safetensors"))
         self.optimizer.load_state_dict(torch.load(os.path.join(checkpoint_dir, "optimizer.pt")))
 
-    def save_checkpoint(self, checkpoint_dir: str) -> dict | None:
+    def save_checkpoint(self, checkpoint_dir: str) -> dict[Any, Any]:
         """Save model and optimizer state to checkpoint."""
         safe_save_model(self.model, os.path.join(checkpoint_dir, "model.safetensors"))
         torch.save(self.optimizer.state_dict(), os.path.join(checkpoint_dir, "optimizer.pt"))
-        return checkpoint_dir
+        return {"checkpoint_dir": checkpoint_dir}
diff --git a/src/stimulus/learner/raytune_parser.py b/src/stimulus/learner/raytune_parser.py
index c825967f..dbbaf0f8 100644
--- a/src/stimulus/learner/raytune_parser.py
+++ b/src/stimulus/learner/raytune_parser.py
@@ -2,22 +2,47 @@
 
 import json
 import os
+from typing import Any, TypedDict, cast
 
+import pandas as pd
 import torch
+from ray.tune import ExperimentAnalysis
 from safetensors.torch import load_file as safe_load_file
 from safetensors.torch import save_file as safe_save_file
 
 
+class RayTuneResult(TypedDict):
+    """TypedDict for storing Ray Tune optimization results."""
+
+    config: dict[str, Any]
+    checkpoint: str
+    metrics_dataframe: pd.DataFrame
+
+
+class RayTuneMetrics(TypedDict):
+    """TypedDict for storing Ray Tune metrics results."""
+
+    checkpoint: str
+    metrics_dataframe: pd.DataFrame
+
+
+class RayTuneOptimizer(TypedDict):
+    """TypedDict for storing Ray Tune optimizer state."""
+
+    checkpoint: str
+
+
 class TuneParser:
     """Parser class for Ray Tune results to extract best configurations and model weights."""
 
-    def __init__(self, results: object) -> None:
+    def __init__(self, results: ExperimentAnalysis) -> None:
         """`results` is the output of ray.tune."""
         self.results = results
 
-    def get_best_config(self) -> dict:
+    def get_best_config(self) -> dict[str, Any]:
         """Get the best config from the results."""
-        return self.results.get_best_result().config
+        best_result = cast(RayTuneResult, self.results.best_result)
+        return best_result["config"]
 
     def save_best_config(self, output: str) -> None:
         """Save the best config to a file.
@@ -29,7 +54,7 @@ def save_best_config(self, output: str) -> None:
         with open(output, "w") as f:
             json.dump(config, f, indent=4)
 
-    def fix_config_values(self, config: dict) -> dict:
+    def fix_config_values(self, config: dict[str, Any]) -> dict[str, Any]:
         """Correct config values.
 
         Args:
@@ -51,25 +76,28 @@ def fix_config_values(self, config: dict) -> dict:
 
     def save_best_metrics_dataframe(self, output: str) -> None:
         """Save the dataframe with the metrics at each iteration of the best sample to a file."""
-        df = self.results.get_best_result().metrics_dataframe
-        columns = [col for col in df.columns if "config" not in col]
-        df = df[columns]
-        df.to_csv(output, index=False)
+        best_result = cast(RayTuneMetrics, self.results.best_result)
+        metrics_df = best_result["metrics_dataframe"]
+        columns = [col for col in metrics_df.columns if "config" not in col]
+        metrics_df = metrics_df[columns]
+        metrics_df.to_csv(output, index=False)
 
-    def get_best_model(self) -> dict:
+    def get_best_model(self) -> dict[str, torch.Tensor]:
         """Get the best model weights from the results."""
-        checkpoint = self.results.get_best_result().checkpoint.to_directory()
-        checkpoint = os.path.join(checkpoint, "model.safetensors")
+        best_result = cast(RayTuneMetrics, self.results.best_result)
+        checkpoint_dir = best_result["checkpoint"]
+        checkpoint = os.path.join(checkpoint_dir, "model.safetensors")
         return safe_load_file(checkpoint)
 
     def save_best_model(self, output: str) -> None:
         """Save the best model weights to a file."""
         safe_save_file(self.get_best_model(), output)
 
-    def get_best_optimizer(self) -> dict:
+    def get_best_optimizer(self) -> dict[str, Any]:
         """Get the best optimizer state from the results."""
-        checkpoint = self.results.get_best_result().checkpoint.to_directory()
-        checkpoint = os.path.join(checkpoint, "optimizer.pt")
+        best_result = cast(RayTuneOptimizer, self.results.best_result)
+        checkpoint_dir = best_result["checkpoint"]
+        checkpoint = os.path.join(checkpoint_dir, "optimizer.pt")
         return torch.load(checkpoint)
 
     def save_best_optimizer(self, output: str) -> None:
diff --git a/src/stimulus/utils/launch_utils.py b/src/stimulus/utils/launch_utils.py
index 24b12b96..13574a77 100644
--- a/src/stimulus/utils/launch_utils.py
+++ b/src/stimulus/utils/launch_utils.py
@@ -27,7 +27,11 @@ def import_class_from_file(file_path: str) -> type:
     # Create a module from the file path
     # In summary, these three lines of code are responsible for creating a module specification based on a file location, creating a module object from that specification, and then executing the module's code to populate the module object with the definitions from the Python file.
     spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Could not create module spec for {file_path}")
     module = importlib.util.module_from_spec(spec)
+    if spec.loader is None:
+        raise ImportError(f"Module spec has no loader for {file_path}")
     spec.loader.exec_module(module)
 
     # Find the class dynamically
@@ -67,7 +71,7 @@ def memory_split_for_ray_init(memory_str: Union[str, None]) -> tuple[float, floa
         tuple[float, float]: A tuple containing (store_memory, memory) in bytes.
     """
     if memory_str is None:
-        return None, None
+        return 0.0, 0.0
 
     units = {"B": 1, "K": 2**10, "M": 2**20, "G": 2**30, "T": 2**40, "P": 2**50}
 
diff --git a/src/stimulus/utils/performance.py b/src/stimulus/utils/performance.py
index c297bd79..2ac83df2 100644
--- a/src/stimulus/utils/performance.py
+++ b/src/stimulus/utils/performance.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import torch
+from numpy.typing import NDArray
 from scipy.stats import spearmanr
 from sklearn.metrics import (
     average_precision_score,
@@ -41,7 +42,7 @@ class Performance:
     metrics.
     """
 
-    def __init__(self, labels: Any, predictions: Any, metric: str = "rocauc") -> float:
+    def __init__(self, labels: Any, predictions: Any, metric: str = "rocauc") -> None:
         """Initialize Performance class with labels, predictions and metric type.
 
         Args:
@@ -49,39 +50,43 @@ def __init__(self, labels: Any, predictions: Any, metric: str = "rocauc") -> flo
             predictions: Model predictions
             metric: Type of metric to compute (default: "rocauc")
         """
-        labels = self.data2array(labels)
-        predictions = self.data2array(predictions)
-        labels, predictions = self.handle_multiclass(labels, predictions)
-        if labels.shape != predictions.shape:
+        labels_arr = self.data2array(labels)
+        predictions_arr = self.data2array(predictions)
+        labels_arr, predictions_arr = self.handle_multiclass(labels_arr, predictions_arr)
+        if labels_arr.shape != predictions_arr.shape:
             raise ValueError(
-                f"The labels have shape {labels.shape} whereas predictions have shape {predictions.shape}.",
+                f"The labels have shape {labels_arr.shape} whereas predictions have shape {predictions_arr.shape}.",
             )
         function = getattr(self, metric)
-        self.val = function(labels, predictions)
+        self.val = function(labels_arr, predictions_arr)
 
-    def data2array(self, data: Any) -> np.array:
+    def data2array(self, data: Any) -> NDArray[np.float64]:
         """Convert input data to numpy array.
 
         Args:
             data: Input data in various formats
 
         Returns:
-            np.array: Converted numpy array
+            NDArray[np.float64]: Converted numpy array
 
         Raises:
             ValueError: If input data type is not supported
         """
         if isinstance(data, list):
-            return np.array(data)
+            return np.array(data, dtype=np.float64)
         if isinstance(data, np.ndarray):
-            return data
+            return data.astype(np.float64)
         if isinstance(data, torch.Tensor):
-            return data.detach().cpu().numpy()
+            return data.detach().cpu().numpy().astype(np.float64)
         if isinstance(data, (int, float)):
-            return np.array([data])
+            return np.array([data], dtype=np.float64)
         raise ValueError(f"The data must be a list, np.array, torch.Tensor, int or float. Instead it is {type(data)}")
 
-    def handle_multiclass(self, labels: np.array, predictions: np.array) -> tuple[np.array, np.array]:
+    def handle_multiclass(
+        self,
+        labels: NDArray[np.float64],
+        predictions: NDArray[np.float64],
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
         """Handle the case of multiclass classification.
 
         TODO currently only two class predictions are handled. Needs to handle the other scenarios.
@@ -98,34 +103,34 @@ def handle_multiclass(self, labels: np.array, predictions: np.array) -> tuple[np
         # other scenarios not implemented yet
         raise ValueError(f"Labels have shape {labels.shape} and predictions have shape {predictions.shape}.")
 
-    def rocauc(self, labels: np.array, predictions: np.array) -> float:
+    def rocauc(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute ROC AUC score."""
-        return roc_auc_score(labels, predictions)
+        return float(roc_auc_score(labels, predictions))
 
-    def prauc(self, labels: np.array, predictions: np.array) -> float:
+    def prauc(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute PR AUC score."""
-        return average_precision_score(labels, predictions)
+        return float(average_precision_score(labels, predictions))
 
-    def mcc(self, labels: np.array, predictions: np.array) -> float:
+    def mcc(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute Matthews Correlation Coefficient."""
-        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
-        return matthews_corrcoef(labels, predictions)
+        predictions_binary = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
+        return float(matthews_corrcoef(labels, predictions_binary))
 
-    def f1score(self, labels: np.array, predictions: np.array) -> float:
+    def f1score(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute F1 score."""
-        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
-        return f1_score(labels, predictions)
+        predictions_binary = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
+        return float(f1_score(labels, predictions_binary))
 
-    def precision(self, labels: np.array, predictions: np.array) -> float:
+    def precision(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute precision score."""
-        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
-        return precision_score(labels, predictions)
+        predictions_binary = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
+        return float(precision_score(labels, predictions_binary))
 
-    def recall(self, labels: np.array, predictions: np.array) -> float:
+    def recall(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute recall score."""
-        predictions = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
-        return recall_score(labels, predictions)
+        predictions_binary = np.array([1 if p > BINARY_THRESHOLD else 0 for p in predictions])
+        return float(recall_score(labels, predictions_binary))
 
-    def spearmanr(self, labels: np.array, predictions: np.array) -> float:
+    def spearmanr(self, labels: NDArray[np.float64], predictions: NDArray[np.float64]) -> float:
         """Compute Spearman correlation coefficient."""
-        return spearmanr(labels, predictions)[0]
+        return float(spearmanr(labels, predictions)[0])
diff --git a/src/stimulus/utils/yaml_data.py b/src/stimulus/utils/yaml_data.py
index 30e93f38..cfb550fc 100644
--- a/src/stimulus/utils/yaml_data.py
+++ b/src/stimulus/utils/yaml_data.py
@@ -16,7 +16,7 @@ class YamlColumnsEncoder(BaseModel):
     """Model for column encoder configuration."""
 
     name: str
-    params: Optional[dict[str, Union[str, list]]]  # Allow both string and list values
+    params: Optional[dict[str, Union[str, list[Any]]]]  # Allow both string and list values
 
 
 class YamlColumns(BaseModel):
@@ -32,7 +32,7 @@ class YamlTransformColumnsTransformation(BaseModel):
     """Model for column transformation configuration."""
 
     name: str
-    params: Optional[dict[str, Union[list, float]]]  # Allow both list and float values
+    params: Optional[dict[str, Union[list[Any], float]]]  # Allow both list and float values
 
 
 class YamlTransformColumns(BaseModel):
@@ -60,7 +60,7 @@ def validate_param_lists_across_columns(cls, columns: list[YamlTransformColumns]
             The validated columns list
         """
         # Get all parameter list lengths across all columns and transformations
-        all_list_lengths = set()
+        all_list_lengths: set[int] = set()
 
         for column in columns:
             for transformation in column.transformations:
@@ -251,8 +251,8 @@ def dump_yaml_list_into_files(
     base_name: str,
 ) -> None:
     """Dumps a list of YAML configurations into separate files with custom formatting."""
-    # Disable YAML aliases to prevent reference-style output
-    yaml.Dumper.ignore_aliases = lambda *args: True
+    # 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."""
@@ -272,15 +272,22 @@ def custom_representer(dumper: yaml.Dumper, data: Any) -> yaml.Node:
     class CustomDumper(yaml.Dumper):
         """Custom YAML dumper that adds extra formatting controls."""
 
-        def write_line_break(self, data: Any = None) -> None:
+        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)
+            super().write_line_break(_data)
             if len(self.indents) <= 1:  # At root level
-                super().write_line_break(data)
+                super().write_line_break(_data)
 
-        def increase_indent(self, *, flow: bool = False, indentless: bool = False) -> bool:
+        def increase_indent(self, *, flow: bool = False, indentless: bool = False) -> None:  # type: ignore[override]
             """Ensure consistent indentation by preventing indentless sequences."""
-            return super().increase_indent(flow=flow, indentless=indentless)
+            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)
@@ -292,7 +299,7 @@ def increase_indent(self, *, flow: bool = False, indentless: bool = False) -> bo
         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 = {}
+                processed_dict: dict[str, Any] = {}
                 for key, value in input_dict.items():
                     if key == "encoder" and isinstance(value, list):
                         processed_dict[key] = []
@@ -333,14 +340,14 @@ def fix_params(input_dict: dict[str, Any]) -> dict[str, Any]:
             )
 
 
-def check_yaml_schema(config_yaml: str) -> str:
+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 (dict): The dict containing the fields of the yaml configuration file
+        config_yaml: The YamlConfigDict containing the fields of the yaml configuration file
 
     Returns:
         str: Empty string if validation succeeds
diff --git a/src/stimulus/utils/yaml_model_schema.py b/src/stimulus/utils/yaml_model_schema.py
index f6b68619..c7078ab7 100644
--- a/src/stimulus/utils/yaml_model_schema.py
+++ b/src/stimulus/utils/yaml_model_schema.py
@@ -3,6 +3,7 @@
 import random
 from collections.abc import Callable
 from copy import deepcopy
+from typing import Any
 
 import yaml
 from ray import tune
@@ -25,7 +26,7 @@ def __init__(self, config_path: str) -> None:
             self.config = yaml.safe_load(f)
             self.config = self.convert_config_to_ray(self.config)
 
-    def raytune_space_selector(self, mode: Callable, space: list) -> Callable:
+    def raytune_space_selector(self, mode: Callable, space: list) -> dict[str, Any]:
         """Convert space parameters to Ray Tune format based on the mode.
 
         Args:
@@ -46,7 +47,7 @@ def raytune_space_selector(self, mode: Callable, space: list) -> Callable:
 
         raise NotImplementedError(f"Mode {mode.__name__} not implemented yet")
 
-    def raytune_sample_from(self, mode: Callable, param: dict) -> Callable:
+    def raytune_sample_from(self, mode: Callable, param: dict) -> dict[str, Any]:
         """Apply tune.sample_from to a given custom sampling function.
 
         Args:
@@ -64,7 +65,7 @@ def raytune_sample_from(self, mode: Callable, param: dict) -> Callable:
 
         raise NotImplementedError(f"Function {param['function']} not implemented yet")
 
-    def convert_raytune(self, param: dict) -> dict:
+    def convert_raytune(self, param: dict) -> dict[str, Any]:
         """Convert parameter configuration to Ray Tune format.
 
         Args:
@@ -130,7 +131,7 @@ def get_config(self) -> dict:
         return self.config
 
     @staticmethod
-    def sampint(sample_space: list, n_space: list) -> list:
+    def sampint(sample_space: list, n_space: list) -> list[int]:
         """Return a list of n random samples from the sample_space.
 
         This function is useful for sampling different numbers of layers,
@@ -148,7 +149,7 @@ def sampint(sample_space: list, n_space: list) -> list:
             This is acceptable for hyperparameter sampling but should not be
             used for security-critical purposes (S311 fails when linting).
         """
-        sample_space = range(sample_space[0], sample_space[1] + 1)
-        n_space = range(n_space[0], n_space[1] + 1)
-        n = random.choice(tuple(n_space))  # noqa: S311
-        return random.sample(tuple(sample_space), n)
+        sample_space_list = list(range(sample_space[0], sample_space[1] + 1))
+        n_space_list = list(range(n_space[0], n_space[1] + 1))
+        n = random.choice(n_space_list)  # noqa: S311
+        return random.sample(sample_space_list, n)
diff --git a/tests/cli/test_split_yaml.py b/tests/cli/test_split_yaml.py
index f44d5c93..ad56f2b7 100644
--- a/tests/cli/test_split_yaml.py
+++ b/tests/cli/test_split_yaml.py
@@ -3,6 +3,7 @@
 import hashlib
 import os
 import tempfile
+from typing import Any, Callable
 
 import pytest
 
@@ -33,7 +34,7 @@ def wrong_yaml_path() -> str:
 @pytest.mark.parametrize(("yaml_type", "error"), test_cases)
 def test_split_yaml(
     request: pytest.FixtureRequest,
-    snapshot: pytest.fixture,
+    snapshot: Callable[[], Any],
     yaml_type: str,
     error: Exception | None,
 ) -> None:
@@ -41,10 +42,10 @@ def test_split_yaml(
     yaml_path = request.getfixturevalue(yaml_type)
     tmpdir = tempfile.gettempdir()
     if error:
-        with pytest.raises(error):
+        with pytest.raises(error):  # type: ignore[call-overload]
             main(yaml_path, tmpdir)
     else:
-        assert main(yaml_path, tmpdir) is None  # this is to assert that the function does not raise any exceptions
+        main(yaml_path, tmpdir)  # main() returns None, no need to assert
         files = os.listdir(tmpdir)
         test_out = [f for f in files if f.startswith("test_")]
         hashes = []
diff --git a/tests/data/encoding/test_encoders.py b/tests/data/encoding/test_encoders.py
index 5b970488..6422b480 100644
--- a/tests/data/encoding/test_encoders.py
+++ b/tests/data/encoding/test_encoders.py
@@ -36,11 +36,6 @@ def encoder_lowercase() -> TextOneHotEncoder:
 
     # ---- Test for initialization ---- #
 
-    def test_init_with_non_string_alphabet_raises_type_error(self) -> None:
-        """Test initialization with non-string alphabet raises TypeError."""
-        with pytest.raises(TypeError, match="Expected a string input for alphabet"):
-            TextOneHotEncoder(alphabet=["a", "c", "g", "t"])
-
     def test_init_with_string_alphabet(self) -> None:
         """Test initialization with valid string alphabet."""
         encoder = TextOneHotEncoder(alphabet="acgt")
@@ -56,14 +51,14 @@ def test_sequence_to_array_with_non_string_input(
     ) -> None:
         """Test _sequence_to_array with non-string input raises TypeError."""
         with pytest.raises(TypeError, match="Expected string input for sequence"):
-            encoder_default._sequence_to_array(1234)
+            encoder_default._sequence_to_array(1234)  # type: ignore[arg-type]
 
     def test_sequence_to_array_returns_correct_shape(
         self,
         encoder_default: TextOneHotEncoder,
     ) -> None:
         """Test _sequence_to_array returns array of correct shape."""
-        seq = "acgt"
+        seq: str = "acgt"
         arr = encoder_default._sequence_to_array(seq)
         assert arr.shape == (4, 1)
         assert (arr.flatten() == list(seq)).all()
@@ -185,6 +180,7 @@ def test_decode_unknown_characters(self, encoder_default: TextOneHotEncoder) ->
         # In the given code, it returns an empty decode for that position. So let's assume it becomes ''.
         # That means we might get "acgt" with a missing final char or a placeholder.
         # Let's do a partial check:
+        assert isinstance(decoded, str)
         assert decoded.startswith("acgt")
 
     def test_decode_multiple_sequences(self, encoder_default: TextOneHotEncoder) -> None:
@@ -259,11 +255,11 @@ def test_encode_all_single_float(self, float_encoder: NumericEncoder) -> None:
         Args:
             float_encoder: Float-based encoder instance
         """
-        input_val = 2.71
+        input_val = [2.71]
         output = float_encoder.encode_all(input_val)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
         assert output.numel() == 1, "Tensor should have exactly one element."
-        assert output.item() == pytest.approx(input_val), "Encoded value does not match the input."
+        assert output.item() == pytest.approx(input_val[0]), "Encoded value does not match the input."
 
     def test_encode_all_single_int(self, int_encoder: NumericEncoder) -> None:
         """Test encode_all when given a single int.
@@ -273,11 +269,11 @@ def test_encode_all_single_int(self, int_encoder: NumericEncoder) -> None:
         Args:
             int_encoder: Integer-based encoder instance
         """
-        input_val = 2
+        input_val = [2.0]
         output = int_encoder.encode_all(input_val)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
         assert output.numel() == 1, "Tensor should have exactly one element."
-        assert output.item() == input_val
+        assert output.item() == int(input_val[0])
 
     def test_encode_all_multi_float(self, float_encoder: NumericEncoder) -> None:
         """Test encode_all with a list of floats."""
@@ -291,7 +287,7 @@ def test_encode_all_multi_float(self, float_encoder: NumericEncoder) -> None:
 
     def test_encode_all_multi_int(self, int_encoder: NumericEncoder) -> None:
         """Test encode_all with a list of integers."""
-        input_vals = [3, 4]
+        input_vals = [3.0, 4.0]
         output = int_encoder.encode_all(input_vals)
         assert isinstance(output, torch.Tensor), "Output should be a torch.Tensor."
         assert output.dtype == torch.int32, "Tensor dtype should be int32."
diff --git a/tests/data/transform/test_data_transformers.py b/tests/data/transform/test_data_transformers.py
index 59617a0d..91215a89 100644
--- a/tests/data/transform/test_data_transformers.py
+++ b/tests/data/transform/test_data_transformers.py
@@ -50,7 +50,7 @@ def uniform_text_masker() -> DataTransformerTest:
     """Return a UniformTextMasker test object."""
     np.random.seed(42)  # Set seed before creating transformer
     transformer = UniformTextMasker(mask="N", probability=0.1)
-    params = {}  # Remove seed from params
+    params: dict[str, Any] = {}  # Remove seed from params
     single_input = "ACGTACGT"
     expected_single_output = "ACGTACNT"
     multiple_inputs = ["ATCGATCGATCG", "ATCG"]
@@ -70,7 +70,7 @@ def gaussian_noise() -> DataTransformerTest:
     """Return a GaussianNoise test object."""
     np.random.seed(42)  # Set seed before creating transformer
     transformer = GaussianNoise(mean=0, std=1)
-    params = {}  # Remove seed from params
+    params: dict[str, Any] = {}  # Remove seed from params
     single_input = 5.0
     expected_single_output = 5.4967141530112327
     multiple_inputs = [1.0, 2.0, 3.0]
@@ -90,7 +90,7 @@ def gaussian_chunk() -> DataTransformerTest:
     """Return a GaussianChunk test object."""
     np.random.seed(42)  # Set seed before creating transformer
     transformer = GaussianChunk(chunk_size=2)
-    params = {}  # Remove seed from params
+    params: dict[str, Any] = {}  # Remove seed from params
     single_input = "ACGT"
     expected_single_output = "CG"
     multiple_inputs = ["ACGT", "TGCA"]
@@ -165,7 +165,7 @@ def test_transform_multiple(self, request: Any, test_data_name: DataTransformerT
         """Test transforming multiple floats."""
         test_data = request.getfixturevalue(test_data_name)
         transformed_data = test_data.transformer.transform_all(test_data.multiple_inputs, **test_data.params)
-        assert isinstance(transformed_data, np.ndarray)
+        assert isinstance(transformed_data, list)
         for item in transformed_data:
             assert isinstance(item, float)
         assert len(transformed_data) == len(test_data.expected_multiple_outputs)
diff --git a/tests/test_model/dnatofloat_model.py b/tests/test_model/dnatofloat_model.py
index f1f84770..9e0c503e 100644
--- a/tests/test_model/dnatofloat_model.py
+++ b/tests/test_model/dnatofloat_model.py
@@ -4,6 +4,7 @@
 
 import torch
 from torch import nn
+from torch.optim import Optimizer
 
 
 class ModelSimple(torch.nn.Module):
@@ -23,7 +24,7 @@ def __init__(self, kernel_size: int = 3, pool_size: int = 2) -> None:
         self.pool = nn.MaxPool1d(pool_size, pool_size)
         self.linear = nn.Linear(49, 1)
 
-    def forward(self, hello: torch.Tensor) -> dict:
+    def forward(self, hello: torch.Tensor) -> dict[str, torch.Tensor]:
         """Forward pass of the model.
 
         It should return the output as a dictionary, with the same keys as `y`.
@@ -32,7 +33,7 @@ def forward(self, hello: torch.Tensor) -> dict:
         x = self.conv1(x)
         x = self.pool(x)
         x = self.linear(x)
-        return x.squeeze()
+        return {"output": x.squeeze()}
 
     def compute_loss(self, output: torch.Tensor, hola: torch.Tensor, loss_fn: Callable) -> torch.Tensor:
         """Compute the loss.
@@ -45,12 +46,12 @@ def compute_loss(self, output: torch.Tensor, hola: torch.Tensor, loss_fn: Callab
 
     def batch(
         self,
-        x: dict,
-        y: dict,
-        loss_fn1: Callable,
-        loss_fn2: Callable,
-        optimizer: Optional[Callable] = None,
-    ) -> tuple[torch.Tensor, dict]:
+        x: dict[str, torch.Tensor],
+        y: dict[str, torch.Tensor],
+        loss_fn1: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        loss_fn2: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        optimizer: Optional[Optimizer] = None,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         """Perform one batch step.
 
         `x` is a dictionary with the input tensors.
@@ -63,7 +64,7 @@ def batch(
         TODO currently only returning loss1, but we could potentially summarize loss1 and loss2 in some way.
         However, note that both loss1 and loss2 are participating in the backward propagation, one after another.
         """
-        output = self(**x)
+        output = self(**x)["output"]
         loss1 = self.compute_loss(output, **y, loss_fn=loss_fn1)
         loss2 = self.compute_loss(output, **y, loss_fn=loss_fn2)
 
diff --git a/tests/test_model/titanic_model.py b/tests/test_model/titanic_model.py
index 65617b27..22a2d21b 100644
--- a/tests/test_model/titanic_model.py
+++ b/tests/test_model/titanic_model.py
@@ -4,6 +4,7 @@
 
 import torch
 from torch import nn
+from torch.optim import Optimizer
 
 
 class ModelTitanic(torch.nn.Module):
@@ -43,7 +44,7 @@ def forward(
         parch: torch.Tensor,
         fare: torch.Tensor,
         embarked: torch.Tensor,
-    ) -> dict:
+    ) -> torch.Tensor:
         """Forward pass of the model.
 
         It should return the output as a dictionary, with the same keys as `y`.
@@ -68,11 +69,11 @@ def compute_loss(self, output: torch.Tensor, survived: torch.Tensor, loss_fn: Ca
 
     def batch(
         self,
-        x: dict,
-        y: dict,
+        x: dict[str, torch.Tensor],
+        y: dict[str, torch.Tensor],
         loss_fn: Callable,
-        optimizer: Optional[Callable] = None,
-    ) -> tuple[torch.Tensor, dict]:
+        optimizer: Optional[Optimizer] = None,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         """Perform one batch step.
 
         `x` is a dictionary with the input tensors.