sinapsis-data-tools 0.2.26__tar.gz → 0.2.28__tar.gz

{sinapsis_data_tools-0.2.26 → sinapsis_data_tools-0.2.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sinapsis-data-tools
-Version: 0.2.26
+Version: 0.2.28
 Summary: Module that contains different packages to perform data-related operations through Sinapsis templates.
 Author-email: SinapsisAI <dev@sinapsis.tech>
 Project-URL: Homepage, https://sinapsis.tech

{sinapsis_data_tools-0.2.26 → sinapsis_data_tools-0.2.28}/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/__init__.py

@@ -7,13 +7,15 @@ from sinapsis_core.template_base import Template
 _root_lib_path: str = "sinapsis_data_analysis.templates"
 
 _ADDITIONAL_TEMPLATE_MODULES = [
-    f"{_root_lib_path}.sklearn_manifold",
-    f"{_root_lib_path}.sklearn_train",
-    f"{_root_lib_path}.xgboost_train",
+    f"{_root_lib_path}.training.sklearn_manifold",
+    f"{_root_lib_path}.training.sklearn_train",
+    f"{_root_lib_path}.training.xgboost_train",
+
 ]
 _template_lookup: dict = {
-    "SKLearnInference": f"{_root_lib_path}.sklearn_inference",
-    "XGBoostInference": f"{_root_lib_path}.xgboost_inference",
+    "SKLearnInference": f"{_root_lib_path}.inference.sklearn_inference",
+    "XGBoostInference": f"{_root_lib_path}.inference.xgboost_inference",
+    "CategoricalToNumerical": f"{_root_lib_path}.transformation.categorical_to_numerical",
 }
 for t_module in _ADDITIONAL_TEMPLATE_MODULES:
     _template_lookup |= _import_template_package(t_module)

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/inference}/ml_base_inference.py

@@ -7,7 +7,7 @@ import numpy as np
 from sinapsis_core.data_containers.data_packet import DataContainer
 from sinapsis_core.template_base.base_models import TemplateAttributes
 from sinapsis_core.template_base.template import Template
-from sinapsis_core.utils.env_var_keys import SINAPSIS_CACHE_DIR
+from sinapsis_core.utils.env_var_keys import WORKING_DIR
 
 
 class MLBaseInference(Template):
@@ -26,8 +26,9 @@ class MLBaseInference(Template):
         """
 
         model_path: str
-        root_dir : str = SINAPSIS_CACHE_DIR
+        root_dir : str = WORKING_DIR
         generic_field_key: str
+        target_key : str
 
     def __init__(self, attributes: TemplateAttributes) -> None:
         super().__init__(attributes)
@@ -42,7 +43,7 @@ class MLBaseInference(Template):
         Returns:
             Any: The data from the generic field
         """
-        return self._get_generic_data(container, self.attributes.generic_field_key)
+        return container.data_frames[-1].content
 
     @staticmethod
     def data_is_valid(data: Any) -> bool:
@@ -68,7 +69,7 @@ class MLBaseInference(Template):
             Any: The preprocessed data
         """
         try:
-            data.pop("target")
+            data.pop(self.attributes.target_key)
         except (KeyError, IndexError):
             self.logger.info("No target column")
         return data
@@ -95,7 +96,6 @@ class MLBaseInference(Template):
         Returns:
             np.ndarray: The model's predictions
         """
-
         return self.model.predict(data)
 
     def execute(self, container: DataContainer) -> DataContainer:

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/inference}/sklearn_inference.py

@@ -3,9 +3,8 @@ from typing import Any
 
 import joblib
 from sinapsis_core.template_base.base_models import UIPropertiesMetadata
-
 from sinapsis_data_analysis.helpers.tags import Tags
-from sinapsis_data_analysis.templates.ml_base_inference import MLBaseInference
+from sinapsis_data_analysis.templates.inference.ml_base_inference import MLBaseInference
 
 
 class SKLearnInference(MLBaseInference):

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/inference}/xgboost_inference.py

@@ -2,9 +2,8 @@
 
 
 from sinapsis_core.template_base.base_models import UIPropertiesMetadata
