Remove yaml from classes

src/stimulus/cli/check_model.py

@@ -121,7 +121,7 @@ def main(
     """
     with open(data_config_path) as file:
         data_config = yaml.safe_load(file)
-        data_config = yaml_data.YamlSplitTransformDict(**data_config)
+        data_config = yaml_data.SplitTransformDict(**data_config)
 
     with open(model_config_path) as file:
         model_config = yaml.safe_load(file)

src/stimulus/cli/shuffle_csv.py

@@ -6,7 +6,7 @@
 import yaml
 
 from stimulus.data.data_handlers import DatasetProcessor
-from stimulus.utils.yaml_data import YamlSplitTransformDict
+from stimulus.utils.yaml_data import SplitTransformDict
 
 
 def get_args() -> argparse.Namespace:
@@ -56,7 +56,7 @@ def main(data_csv: str, config_yaml: str, out_path: str) -> None:
     """
     # read the yaml file
     with open(config_yaml) as f:
-        data_config: YamlSplitTransformDict = YamlSplitTransformDict(
+        data_config: SplitTransformDict = SplitTransformDict(
             **yaml.safe_load(f),
         )
     # create a DatasetProcessor object from the config and the csv

src/stimulus/cli/split_csv.py

@@ -7,7 +7,7 @@
 
 from stimulus.data.data_handlers import DatasetProcessor, SplitManager
 from stimulus.data.loaders import SplitLoader
-from stimulus.utils.yaml_data import YamlSplitConfigDict
+from stimulus.utils.yaml_data import SplitConfigDict
 
 
 def get_args() -> argparse.Namespace:
@@ -70,7 +70,7 @@ def main(
     # create a split manager from the config
     split_config = processor.dataset_manager.config.split
     with open(config_yaml) as f:
-        yaml_config = YamlSplitConfigDict(**yaml.safe_load(f))
+        yaml_config = SplitConfigDict(**yaml.safe_load(f))
     split_loader = SplitLoader(seed=yaml_config.global_params.seed)
     split_loader.initialize_splitter_from_config(split_config)
     split_manager = SplitManager(split_loader)

src/stimulus/cli/split_split.py

@@ -12,9 +12,9 @@
 import yaml
 
 from stimulus.utils.yaml_data import (
-    YamlConfigDict,
-    YamlSplitConfigDict,
-    check_yaml_schema,
+    ConfigDict,
+    SplitConfigDict,
+    check_schema,
     dump_yaml_list_into_files,
     generate_split_configs,
 )
@@ -64,13 +64,13 @@ def main(config_yaml: str, out_dir_path: str) -> None:
     with open(config_yaml) as conf_file:
         yaml_config = yaml.safe_load(conf_file)
 
-    yaml_config_dict: YamlConfigDict = YamlConfigDict(**yaml_config)
+    yaml_config_dict: ConfigDict = ConfigDict(**yaml_config)
     # check if the yaml schema is correct
     # FIXME: isn't it redundant to check and already class with pydantic ?
-    check_yaml_schema(yaml_config_dict)
+    check_schema(yaml_config_dict)
 
     # generate the yaml files per split
-    split_configs: list[YamlSplitConfigDict] = generate_split_configs(yaml_config_dict)
+    split_configs: list[SplitConfigDict] = generate_split_configs(yaml_config_dict)
 
     # dump all the YAML configs into files
     dump_yaml_list_into_files(split_configs, out_dir_path, "test_split")

src/stimulus/cli/split_transforms.py

@@ -12,8 +12,8 @@
 import yaml
 
 from stimulus.utils.yaml_data import (
-    YamlSplitConfigDict,
-    YamlSplitTransformDict,
+    SplitConfigDict,
+    SplitTransformDict,
     dump_yaml_list_into_files,
     generate_split_transform_configs,
 )
@@ -60,10 +60,10 @@ def main(config_yaml: str, out_dir_path: str) -> None:
     with open(config_yaml) as conf_file:
         yaml_config = yaml.safe_load(conf_file)
 
-    yaml_config_dict: YamlSplitConfigDict = YamlSplitConfigDict(**yaml_config)
+    yaml_config_dict: SplitConfigDict = SplitConfigDict(**yaml_config)
 
     # Generate the yaml files for each transform
-    split_transform_configs: list[YamlSplitTransformDict] = generate_split_transform_configs(yaml_config_dict)
+    split_transform_configs: list[SplitTransformDict] = generate_split_transform_configs(yaml_config_dict)
 
     # Dump all the YAML configs into files
     dump_yaml_list_into_files(split_transform_configs, out_dir_path, "test_transforms")

src/stimulus/cli/transform_csv.py

@@ -7,7 +7,7 @@
 
 from stimulus.data.data_handlers import DatasetProcessor, TransformManager
 from stimulus.data.loaders import TransformLoader
-from stimulus.utils.yaml_data import YamlSplitConfigDict
+from stimulus.utils.yaml_data import SplitConfigDict
 
 
 def get_args() -> argparse.Namespace:
@@ -55,7 +55,7 @@ def main(data_csv: str, config_yaml: str, out_path: str) -> None:
     # initialize the transform manager
     transform_config = processor.dataset_manager.config.transforms
     with open(config_yaml) as f:
-        yaml_config = YamlSplitConfigDict(**yaml.safe_load(f))
+        yaml_config = SplitConfigDict(**yaml.safe_load(f))
     transform_loader = TransformLoader(seed=yaml_config.global_params.seed)
     transform_loader.initialize_column_data_transformers_from_config(transform_config)
     transform_manager = TransformManager(transform_loader)

src/stimulus/cli/tuning.py

@@ -151,7 +151,7 @@ def get_args() -> argparse.Namespace:
 def main(
     model_path: str,
     data_path: str,
-    data_config: yaml_data.YamlSplitTransformDict,
+    data_config: yaml_data.SplitTransformDict,
     model_config_path: str,
     initial_weights: str | None = None,  # noqa: ARG001
     ray_results_dirpath: str | None = None,
@@ -167,7 +167,7 @@ def main(
     Args:
         data_path: Path to input data file.
         model_path: Path to model file.
-        data_config: A YamlSplitTransformObject
+        data_config: A SplitTransformObject
         model_config_path: Path to model config file.
         initial_weights: Optional path to initial weights.
         ray_results_dirpath: Directory for ray results.
@@ -239,10 +239,10 @@ def run() -> None:
     """Run the model checking script."""
     ray.init(address="auto", ignore_reinit_error=True)
     args = get_args()
-    # Try to convert the configuration file to a YamlSplitTransformDict
-    config_dict: yaml_data.YamlSplitTransformDict
+    # Try to convert the configuration file to a SplitTransformDict
+    config_dict: yaml_data.SplitTransformDict
     with open(args.data_config) as f:
-        config_dict = yaml_data.YamlSplitTransformDict(**yaml.safe_load(f))
+        config_dict = yaml_data.SplitTransformDict(**yaml.safe_load(f))
     main(
         data_path=args.data,
         model_path=args.model,

src/stimulus/data/data_handlers.py

@@ -50,11 +50,11 @@ class DatasetManager:
 
     def __init__(
         self,
-        config_dict: yaml_data.YamlSplitConfigDict,
+        config_dict: yaml_data.SplitConfigDict,
     ) -> None:
         """Initialize the DatasetManager."""
         # self.config = self._load_config(config_path)
-        self.config: yaml_data.YamlSplitTransformDict = config_dict
+        self.config: yaml_data.SplitTransformDict = config_dict
         self.column_categories = self.categorize_columns_by_type()
 
     def categorize_columns_by_type(self) -> dict:
@@ -95,7 +95,7 @@ def categorize_columns_by_type(self) -> dict:
         return {"input": input_columns, "label": label_columns, "meta": meta_columns}
 
     # TODO: Remove or change this function as the config is now preloaded
-    def _load_config(self, config_path: str) -> yaml_data.YamlSplitConfigDict:
+    def _load_config(self, config_path: str) -> yaml_data.SplitConfigDict:
         """Loads and parses a YAML configuration file.
 
         Args:
@@ -112,8 +112,8 @@ def _load_config(self, config_path: str) -> yaml_data.YamlSplitConfigDict:
         """
         with open(config_path) as file:
             # FIXME: cette fonction est appellé pour test_shuffle_csv et test_tune
-            return yaml_data.YamlSplitConfigDict(**yaml.safe_load(file))
-            return yaml_data.YamlSplitTransformDict(**yaml.safe_load(file))
+            return yaml_data.SplitConfigDict(**yaml.safe_load(file))
+            return yaml_data.SplitTransformDict(**yaml.safe_load(file))
 
     def get_split_columns(self) -> list[str]:
         """Get the columns that are used for splitting."""
@@ -282,13 +282,13 @@ class DatasetHandler:
 
     def __init__(
         self,
-        data_config: yaml_data.YamlSplitTransformDict,
+        data_config: yaml_data.SplitTransformDict,
         csv_path: str,
     ) -> None:
         """Initialize the DatasetHandler with required config.
 
         Args:
-            data_config (yaml_data.YamlSplitTransformDict): A YamlSplitTransformDict object holding the config.
+            data_config (yaml_data.SplitTransformDict): A SplitTransformDict object holding the config.
             csv_path (str): Path to the CSV data file.
         """
         self.dataset_manager = DatasetManager(data_config)
@@ -343,7 +343,11 @@ def save(self, path: str) -> None:
 class DatasetProcessor(DatasetHandler):
     """Class for loading dataset, applying transformations and splitting."""
 
-    def __init__(self, data_config: yaml_data.YamlSplitTransformDict, csv_path: str) -> None:
+    def __init__(
+        self,
+        data_config: yaml_data.SplitTransformDict,
+        csv_path: str,
+    ) -> None:
         """Initialize the DatasetProcessor."""
         super().__init__(data_config, csv_path)
 
@@ -416,7 +420,7 @@ class DatasetLoader(DatasetHandler):
 
     def __init__(
         self,
-        data_config: yaml_data.YamlSplitTransformDict,
+        data_config: yaml_data.SplitTransformDict,
         csv_path: str,
         encoder_loader: loaders.EncoderLoader,
         split: Union[int, None] = None,

src/stimulus/data/handlertorch.py

@@ -5,23 +5,23 @@
 from torch.utils.data import Dataset
 
 from stimulus.data import data_handlers, loaders
-from stimulus.utils.yaml_data import YamlSplitTransformDict
+from stimulus.utils.yaml_data import SplitTransformDict
 
 
 class TorchDataset(Dataset):
     """Class for creating a torch dataset."""
 
     def __init__(
         self,
-        data_config: YamlSplitTransformDict,
+        data_config: SplitTransformDict,
         csv_path: str,
         encoder_loader: loaders.EncoderLoader,
         split: Optional[int] = None,
     ) -> None:
         """Initialize the TorchDataset.
 
         Args:
-            data_config: A YamlSplitTransformDict holding the configuration.
+            data_config: A SplitTransformDict holding the configuration.
             csv_path: Path to the CSV data file
             encoder_loader: Encoder loader instance
             split: Optional tuple containing split information

src/stimulus/data/loaders.py

@@ -31,12 +31,12 @@ def __init__(self, seed: Optional[float] = None) -> None:
 
     def initialize_column_encoders_from_config(
         self,
-        column_config: yaml_data.YamlColumns,
+        column_config: yaml_data.Columns,
     ) -> None:
         """Build the loader from a config dictionary.
 
         Args:
-            column_config (yaml_data.YamlColumns): Configuration dictionary containing field names (column_name) and their encoder specifications.
+            column_config (yaml_data.Columns): 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)
@@ -176,12 +176,12 @@ def set_data_transformer_as_attribute(
 
     def initialize_column_data_transformers_from_config(
         self,
-        transform_config: yaml_data.YamlTransform,
+        transform_config: yaml_data.Transform,
     ) -> None:
         """Build the loader from a config dictionary.
 
         Args:
-            transform_config (yaml_data.YamlTransform): Configuration dictionary containing transforms configurations.
+            transform_config (yaml_data.Transform): Configuration dictionary containing transforms configurations.
 
         Example:
             Given a YAML config like:
@@ -282,12 +282,12 @@ def set_splitter_as_attribute(self, splitter: Any) -> None:
 
     def initialize_splitter_from_config(
         self,
-        split_config: yaml_data.YamlSplitConfigDict,
+        split_config: yaml_data.SplitConfigDict,
     ) -> None:
         """Build the loader from a config dictionary.
 
         Args:
-            split_config (yaml_data.YamlSplitConfigDict): Configuration dictionary containing split configurations.
+            split_config (yaml_data.SplitConfigDict): Configuration dictionary containing split configurations.
         """
         splitter = self.get_splitter(split_config.split_method, split_config.params)
         self.set_splitter_as_attribute(splitter)

src/stimulus/learner/raytune_learner.py

@@ -20,7 +20,7 @@
 from stimulus.data.loaders import EncoderLoader
 from stimulus.learner.predict import PredictWrapper
 from stimulus.utils.generic_utils import set_general_seeds
-from stimulus.utils.yaml_data import YamlSplitTransformDict
+from stimulus.utils.yaml_data import SplitTransformDict
 from stimulus.utils.yaml_model_schema import RayTuneModel
 
 
@@ -36,7 +36,7 @@ class TuneWrapper:
     def __init__(
         self,
         model_config: RayTuneModel,
-        data_config: YamlSplitTransformDict,
+        data_config: SplitTransformDict,
         model_class: nn.Module,
         data_path: str,
         encoder_loader: EncoderLoader,
@@ -119,7 +119,7 @@ def __init__(
 
     def tuner_initialization(
         self,
-        data_config: YamlSplitTransformDict,
+        data_config: SplitTransformDict,
         data_path: str,
         encoder_loader: EncoderLoader,
         *,
@@ -318,7 +318,8 @@ def objective(self) -> dict[str, float]:
             **{"train_" + metric: value for metric, value in predict_train.compute_metrics(metrics).items()},
         }
 
-    def export_model(self, export_dir: str | None = None) -> None:  # type: ignore[override]
+    # type: ignore[override]
+    def export_model(self, export_dir: str | None = None) -> None:
         """Export model to safetensors format."""
         if export_dir is None:
             return
@@ -329,11 +330,19 @@ def load_checkpoint(self, checkpoint: dict[Any, Any] | None) -> None:
         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")))
+        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[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"))
+        torch.save(
+            self.optimizer.state_dict(),
+            os.path.join(checkpoint_dir, "optimizer.pt"),
+        )
         return {"checkpoint_dir": checkpoint_dir}