-
 from sinapsis_data_analysis.helpers.tags import Tags
-from sinapsis_data_analysis.templates.sklearn_inference import SKLearnInference
+from sinapsis_data_analysis.templates.inference.sklearn_inference import SKLearnInference
 
 
 class XGBoostInference(SKLearnInference):

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/training}/ml_base_training.py

@@ -4,11 +4,14 @@ from abc import abstractmethod
 from typing import Any
 
 import numpy as np
-from sinapsis_core.data_containers.data_packet import DataContainer
+from sinapsis_core.data_containers.data_packet import DataContainer, DataFramePacket
 from sinapsis_core.template_base.base_models import TemplateAttributes
 from sinapsis_core.template_base.dynamic_template import BaseDynamicWrapperTemplate
 from sinapsis_core.utils.env_var_keys import WORKING_DIR
-from sinapsis_data_readers.templates.datasets_readers.dataset_splitter import TabularDatasetSplit
+from sinapsis_data_analysis.helpers.model_metrics import (
+    ModelMetrics,
+    ModelPredictionResults,
+)
 from sklearn.base import is_classifier, is_regressor
 from sklearn.metrics import (
     accuracy_score,
@@ -20,21 +23,14 @@ from sklearn.metrics import (
     recall_score,
 )
 
-from sinapsis_data_analysis.helpers.model_metrics import (
-    ModelMetrics,
-    ModelPredictionResults,
-)
-
 
 class MLBaseAttributes(TemplateAttributes):
     """Base attributes for machine learning model templates.
 
     Attributes:
-        generic_field_key (str): Key of the generic field where datasets are stored.
         model_save_path (str): Path where the trained model will be saved.
     """
 
-    generic_field_key: str | None = None
     root_dir : str = WORKING_DIR
     model_save_path: str
 
@@ -69,7 +65,9 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
         Returns:
             Any: The dataset from the generic field.
         """
-        return self._get_generic_data(container, self.attributes.generic_field_key)
+
+
+        return container.data_frames
 
     @staticmethod
     def dataset_is_valid(dataset: Any) -> bool:
@@ -83,7 +81,7 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
         """
         return dataset is not None
 
-    def process_dataset(self, dataset: TabularDatasetSplit | dict) -> tuple | None:
+    def process_dataset(self, dataset: list[DataFramePacket]) -> tuple | None:
         """
         Extracts x_train, y_train, x_test, y_test from the dataset
 
@@ -94,18 +92,23 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
             tuple | None: A tuple containing (x_train, y_train, x_test, y_test)
                 or None if the dataset doesn't have the expected attributes
         """
-        if isinstance(dataset, dict):
-            dataset = TabularDatasetSplit(**dataset)
-        try:
-            x_train = dataset.x_train
-            y_train = dataset.y_train
-            x_test = dataset.x_test
-            y_test = dataset.y_test
-
-            return x_train, y_train, x_test, y_test
-        except AttributeError:
-            self.logger.warning("Dataset doesn't have the expected attributes")
-            return None
+        x_train, y_train, x_test, y_test = None, None, None, None
+        for set in dataset:
+            if "x_train" in set.source:
+                x_train = set.content
+            elif "x_test" in set.source:
+                x_test = set.content
+            elif "y_train" in set.source:
+                y_train = set.content
+            elif "y_test" in set.source:
+                y_test = set.content
+            else:
+                if "x_dataset" in set.source:
+                    x_train =  set.content
+                elif "y_dataset" in set.source:
+                    y_train = set.content
+
+        return x_train, y_train, x_test, y_test
 
     def train_model(self, x_train: Any, y_train: Any) -> None:
         """Train the model using the training data
@@ -114,6 +117,7 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
             x_train (Any): The training features
             y_train (Any): The training targets
         """
+
         self.trained_model = self.model.fit(x_train, y_train)
 
     @staticmethod
@@ -152,7 +156,7 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
 
         return metrics
 
-    def calculate_metrics(self, y_true: np.ndarray, y_pred: np.ndarray) -> ModelMetrics:
+    def calculate_metrics(self, y_true: np.ndarray, y_pred: np.ndarray) -> ModelMetrics | None:
         """
         Detects whether the model is a classifier or regressor and calculates
         the appropriate metrics
@@ -169,7 +173,7 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
                 return self.calculate_classification_metrics(y_true, y_pred)
             elif is_regressor(self.trained_model):
                 return self.calculate_regression_metrics(y_true, y_pred)
-        return ModelMetrics()
+        return None
 
     def generate_predictions(self, x_test: np.ndarray, y_test: np.ndarray) -> ModelPredictionResults | None:
         """
@@ -183,7 +187,7 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
         Returns:
             ModelPredictionResults: Object containing predictions and metrics
         """
-        if self.trained_model is not None:
+        if self.trained_model is not None and x_test is not None:
             predictions = self.trained_model.predict(x_test)
 
             metrics = self.calculate_metrics(y_test, predictions)
@@ -204,8 +208,8 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
             ModelPredictionResults: Object containing predictions and metrics.
         """
         x_train, y_train, x_test, y_test = processed_data
-
         self.train_model(x_train, y_train)
+
         return self.generate_predictions(x_test, y_test)
 
     def save_model(self) -> None:
@@ -221,7 +225,7 @@ class MLBaseTraining(BaseDynamicWrapperTemplate):
         try:
             os.makedirs(os.path.dirname(full_path), exist_ok=True)
             self._save_model_implementation(full_path)
-            self.logger.info(f"Model saved at {self.attributes.model_save_path}")
+            self.logger.info(f"Model saved at {full_path}")
         except (MemoryError, TypeError) as e:
             self.logger.error(f"Error saving model: {e}")
 

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/training}/sklearn_manifold.py

@@ -1,5 +1,4 @@
 # -*- coding: utf-8 -*-
-from typing import cast
 
 import numpy as np
 import pandas as pd
@@ -13,12 +12,8 @@ from sinapsis_core.template_base.dynamic_template import (
 )
 from sinapsis_core.template_base.dynamic_template_factory import make_dynamic_template
 from sinapsis_core.utils.env_var_keys import SINAPSIS_BUILD_DOCS
-from sinapsis_data_readers.templates.datasets_readers.dataset_splitter import (
-    TabularDatasetSplit,
-)
-from sklearn import manifold
-
 from sinapsis_data_analysis.helpers.tags import Tags
+from sklearn import manifold
 
 
 class ManifoldResults(BaseModel):
@@ -42,6 +37,14 @@ class SKLearnManifold(BaseDynamicWrapperTemplate):
     TSNE, MDS, Isomap, etc.
     """
 
+    class AttributesBaseModel(TemplateAttributes):
+        """Attributes for the MLBaseInference template.
+
+        Attributes:
+            target_key (str): Key of the generic field where data is stored.
+        """
+
+        target_key : str =  'target'
     WrapperEntry = WrapperEntryConfig(
         wrapped_object=manifold,
         signature_from_doc_string=True,
@@ -54,15 +57,6 @@ class SKLearnManifold(BaseDynamicWrapperTemplate):
         tags=[Tags.DATA_ANALYSIS, Tags.DYNAMIC, Tags.MANIFOLD, Tags.SKLEARN, Tags.MODELS],
     )
 
-    class AttributesBaseModel(TemplateAttributes):
-        """Attributes for the SKLearnManifold template.
-
-        Attributes:
-            generic_field_key (str): Key of the generic field
-                where the input data is stored.
-        """
-
-        generic_field_key: str
 
     def __init__(self, attributes: TemplateAttributes) -> None:
         super().__init__(attributes)
@@ -83,7 +77,7 @@ class SKLearnManifold(BaseDynamicWrapperTemplate):
         array_data = np.array(feature_arrays)
         return array_data.reshape(array_data.shape[0], -1)
 
-    def get_dataset(self, container: DataContainer) -> TabularDatasetSplit | None:
+    def get_dataset(self, container: DataContainer) -> list:
         """Get the dataset from the data container
 
         Args:
@@ -93,13 +87,10 @@ class SKLearnManifold(BaseDynamicWrapperTemplate):
             TabularDatasetSplit | None: The dataset from the generic field,
                 or None if not found
         """
-        dataset = self._get_generic_data(container, self.attributes.generic_field_key)
-        dataset = cast(TabularDatasetSplit, dataset)
-        if dataset:
-            return dataset
-        return None
+        return container.data_frames
 
-    def process_dataset(self, dataset: TabularDatasetSplit) -> ManifoldResults | None:
+
+    def process_dataset(self, dataset: list) -> ManifoldResults | None:
         """
         Extracts the training data, reshapes it, and applies the
         manifold learning transformation
@@ -111,11 +102,17 @@ class SKLearnManifold(BaseDynamicWrapperTemplate):
             ManifoldResults | None: Results of the manifold transformation,
                 or None if the dataset is empty
         """
-        x_train = dataset.x_train
-        y_train = dataset.y_train
+        x_train, y_train = None, None
+        for set in dataset:
+            if "x_train" in set.source:
+                x_train = set.content
+
+            elif "y_train" in set.source:
+                y_train = set.content
+            else:
+                y_train = set.content[self.attributes.target_key]
+                x_train = set.content.pop(self.attributes.target_key)
 
-        if x_train is None or x_train.empty:
-            return None
 
         x_train_reshaped = self.reshape_arrays(x_train)
         x_transformed = self.manifold_model.fit_transform(x_train_reshaped)
@@ -141,7 +138,7 @@ class SKLearnManifold(BaseDynamicWrapperTemplate):
         results = self.process_dataset(dataset)
 
         if results is not None:
-            self._set_generic_data(container, results)
+            self._set_generic_data(container, results.model_dump())
 
         return container
 

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/training}/sklearn_train.py

@@ -5,8 +5,6 @@ from sinapsis_core.template_base.base_models import UIPropertiesMetadata
 from sinapsis_core.template_base.dynamic_template import WrapperEntryConfig
 from sinapsis_core.template_base.dynamic_template_factory import make_dynamic_template
 from sinapsis_core.utils.env_var_keys import SINAPSIS_BUILD_DOCS
-from sklearn import linear_model, neighbors, neural_network, svm, tree
-
 from sinapsis_data_analysis.helpers.excluded_models import (
     excluded_linear_models,
     excluded_neighbors_models,
@@ -14,7 +12,8 @@ from sinapsis_data_analysis.helpers.excluded_models import (
     excluded_tree_models,
 )
 from sinapsis_data_analysis.helpers.tags import Tags
-from sinapsis_data_analysis.templates.ml_base_training import MLBaseTraining
+from sinapsis_data_analysis.templates.training.ml_base_training import MLBaseTraining
+from sklearn import linear_model, neighbors, neural_network, svm, tree
 
 
 class SKLearnLinearModelsTrain(MLBaseTraining):

{sinapsis_data_tools-0.2.26/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates → sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/training}/xgboost_train.py

@@ -5,9 +5,8 @@ from sinapsis_core.template_base.base_models import UIPropertiesMetadata
 from sinapsis_core.template_base.dynamic_template import WrapperEntryConfig
 from sinapsis_core.template_base.dynamic_template_factory import make_dynamic_template
 from sinapsis_core.utils.env_var_keys import SINAPSIS_BUILD_DOCS
-
 from sinapsis_data_analysis.helpers.tags import Tags
-from sinapsis_data_analysis.templates.sklearn_train import SKLearnLinearModelsTrain
+from sinapsis_data_analysis.templates.training.sklearn_train import SKLearnLinearModelsTrain
 
 INCLUDED_MODELS = [
     "XGBClassifier",

sinapsis_data_tools-0.2.28/packages/sinapsis_data_analysis/src/sinapsis_data_analysis/templates/transformation/categorical_to_numerical.py

@@ -0,0 +1,45 @@
+import pandas as pd
+from sinapsis_core.data_containers.data_packet import DataContainer, DataFramePacket
+from sinapsis_core.template_base.base_models import TemplateAttributes
+from sinapsis_core.template_base.template import Template
+
+
+class CategoricalToNumerical(Template):
+	class AttributesBaseModel(TemplateAttributes):
+		generic_key: str
+
+	@staticmethod
+	def map_categorical_to_numerical(df: pd.DataFrame| pd.Series)->tuple:
+
+		if isinstance(df, pd.Series):
+			categorical_cols =  [df.name] if df.dtype in ["object", "string", "category"] else []
+		else:
+			categorical_cols =  df.select_dtypes(include=["object", "string", "category"]).columns.tolist()
+
+
+		category_maps = {}
+
+		for col in categorical_cols:
+			categories = df[col].astype("string").unique()
+			category_maps[col] = {cat: idx for idx, cat in enumerate(categories)}
+		for col, mapping in category_maps.items():
+			df[col] = df[col].map(mapping)
+		return df, category_maps
+	@staticmethod
+	def unmap_numerical_to_categorical(df: pd.DataFrame| pd.Series, categories: dict):
+		inverse_maps = {col: {v: k for k, v in mapping.items()} for col, mapping in categories.items()}
+
+		for col in inverse_maps:
+			df[col] = df[col].map(inverse_maps[col])
+		return df
+
+
+	def execute(self, container: DataContainer) -> DataContainer:
+		for data_frame in container.data_frames:
+
+			transformed_dataset, labels = self.map_categorical_to_numerical(data_frame.content)
+			container.data_frames.append(DataFramePacket(content=transformed_dataset, generic_data = labels))
+
+		return container
+
+

{sinapsis_data_tools-0.2.26 → sinapsis_data_tools-0.2.28}/packages/sinapsis_data_readers/src/sinapsis_data_readers/templates/datasets_readers/csv_datasets.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import os
 
-from sinapsis_core.data_containers.data_packet import DataContainer, TimeSeriesPacket
+from sinapsis_core.data_containers.data_packet import DataContainer, DataFramePacket, TimeSeriesPacket
 from sinapsis_core.template_base.base_models import TemplateAttributes, TemplateAttributeType
 from sinapsis_core.template_base.template import Template
 from sinapsis_core.utils.env_var_keys import SINAPSIS_CACHE_DIR
@@ -25,6 +25,7 @@ class CSVDatasetReader(Template):
             packet = TimeSeriesPacket(content=self.csv_file)
             container.time_series.append(packet)
         else:
-            self._set_generic_data(container, self.csv_file)
+            packet = DataFramePacket(content=self.csv_file)
+            container.data_frames.append(packet)
 
         return container

sinapsis_data_tools-0.2.28/packages/sinapsis_data_readers/src/sinapsis_data_readers/templates/datasets_readers/dataset_splitter.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+
+import pandas as pd
+from sinapsis_core.data_containers.data_packet import DataContainer, DataFramePacket
+from sinapsis_core.template_base import Template
+from sinapsis_core.template_base.base_models import TemplateAttributes
+from sklearn.model_selection import train_test_split
+
+
+class CSVDatasetSplitter(Template):
+    """
+    Template to split a tabular data set into test and train samples.
+    The template retrieves the dataset from the dataframe packet of the
+    container and stores the features and targets as new dataframes, with source
+    indicating whether train or test samples
+    """
+
+    class AttributesBaseModel(TemplateAttributes):
+        target_key: str = "target"  # labels
+        feature_key: str | None  = None  # arrays
+        random_state: int = 42
+        train_size: float = 0.2
+
+
+
+
+    def extract_x_y_from_packet(self, data_frame: DataFramePacket) -> tuple[pd.DataFrame, pd.DataFrame]:
+
+
+        target = data_frame.content.get(self.attributes.target_key)
+        feature = data_frame.content.get(self.attributes.feature_key) \
+            if self.attributes.feature_key \
+            else data_frame.drop(columns=[self.attributes.target_key])
+        return feature, target
+
+    def split_dataset(self, x_data:pd.DataFrame, y_data:pd.DataFrame) -> dict:
+
+        x_train, x_test, y_train, y_test = train_test_split(
+            x_data,
+            y_data,
+            train_size=self.attributes.train_size,
+            test_size=1 - self.attributes.train_size,
+            random_state=self.attributes.random_state,
+        )
+        data_map = {
+            "x_train": x_train,
+            "y_train": y_train,
+            "x_test": x_test,
+            "y_test": y_test
+        }
+        return data_map
+    def assign_to_dataframe_packets(self,container: DataContainer, source: str, data: pd.DataFrame) -> DataContainer:
+        df_packet = DataFramePacket(content=data, source=f"{self.instance_name}_{source}")
+        container.data_frames.append(df_packet)
+        return container
+
+    def execute(self, container: DataContainer) -> DataContainer:
+        for df in container.data_frames:
+            feature, target = self.extract_x_y_from_packet(df)
+            sample = feature, target, None, None
+            if self.attributes.train_size:
+                sample = self.split_dataset(feature, target)
+            for name, data_frame in sample.items():
+                container = self.assign_to_dataframe_packets(container, name, data_frame)
+        return container
+

{sinapsis_data_tools-0.2.26 → sinapsis_data_tools-0.2.28}/packages/sinapsis_data_readers/src/sinapsis_data_readers/templates/datasets_readers/sklearn_datasets.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import pandas as pd
-from sinapsis_core.data_containers.data_packet import DataContainer, TimeSeriesPacket
+from sinapsis_core.data_containers.data_packet import DataContainer, DataFramePacket, TimeSeriesPacket
 from sinapsis_core.template_base import Template
 from sinapsis_core.template_base.base_models import TemplateAttributes, TemplateAttributeType, UIPropertiesMetadata
 from sinapsis_core.template_base.dynamic_template import (
@@ -19,9 +19,6 @@ from sklearn.utils import Bunch
 
 from sinapsis_data_readers.helpers import sklearn_dataset_subset
 from sinapsis_data_readers.helpers.tags import Tags
-from sinapsis_data_readers.templates.datasets_readers.dataset_splitter import (
-    TabularDatasetSplit,
-)
 
 TARGET: str = "target"
 
@@ -137,39 +134,40 @@ class SKLearnDatasets(BaseDynamicWrapperTemplate):
         results: pd.DataFrame, feature_name_cols: list, target_name_cols: list, n_features: int, split_size: float
     ) -> dict:
         """Method to split the dataset into training and testing samples"""
-        if feature_name_cols:
+        if feature_name_cols is not None:
             X = results[feature_name_cols]
             y = results[target_name_cols]
         else:
             X = results.iloc[:, :n_features]
             y = results.iloc[:, n_features:]
 
-
         x_train, x_test, y_train, y_test = train_test_split(X, y, train_size=split_size, random_state=0)
-        split_data = TabularDatasetSplit(
-            x_train=pd.DataFrame(x_train),
-            x_test=pd.DataFrame(x_test),
-            y_train=pd.DataFrame(y_train),
-            y_test=pd.DataFrame(y_test),
-        )
-
-        return split_data.model_dump()
+        data_map = {
+            "x_train": x_train,
+            "y_train": y_train,
+            "x_test": x_test,
+            "y_test": y_test
+        }
+        return data_map
 
     def execute(self, container: DataContainer) -> DataContainer:
         sklearn_dataset = self.wrapped_callable.__func__(**self.dataset_attributes.model_dump())
         dataset, feature_columns, target_columns, n_features = self.parse_results(sklearn_dataset)
+
         if self.attributes.store_as_time_series:
             time_series_packet = TimeSeriesPacket(content=dataset)
             container.time_series.append(time_series_packet)
 
         if self.attributes.split_dataset:
-            split_dataset = self.split_dataset(
+            data_map = self.split_dataset(
                 dataset, feature_columns, target_columns, n_features, split_size=self.attributes.train_size
             )
-            self._set_generic_data(container, split_dataset)
-        if sklearn_dataset and not self.attributes.split_dataset:
-            self._set_generic_data(container, dataset)
 
+            for name, df  in data_map.items():
+                container.data_frames.append(DataFramePacket(content=df, source=f"{self.instance_name}_{name}"))
+
+        if sklearn_dataset and not self.attributes.split_dataset:
+            container.data_frames.append(DataFramePacket(content=dataset, source=self.instance_name))
         return container