Repository: DashAISoftware/DashAI (language: python)

Issue:
Change splits in Experiment configuration
# Summary

<!--
Include a short summary of the changes made to the platforn and list any dependencies
of the pr - Required.
-->
Adds back-end logic to support the change of splits configuration in Prepare Dataset Step.

## Type of change

<!-- Indicate the type of change. Delete those that do not apply  - Required -->

- Front end new feature.
- Back end new feature.
- Refactoring.

## Changes

<!--
Indicate the changes/fixes that you want to merge - Required.

- Bug 1
- Bug 2
- Feature 1
- Feature 2
-->
- Shows default splits in Prepare Dataset Step
- Save split configuration adding status to know is the splits have changed and if the new configuration is radom splits
- Use this new configuration in ModelJob flow by calling a new function update_dataset_splits
- Include tests for update_dataset_splits

## How to Test

<!--
Describe the tests or executions that you ran to verify your changes and provide
instructions to reproduce them - Required.
-->
1.- To check test run ```pytest tests/```
2.- To check correct functionality of this new feature from the front:
- Run ```python -m DashAI```
- Go to ```dashai/front``` and run ```yarn start```
- Create new experiments with different configuration of splits and verify the diffence between the results.



Discussion / hints:

Fix patch:
diff --git a/DashAI/back/api/api_v1/endpoints/datasets.py b/DashAI/back/api/api_v1/endpoints/datasets.py
index 0c97ff55..5ca68a99 100644
--- a/DashAI/back/api/api_v1/endpoints/datasets.py
+++ b/DashAI/back/api/api_v1/endpoints/datasets.py
@@ -1,8 +1,7 @@
-import json
 import logging
 import os
 import shutil
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict
 
 from dependency_injector.wiring import Provide, inject
 from fastapi import APIRouter, Depends, File, Form, Response, UploadFile, status
@@ -11,11 +10,23 @@
 from sqlalchemy.orm import Session
 from typing_extensions import ContextManager
 
-from DashAI.back.api.api_v1.schemas.datasets_params import DatasetParams
+from DashAI.back.api.api_v1.schemas.datasets_params import (
+    DatasetParams,
+    DatasetUpdateParams,
+)
 from DashAI.back.api.utils import parse_params
 from DashAI.back.containers import Container
-from DashAI.back.dataloaders.classes.dashai_dataset import save_dataset
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    DashAIDataset,
+    get_columns_spec,
+    get_dataset_info,
+    load_dataset,
+    save_dataset,
+    split_dataset,
+    split_indexes,
+    to_dashai_dataset,
+    update_columns_spec,
+)
 from DashAI.back.dependencies.database.models import Dataset
 from DashAI.back.dependencies.registry import ComponentRegistry
 
@@ -104,6 +115,126 @@ async def get_dataset(
     return dataset
 
 
+@router.get("/{dataset_id}/sample")
+@inject
+async def get_sample(
+    dataset_id: int,
+    session_factory: Callable[..., ContextManager[Session]] = Depends(
+        Provide[Container.db.provided.session]
+    ),
+):
+    """Return the dataset with id dataset_id from the database.
+
+    Parameters
+    ----------
+    dataset_id : int
+        id of the dataset to query.
+
+    Returns
+    -------
+    Dict
+        A Dict with a sample of 10 rows
+    """
+    with session_factory() as db:
+        try:
+            file_path = db.get(Dataset, dataset_id).file_path
+            if not file_path:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Dataset not found",
+                )
+            dataset: DashAIDataset = load_dataset(f"{file_path}/dataset")
+            sample = dataset["train"].sample(n=10)
+        except exc.SQLAlchemyError as e:
+            logger.exception(e)
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal database error",
+            ) from e
+    return sample
+
+
+@router.get("/{dataset_id}/info")
+@inject
+async def get_info(
+    dataset_id: int,
+    session_factory: Callable[..., ContextManager[Session]] = Depends(
+        Provide[Container.db.provided.session]
+    ),
+):
+    """Return the dataset with id dataset_id from the database.
+
+    Parameters
+    ----------
+    dataset_id : int
+        id of the dataset to query.
+
+    Returns
+    -------
+    JSON
+        JSON with the specified dataset id.
+    """
+    with session_factory() as db:
+        try:
+            dataset = db.get(Dataset, dataset_id)
+            if not dataset:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Dataset not found",
+                )
+            info = get_dataset_info(f"{dataset.file_path}/dataset")
+        except exc.SQLAlchemyError as e:
+            logger.exception(e)
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal database error",
+            ) from e
+    return info
+
+
+@router.get("/{dataset_id}/types")
+@inject
+async def get_types(
+    dataset_id: int,
+    session_factory: Callable[..., ContextManager[Session]] = Depends(
+        Provide[Container.db.provided.session]
+    ),
+):
+    """Return the dataset with id dataset_id from the database.
+
+    Parameters
+    ----------
+    dataset_id : int
+        id of the dataset to query.
+
+    Returns
+    -------
+    Dict
+        Dict containing column names and types.
+    """
+    with session_factory() as db:
+        try:
+            file_path = db.get(Dataset, dataset_id).file_path
+            if not file_path:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Dataset not found",
+                )
+            columns_spec = get_columns_spec(f"{file_path}/dataset")
+            if not columns_spec:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Error while loading column types.",
+                )
+        except exc.SQLAlchemyError as e:
+            logger.exception(e)
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal database error",
+            ) from e
+    return columns_spec
+
+
 @router.post("/", status_code=status.HTTP_201_CREATED)
 @inject
 async def upload_dataset(
@@ -169,30 +300,24 @@ async def upload_dataset(
             temp_path=str(folder_path),
             params=parsed_params.dataloader_params.dict(),
         )
-        columns = dataset["train"].column_names
-        outputs_columns = parsed_params.outputs_columns
-
-        if len(outputs_columns) == 0:
-            outputs_columns = [s for s in columns if s in ["class", "label"]]
-            if not outputs_columns:
-                outputs_columns = [columns[-1]]
 
-        inputs_columns = [x for x in columns if x not in outputs_columns]
-
-        dataset = to_dashai_dataset(dataset, inputs_columns, outputs_columns)
+        dataset = to_dashai_dataset(dataset)
 
         if not parsed_params.splits_in_folders:
-            dataset = dataloader.split_dataset(
-                dataset,
+            n = len(dataset["train"])
+            train_indexes, test_indexes, val_indexes = split_indexes(
+                n,
                 parsed_params.splits.train_size,
                 parsed_params.splits.test_size,
                 parsed_params.splits.val_size,
                 parsed_params.splits.seed,
                 parsed_params.splits.shuffle,
-                parsed_params.splits.stratify,
-                outputs_columns[0],  # Stratify according
-                # to the split is only done in classification,
-                # so it will correspond to the class column.
+            )
+            dataset = split_dataset(
+                dataset["train"],
+                train_indexes=train_indexes,
+                test_indexes=test_indexes,
+                val_indexes=val_indexes,
             )
             dataset_path = folder_path / "dataset"
         logger.debug("Saving dataset in %s", str(dataset_path))
@@ -219,8 +344,6 @@ async def upload_dataset(
             folder_path = os.path.realpath(folder_path)
             dataset = Dataset(
                 name=parsed_params.dataset_name,
-                task_name=parsed_params.task_name,
-                feature_names=json.dumps(inputs_columns),
                 file_path=folder_path,
             )
             db.add(dataset)
@@ -296,8 +419,7 @@ async def delete_dataset(
 @inject
 async def update_dataset(
     dataset_id: int,
-    name: Union[str, None] = None,
-    task_name: Union[str, None] = None,
+    params: DatasetUpdateParams,
     session_factory: Callable[..., ContextManager[Session]] = Depends(
         Provide[Container.db.provided.session]
     ),
@@ -324,11 +446,10 @@ async def update_dataset(
     with session_factory() as db:
         try:
             dataset = db.get(Dataset, dataset_id)
-            if name:
-                setattr(dataset, "name", name)
-            if task_name:
-                setattr(dataset, "task_name", task_name)
-            if name or task_name:
+            if params.columns:
+                update_columns_spec(f"{dataset.file_path}/dataset", params.columns)
+            if params.name:
+                setattr(dataset, "name", params.name)
                 db.commit()
                 db.refresh(dataset)
                 return dataset
diff --git a/DashAI/back/api/api_v1/endpoints/experiments.py b/DashAI/back/api/api_v1/endpoints/experiments.py
index 037f95a7..3d4981c2 100644
--- a/DashAI/back/api/api_v1/endpoints/experiments.py
+++ b/DashAI/back/api/api_v1/endpoints/experiments.py
@@ -7,9 +7,18 @@
 from sqlalchemy import exc
 from sqlalchemy.orm import Session
 
-from DashAI.back.api.api_v1.schemas.experiments_params import ExperimentParams
+from DashAI.back.api.api_v1.schemas.experiments_params import (
+    ColumnsValidationParams,
+    ExperimentParams,
+)
 from DashAI.back.containers import Container
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    load_dataset,
+    parse_columns_indices,
+)
 from DashAI.back.dependencies.database.models import Dataset, Experiment
+from DashAI.back.dependencies.registry import ComponentRegistry
+from DashAI.back.tasks.base_task import BaseTask
 
 logging.basicConfig(level=logging.DEBUG)
 log = logging.getLogger(__name__)
@@ -90,6 +99,68 @@ async def get_experiment(
         return experiment
 
 
+@router.post("/validation")
+@inject
+async def validate_columns(
+    params: ColumnsValidationParams,
+    component_registry: ComponentRegistry = Depends(
+        Provide[Container.component_registry]
+    ),
+    session_factory: Callable[..., ContextManager[Session]] = Depends(
+        Provide[Container.db.provided.session]
+    ),
+):
+    with session_factory() as db:
+        try:
+            dataset = db.get(Dataset, params.dataset_id)
+            if not dataset:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Dataset not found",
+                )
+            datasetdict = load_dataset(f"{dataset.file_path}/dataset")
+            if not datasetdict:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Error while loading the dataset.",
+                )
+            inputs_names = parse_columns_indices(
+                datasetdict=datasetdict, indices=params.inputs_columns
+            )
+            outputs_names = parse_columns_indices(
+                datasetdict=datasetdict, indices=params.outputs_columns
+            )
+
+        except exc.SQLAlchemyError as e:
+            log.exception(e)
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Internal database error",
+            ) from e
+    if params.task_name not in component_registry:
+        raise HTTPException(
+            status_code=404,
+            detail=f"Task {params.task_name} not found in the registry.",
+        )
+    task: BaseTask = component_registry[params.task_name]["class"]()
+    validation_response = {}
+    try:
+        prepared_dataset = task.prepare_for_task(
+            datasetdict=datasetdict, outputs_columns=outputs_names
+        )
+        task.validate_dataset_for_task(
+            dataset=prepared_dataset,
+            dataset_name=dataset.name,
+            input_columns=inputs_names,
+            output_columns=outputs_names,
+        )
+        validation_response["dataset_status"] = "valid"
+    except (TypeError, ValueError) as e:
+        validation_response["dataset_status"] = "invalid"
+        validation_response["error"] = str(e)
+    return validation_response
+
+
 @router.post("/", status_code=status.HTTP_201_CREATED)
 @inject
 async def create_experiment(
@@ -125,10 +196,25 @@ async def create_experiment(
                 raise HTTPException(
                     status_code=status.HTTP_404_NOT_FOUND, detail="Dataset not found"
                 )
+            datasetdict = load_dataset(f"{dataset.file_path}/dataset")
+            if not datasetdict:
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail="Error while loading the dataset.",
+                )
+            inputs_columns = parse_columns_indices(
+                datasetdict=datasetdict, indices=params.input_columns
+            )
+            outputs_columns = parse_columns_indices(
+                datasetdict=datasetdict, indices=params.output_columns
+            )
             experiment = Experiment(
                 dataset_id=params.dataset_id,
                 task_name=params.task_name,
                 name=params.name,
+                input_columns=inputs_columns,
+                output_columns=outputs_columns,
+                splits=params.splits,
             )
             db.add(experiment)
             db.commit()
diff --git a/DashAI/back/api/api_v1/endpoints/predict.py b/DashAI/back/api/api_v1/endpoints/predict.py
index 934f739d..64a7e996 100644
--- a/DashAI/back/api/api_v1/endpoints/predict.py
+++ b/DashAI/back/api/api_v1/endpoints/predict.py
@@ -1,4 +1,3 @@
-import json
 import logging
 from typing import Any, Callable, ContextManager, Dict, List, Union
 
@@ -12,8 +11,8 @@
 
 from DashAI.back.api.api_v1.schemas.predict_params import PredictParams
 from DashAI.back.containers import Container
-from DashAI.back.dataloaders.classes.dataloader import BaseDataLoader, to_dashai_dataset
-from DashAI.back.dependencies.database.models import Dataset as Dt
+from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset
+from DashAI.back.dataloaders.classes.dataloader import BaseDataLoader
 from DashAI.back.dependencies.database.models import Experiment, Run
 from DashAI.back.dependencies.registry import ComponentRegistry
 from DashAI.back.models.base_model import BaseModel
@@ -98,12 +97,6 @@ async def predict(
                     status_code=status.HTTP_404_NOT_FOUND, detail="Experiment not found"
                 )
 
-            # TODO: Use only experiment
-            dat: Dt = db.get(Dt, exp.dataset_id)
-            if not dat:
-                raise HTTPException(
-                    status_code=status.HTTP_404_NOT_FOUND, detail="Dataset not found"
-                )
         except exc.SQLAlchemyError as e:
             logger.exception(e)
             raise HTTPException(
@@ -126,19 +119,19 @@ async def predict(
             status_code=status.HTTP_409_CONFLICT,
             detail="A dataset with this name already exists",
         ) from e
+
     dataloader: BaseDataLoader = component_registry["JSONDataLoader"]["class"]()
     raw_dataset = dataloader.load_data(
         filepath_or_buffer=input_file,
         temp_path=str(tmp_path),
         params={"data_key": "data"},
     )
+    # TODO Extract this Code to DashAIDataset
     input_df = pd.DataFrame(raw_dataset["train"])
-    # TODO: Use feature_names from Experiment
-    input_df = input_df.reindex(columns=json.loads(dat.feature_names))
+    input_df = input_df.reindex(columns=exp.input_columns)
     raw_dataset["train"] = Dataset.from_pandas(input_df)
-
-    # Transform into DashAIDataset
-    dataset = to_dashai_dataset(raw_dataset, raw_dataset["train"].column_names, [])
+    # ---------------------------------------
+    dataset = to_dashai_dataset(raw_dataset)
 
     y_pred = trained_model.predict(dataset["train"])
 
diff --git a/DashAI/back/api/api_v1/schemas/datasets_params.py b/DashAI/back/api/api_v1/schemas/datasets_params.py
index 1f39a2db..8005e826 100644
--- a/DashAI/back/api/api_v1/schemas/datasets_params.py
+++ b/DashAI/back/api/api_v1/schemas/datasets_params.py
@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import Dict, Union
 
 from pydantic import BaseModel
 
@@ -18,10 +18,22 @@ class SplitParams(BaseModel):
 
 
 class DatasetParams(BaseModel):
-    task_name: str
     dataloader: str
     dataset_name: str
-    outputs_columns: List[str] = []
     splits_in_folders: bool = False
     splits: SplitParams
     dataloader_params: DataLoaderParams
+
+
+class ColumnSpecItemParams(BaseModel):
+    type: str
+    dtype: str
+
+
+class ColumnsSpecParams(BaseModel):
+    columns: Dict[str, ColumnSpecItemParams]
+
+
+class DatasetUpdateParams(BaseModel):
+    name: str = None
+    columns: Dict[str, ColumnSpecItemParams] = None
diff --git a/DashAI/back/api/api_v1/schemas/experiments_params.py b/DashAI/back/api/api_v1/schemas/experiments_params.py
index 2713a296..88c1cf68 100644
--- a/DashAI/back/api/api_v1/schemas/experiments_params.py
+++ b/DashAI/back/api/api_v1/schemas/experiments_params.py
@@ -1,3 +1,5 @@
+from typing import List
+
 from pydantic import BaseModel
 
 
@@ -5,3 +7,13 @@ class ExperimentParams(BaseModel):
     dataset_id: int
     task_name: str
     name: str
+    input_columns: List[int]
+    output_columns: List[int]
+    splits: str
+
+
+class ColumnsValidationParams(BaseModel):
+    task_name: str
+    dataset_id: int
+    inputs_columns: List[int]
+    outputs_columns: List[int]
diff --git a/DashAI/back/core/exceptions/__init__.py b/DashAI/back/core/exceptions/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/DashAI/back/dataloaders/classes/dashai_dataset.py b/DashAI/back/dataloaders/classes/dashai_dataset.py
index 94825e3b..19670add 100644
--- a/DashAI/back/dataloaders/classes/dashai_dataset.py
+++ b/DashAI/back/dataloaders/classes/dashai_dataset.py
@@ -2,12 +2,21 @@
 import json
 import os
 import pathlib
-from typing import Dict, List, Literal, Union
+from typing import Dict, List, Literal, Tuple, Union
 
 import numpy as np
+import pyarrow as pa
 from beartype import beartype
-from datasets import ClassLabel, Dataset, DatasetDict, Value, load_from_disk
+from datasets import (
+    ClassLabel,
+    Dataset,
+    DatasetDict,
+    Value,
+    concatenate_datasets,
+    load_from_disk,
+)
 from datasets.table import Table
+from sklearn.model_selection import train_test_split
 
 
 class DashAIDataset(Dataset):
@@ -17,8 +26,6 @@ class DashAIDataset(Dataset):
     def __init__(
         self,
         table: Table,
-        inputs_columns: List[str],
-        outputs_columns: List[str],
         *args,
         **kwargs,
     ):
@@ -28,63 +35,8 @@ def __init__(
         ----------
         table : Table
             Arrow table from which the dataset will be created
-        inputs_columns : List[str]
-            List of input column names.
-        outputs_columns : List[str]
-            List of output column names.
         """
         super().__init__(table, *args, **kwargs)
-        validate_inputs_outputs(self.column_names, inputs_columns, outputs_columns)
-        self._inputs_columns = inputs_columns
-        self._outputs_columns = outputs_columns
-
-    @property
-    @beartype
-    def inputs_columns(self) -> List[str]:
-        """Obtains the list of input columns.
-
-        Returns
-        -------
-        List[str]
-            List of input columns.
-        """
-        return self._inputs_columns
-
-    @inputs_columns.setter
-    @beartype
-    def inputs_columns(self, columns_names: List[str]) -> None:
-        """Set the input columns names.
-
-        Parameters
-        ----------
-        columns_names : List[str]
-            A list with the new input column names.
-        """
-        self._inputs_columns = columns_names
-
-    @property
-    @beartype
-    def outputs_columns(self) -> List[str]:
-        """Obtains the list of output columns.
-
-        Returns
-        -------
-        List[str]
-            List of output columns.
-        """
-        return self._outputs_columns
-
-    @outputs_columns.setter
-    @beartype
-    def outputs_columns(self, columns_names: List[str]) -> None:
-        """Set the output columns names.
-
-        Parameters
-        ----------
-        columns_names : List[str]
-            A list with the new output column names.
-        """
-        self._outputs_columns = columns_names
 
     @beartype
     def cast(self, *args, **kwargs) -> "DashAIDataset":
@@ -96,7 +48,7 @@ def cast(self, *args, **kwargs) -> "DashAIDataset":
             Dataset after cast
         """
         ds = super().cast(*args, **kwargs)
-        return DashAIDataset(ds._data, self.inputs_columns, self.outputs_columns)
+        return DashAIDataset(ds._data)
 
     @beartype
     def save_to_disk(self, dataset_path: str) -> None:
@@ -110,22 +62,6 @@ def save_to_disk(self, dataset_path: str) -> None:
             path where the dataset will be stored
         """
         super().save_to_disk(dataset_path)
-        with open(
-            os.path.join(dataset_path, "dashai_dataset_metadata.json"),
-            "w",
-            encoding="utf-8",
-        ) as dashai_info_file:
-            data_dashai = {
-                "inputs_columns": self.inputs_columns,
-                "outputs_columns": self.outputs_columns,
-            }
-            json.dump(
-                data_dashai,
-                dashai_info_file,
-                indent=2,
-                sort_keys=True,
-                ensure_ascii=False,
-            )
 
     @beartype
     def change_columns_type(self, column_types: Dict[str, str]) -> "DashAIDataset":
@@ -188,20 +124,6 @@ def remove_columns(self, column_names: Union[str, List[str]]) -> "DashAIDataset"
         # Update self with modified dataset attributes
         self.__dict__.update(modified_dataset.__dict__)
 
-        # Update the input and output columns
-        self._inputs_columns = [
-            col for col in self._inputs_columns if col not in column_names
-        ]
-        self._outputs_columns = [
-            col for col in self._outputs_columns if col not in column_names
-        ]
-
-        # Validate that inputs and outputs only contain elements that exist in
-        # names
-        validate_inputs_outputs(
-            self.column_names, self.inputs_columns, self.outputs_columns
-        )
-
         return self
 
     @beartype
@@ -250,14 +172,231 @@ def sample(
         return sample
 
 
+@beartype
+def load_dataset(dataset_path: str) -> DatasetDict:
+    """Load a DashAI dataset from its path.
+
+         This process cast each split into a DashAIdataset object.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path where the dataset is stored.
+
+    Returns
+    -------
+    DatasetDict
+        The loaded dataset.
+    """
+    dataset = load_from_disk(dataset_path=dataset_path)
+
+    for split in dataset:
+        dataset[split] = DashAIDataset(dataset[split].data)
+
+    return dataset
+
+
+@beartype
+def save_dataset(datasetdict: DatasetDict, path: Union[str, pathlib.Path]) -> None:
+    """Save the datasetdict with dashaidatasets inside.
+
+    Parameters
+    ----------
+    datasetdict : DatasetDict
+        The dataset to be saved.
+
+    datasetdict_path : str
+        Path where the dtaaset will be stored.
+
+    """
+    splits = []
+    for split in datasetdict:
+        splits.append(split)
+        datasetdict[split].save_to_disk(f"{path}/{split}")
+
+    with open(
+        os.path.join(path, "dataset_dict.json"), "w", encoding="utf-8"
+    ) as datasetdict_info_file:
+        data = {"splits": splits}
+        json.dump(
+            data,
+            datasetdict_info_file,
+            indent=2,
+            sort_keys=True,
+            ensure_ascii=False,
+        )
+
+
+@beartype
+def check_split_values(
+    train_size: float,
+    test_size: float,
+    val_size: float,
+) -> None:
+    if train_size < 0 or train_size > 1:
+        raise ValueError(
+            "train_size should be in the (0, 1) range "
+            f"(0 and 1 not included), got {val_size}"
+        )
+    if test_size < 0 or test_size > 1:
+        raise ValueError(
+            "test_size should be in the (0, 1) range "
+            f"(0 and 1 not included), got {val_size}"
+        )
+    if val_size < 0 or val_size > 1:
+        raise ValueError(
+            "val_size should be in the (0, 1) range "
+            f"(0 and 1 not included), got {val_size}"
+        )
+
+
+@beartype
+def split_indexes(
+    total_rows: int,
+    train_size: float,
+    test_size: float,
+    val_size: float,
+    seed: Union[int, None] = None,
+    shuffle: bool = True,
+) -> Tuple[List, List, List]:
+    """Generate lists with train, test and validation indexes.
+
+    The algorithm for splitting the dataset is as follows:
+
+    1. The dataset is divided into a training and a test-validation split
+        (sum of test_size and val_size).
+    2. The test and validation set is generated from the test-validation set,
+        where the size of the test-validation set is now considered to be 100%.
+        Therefore, the sizes of the test and validation sets will now be
+        calculated as 100%, i.e. as val_size/(test_size+val_size) and
+        test_size/(test_size+val_size) respectively.
+
+    Example:
+
+    If we split a dataset into 0.8 training, a 0.1 test, and a 0.1 validation,
+    in the first process we split the training data with 80% of the data, and
+    the test-validation data with the remaining 20%; and then in the second
+    process we split this 20% into 50% test and 50% validation.
+
+    Parameters
+    ----------
+    total_rows : int
+        Size of the Dataset.
+    train_size : float
+        Proportion of the dataset for train split (in 0-1).
+    test_size : float
+        Proportion of the dataset for test split (in 0-1).
+    val_size : float
+        Proportion of the dataset for validation split (in 0-1).
+    seed : Union[int, None], optional
+        Set seed to control to enable replicability, by default None
+    shuffle : bool, optional
+        If True, the data will be shuffled when splitting the dataset,
+        by default True.
+
+    Returns
+    -------
+    Tuple[List, List, List]
+        Train, Test and Validation indexes.
+    """
+
+    # Generate shuffled indexes
+    np.random.seed(seed)
+    indexes = np.arange(total_rows)
+
+    test_val = test_size + val_size
+    val_proportion = test_size / test_val
+    train_indexes, test_val_indexes = train_test_split(
+        indexes,
+        train_size=train_size,
+        random_state=seed,
+        shuffle=shuffle,
+    )
+    test_indexes, val_indexes = train_test_split(
+        test_val_indexes,
+        train_size=val_proportion,
+        random_state=seed,
+        shuffle=shuffle,
+    )
+    return list(train_indexes), list(test_indexes), list(val_indexes)
+
+
+@beartype
+def split_dataset(
+    dataset: Dataset,
+    train_indexes: List,
+    test_indexes: List,
+    val_indexes: List,
+) -> DatasetDict:
+    """Split the dataset in train, test and validation subsets.
+
+    Parameters
+    ----------
+    dataset : DatasetDict
+        A HuggingFace DatasetDict containing the dataset to be split.
+    train_indexes : List
+        Train split indexes.
+    test_indexes : List
+        Test split indexes.
+    val_indexes : List
+        Validation split indexes.
+
+
+    Returns
+    -------
+    DatasetDict
+        The split dataset.
+    """
+
+    # Get the number of records
+    n = len(dataset)
+
+    # Convert the indexes into boolean masks
+    train_mask = np.isin(np.arange(n), train_indexes)
+    test_mask = np.isin(np.arange(n), test_indexes)
+    val_mask = np.isin(np.arange(n), val_indexes)
+
+    # Get the underlying table
+    table = dataset.data
+
+    # Create separate tables for each split
+    train_table = table.filter(pa.array(train_mask))
+    test_table = table.filter(pa.array(test_mask))
+    val_table = table.filter(pa.array(val_mask))
+
+    separate_dataset_dict = DatasetDict(
+        {
+            "train": Dataset(train_table),
+            "test": Dataset(test_table),
+            "validation": Dataset(val_table),
+        }
+    )
+
+    dataset = to_dashai_dataset(separate_dataset_dict)
+    return dataset
+
+
+def to_dashai_dataset(dataset: DatasetDict) -> DatasetDict:
+    """
+    Convert all datasets within the DatasetDict to DashAIDataset.
+
+    Returns
+    -------
+    DatasetDict:
+        Datasetdict with datasets converted to DashAIDataset.
+    """
+    for key in dataset:
+        dataset[key] = DashAIDataset(dataset[key].data)
+    return dataset
+
+
 @beartype
 def validate_inputs_outputs(
-    names: List[str],
+    datasetdict: DatasetDict,
     inputs: List[str],
     outputs: List[str],
 ) -> None:
     """Validate the columns to be chosen as input and output.
-
     The algorithm considers those that already exist in the dataset.
 
     Parameters
@@ -269,85 +408,211 @@ def validate_inputs_outputs(
     outputs : List[str]
         List of output column names.
     """
-    if len(inputs) + len(outputs) > len(names):
+    dataset_features = list((datasetdict["train"].features).keys())
+    if len(inputs) == 0 or len(outputs) == 0:
+        raise ValueError(
+            "Inputs and outputs columns lists to validate must not be empty"
+        )
+    if len(inputs) + len(outputs) > len(dataset_features):
         raise ValueError(
             "Inputs and outputs cannot have more elements than names. "
             f"Number of inputs: {len(inputs)}, "
             f"number of outputs: {len(outputs)}, "
-            f"number of names: {len(names)}. "
+            f"number of names: {len(dataset_features)}. "
         )
         # Validate that inputs and outputs only contain elements that exist in names
-    if not set(names).issuperset(set(inputs + outputs)):
+    if not set(dataset_features).issuperset(set(inputs + outputs)):
         raise ValueError(
-            "Inputs and outputs can only contain elements that exist in names."
-        )
-        # Validate that the union of inputs and outputs is equal to names
-    if set(inputs + outputs) != set(names):
-        raise ValueError(
-            "The union of the elements of inputs and outputs list must be equal to "
-            "elements in the list of names."
+            f"Inputs and outputs can only contain elements that exist in names. "
+            f"Extra elements: "
+            f"{', '.join(set(inputs + outputs).difference(set(dataset_features)))}"
         )
 
 
 @beartype
-def load_dataset(dataset_path: str) -> DatasetDict:
-    """Load a datasetdict with dashaidatasets inside.
+def parse_columns_indices(datasetdict: DatasetDict, indices: List[int]) -> List[str]:
+    """Returns the column labes of the dataset that correspond to the indices
+
+    Args:
+        dataset_path (str): Path where the dataset is stored
+        indices (List[int]): List with the indices of the columns
+
+    Returns:
+        List[str]: List with the labels of the columns
+    """
+    dataset_features = list((datasetdict["train"].features).keys())
+    names_list = []
+    for index in indices:
+        if index > len(dataset_features):
+            raise ValueError(
+                f"The list of indices can only contain elements within"
+                f" the amount of columns. "
+                f"Index {index} is greater than the total of columns."
+            )
+        names_list.append(dataset_features[index - 1])
+    return names_list
+
+
+@beartype
+def select_columns(
+    dataset: DatasetDict, input_columns: List[str], output_columns: List[str]
+) -> Tuple[DatasetDict, DatasetDict]:
+    """Divide the dataset into a dataset with only the input columns in it
+    and other dataset only with the output columns
 
     Parameters
     ----------
-    dataset_path : str
-        Path where the dataset is stored.
+    dataset : DatasetDict
+        Dataset to divide
+    input_columns : List[str]
+        List with the input columns labels
+    output_columns : List[str]
+        List with the output columns labels
 
     Returns
     -------
-    DatasetDict
-        The loaded dataset.
+    Tuple[DatasetDict, DatasetDict]
+        Tuple with the separated DatasetDicts x and y
     """
-    dataset = load_from_disk(dataset_path=dataset_path)
-
+    input_columns_dataset = DatasetDict()
+    output_columns_dataset = DatasetDict()
     for split in dataset:
-        path = os.path.join(dataset_path, f"{split}/dashai_dataset_metadata.json")
-        with open(path, "r", encoding="utf-8") as dashai_info_file:
-            dataset_dashai_info = json.load(dashai_info_file)
+        input_columns_dataset[split] = dataset[split].select_columns(input_columns)
+        output_columns_dataset[split] = dataset[split].select_columns(output_columns)
+    return (input_columns_dataset, output_columns_dataset)
 
-        inputs_columns = dataset_dashai_info["inputs_columns"]
-        outputs_columns = dataset_dashai_info["outputs_columns"]
 
-        dataset[split] = DashAIDataset(
-            dataset[split].data,
-            inputs_columns,
-            outputs_columns,
-        )
+@beartype
+def get_columns_spec(dataset_path: str) -> Dict[str, Dict]:
+    """Return the column with their respective types
 
-    return dataset
+    Parameters
+    ----------
+    dataset_path : str
+        Path where the dataset is stored.
+
+    Returns
+    -------
+    Dict
+        Dict with the columns and types
+    """
+    dataset = load_dataset(dataset_path=dataset_path)
+    dataset_features = dataset["train"].features
+    column_types = {}
+    for column in dataset_features:
+        if dataset_features[column]._type == "Value":
+            column_types[column] = {
+                "type": "Value",
+                "dtype": dataset_features[column].dtype,
+            }
+        elif dataset_features[column]._type == "ClassLabel":
+            column_types[column] = {
+                "type": "Classlabel",
+                "dtype": "",
+            }
+    return column_types
 
 
 @beartype
-def save_dataset(datasetdict: DatasetDict, path: Union[str, pathlib.Path]) -> None:
-    """Save the datasetdict with dashaidatasets inside.
+def update_columns_spec(dataset_path: str, columns: Dict) -> DatasetDict:
+    """Return the column with their respective types
 
     Parameters
     ----------
-    datasetdict : DatasetDict
-        The dataset to be saved.
+    dataset_path : str
+        Path where the dataset is stored.
+    columns : Dict
+        Dict with columns and types to change
+    Returns
+    -------
+    Dict
+        Dict with the columns and types
+    """
+    if not isinstance(columns, dict):
+        raise TypeError(f"types should be a dict, got {type(columns)}")
+
+    # Load the dataset from where its stored
+    dataset_dict = load_from_disk(dataset_path=dataset_path)
+    for split in dataset_dict:
+        # Copy the features with the columns ans types
+        new_features = dataset_dict[split].features
+        for column in columns:
+            if columns[column].type == "ClassLabel":
+                names = list(set(dataset_dict[split][column]))
+                new_features[column] = ClassLabel(names=names)
+            elif columns[column].type == "Value":
+                new_features[column] = Value(columns[column].dtype)
+        # Cast the column types with the changes
+        try:
+            dataset_dict[split] = dataset_dict[split].cast(new_features)
+        except ValueError as e:
+            raise ValueError("Error while trying to cast the columns") from e
+    return dataset_dict
 
-    datasetdict_path : str
-        Path where the dtaaset will be stored.
 
+def get_dataset_info(dataset_path: str) -> object:
+    """Return the info of the dataset with the number of rows,
+    number of columns and splits size.
+
+    Parameters
+    ----------
+    dataset_path : str
+        Path where the dataset is stored.
+
+    Returns
+    -------
+    object
+        Dictionary with the information of the dataset
     """
-    splits = []
-    for split in datasetdict:
-        splits.append(split)
-        datasetdict[split].save_to_disk(f"{path}/{split}")
+    dataset = load_dataset(dataset_path=dataset_path)
+    total_rows = sum(split.num_rows for split in dataset.values())
+    total_columns = len(dataset["train"].features)
+    dataset_info = {
+        "total_rows": total_rows,
+        "total_columns": total_columns,
+        "train_size": dataset["train"].num_rows,
+        "test_size": dataset["test"].num_rows,
+        "val_size": dataset["validation"].num_rows,
+    }
+    return dataset_info
 
-    with open(
-        os.path.join(path, "dataset_dict.json"), "w", encoding="utf-8"
-    ) as datasetdict_info_file:
-        data = {"splits": splits}
-        json.dump(
-            data,
-            datasetdict_info_file,
-            indent=2,
-            sort_keys=True,
-            ensure_ascii=False,
+
+@beartype
+def update_dataset_splits(
+    datasetdict: DatasetDict, new_splits: object, is_random: bool
+) -> DatasetDict:
+    """Splits an already separated dataset by concatenating it and applying
+    new splits. The splits could be random by giving numbers between 0 and 1
+    in new_splits parameters and setting the is_random parameter to True, or
+    the could be manually selected by giving lists of indices to new_splits
+    parameter and setting the is_random parameter to False.
+
+    Args:
+        datasetdict (DatasetDict): Dataset to update splits
+        new_splits (object): Object with the new train, test and validation config
+        is_random (bool): If the new splits are random by percentage
+
+    Returns:
+        DatasetDict: New DatasetDict with the new splits configuration
+    """
+    concatenated_dataset = concatenate_datasets(
+        [datasetdict["train"], datasetdict["test"], datasetdict["validation"]]
+    )
+    n = len(concatenated_dataset)
+    if is_random:
+        check_split_values(
+            new_splits["train"], new_splits["test"], new_splits["validation"]
+        )
+        train_indexes, test_indexes, val_indexes = split_indexes(
+            n, new_splits["train"], new_splits["test"], new_splits["validation"]
         )
+    else:
+        train_indexes = new_splits["train"]
+        test_indexes = new_splits["test"]
+        val_indexes = new_splits["validation"]
+    return split_dataset(
+        dataset=concatenated_dataset,
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
+    )
diff --git a/DashAI/back/dataloaders/classes/dataloader.py b/DashAI/back/dataloaders/classes/dataloader.py
index c9b767af..2e60ef38 100644
--- a/DashAI/back/dataloaders/classes/dataloader.py
+++ b/DashAI/back/dataloaders/classes/dataloader.py
@@ -5,16 +5,12 @@
 import os
 import zipfile
 from abc import abstractmethod
-from typing import Any, Dict, Final, List, Union
+from typing import Any, Dict, Final, Union
 
-import numpy as np
-from beartype import beartype
-from datasets import Dataset, DatasetDict
-from sklearn.model_selection import train_test_split
+from datasets import DatasetDict
 from starlette.datastructures import UploadFile
 
 from DashAI.back.config_object import ConfigObject
-from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 
 logger = logging.getLogger(__name__)
 
@@ -89,182 +85,3 @@ def extract_files(self, dataset_path: str, file: UploadFile) -> str:
             with open(files_path, "wb") as f:
                 f.write(file.file.read())
         return files_path
-
-    def _check_split_values(
-        self,
-        dataset: DatasetDict,
-        train_size: float,
-        test_size: float,
-        val_size: float,
-        stratify: bool = False,
-        class_column: Union[str, None] = None,
-    ) -> None:
-        if stratify and class_column is None:
-            raise ValueError("Stratify requires that class_column is not none")
-
-        if (
-            class_column is not None
-            and class_column not in dataset["train"].column_names
-        ):
-            raise ValueError(f"Column '{class_column}' do not exist in the dataset.")
-
-        if train_size < 0 or train_size > 1:
-            raise ValueError(
-                "train_size should be in the (0, 1) range "
-                f"(0 and 1 not included), got {val_size}"
-            )
-        if test_size < 0 or test_size > 1:
-            raise ValueError(
-                "test_size should be in the (0, 1) range "
-                f"(0 and 1 not included), got {val_size}"
-            )
-        if val_size < 0 or val_size > 1:
-            raise ValueError(
-                "val_size should be in the (0, 1) range "
-                f"(0 and 1 not included), got {val_size}"
-            )
-
-    @beartype
-    def split_dataset(
-        self,
-        dataset: DatasetDict,
-        train_size: float,
-        test_size: float,
-        val_size: float,
-        seed: Union[int, None] = None,
-        shuffle: bool = True,
-        stratify: bool = False,
-        class_column: Union[str, None] = None,
-    ) -> DatasetDict:
-        """Split the dataset in train, test and validation subsets.
-
-        The algorithm for splitting the dataset is as follows:
-
-        1. The dataset is divided into a training and a test-validation split
-           (sum of test_size and val_size).
-        2. The test and validation set is generated from the test-validation set,
-           where the size of the test-validation set is now considered to be 100%.
-           Therefore, the sizes of the test and validation sets will now be
-           calculated as 100%, i.e. as val_size/(test_size+val_size) and
-           test_size/(test_size+val_size) respectively.
-
-        Example:
-
-        If we split a dataset into 0.8 training, a 0.1 test, and a 0.1 validation,
-        in the first process we split the training data with 80% of the data, and
-        the test-validation data with the remaining 20%; and then in the second
-        process we split this 20% into 50% test and 50% validation.
-
-        Parameters
-        ----------
-        dataset : DatasetDict
-            A HuggingFace DatasetDict containing the dataset to be split.
-        train_size : float
-            Proportion of the dataset for train split (in 0-1).
-        test_size : float
-            Proportion of the dataset for test split (in 0-1).
-        val_size : float
-            Proportion of the dataset for validation split (in 0-1).
-        seed : Union[int, None], optional
-            Set seed to control to enable replicability, by default None
-        shuffle : bool, optional
-            If True, the data will be shuffled when splitting the dataset,
-            by default True.
-        stratify : bool, optional
-            True indicates the split will be stratified, by default False
-        class_column : str, optional
-            Indicate the column with which to stratify, by default None
-
-        Returns
-        -------
-        DatasetDict
-            The split dataset.
-        """
-        self._check_split_values(
-            dataset, train_size, test_size, val_size, stratify, class_column
-        )
-        inputs_columns = dataset["train"].inputs_columns
-        outputs_columns = dataset["train"].outputs_columns
-
-        # Get the number of records
-        n = len(dataset["train"])
-
-        # Generate shuffled indices
-        np.random.seed(seed)
-        indices = np.arange(n)
-
-        test_val = test_size + val_size
-        val_proportion = test_size / test_val
-
-        # Define stratification array if stratify is True
-        stratify_array = (
-            np.array(dataset["train"][class_column])
-            if stratify and class_column
-            else None
-        )
-
-        # Split the indices
-        train_indices, test_val_indices = train_test_split(
-            indices,
-            train_size=train_size,
-            random_state=seed,
-            stratify=stratify_array,
-            shuffle=shuffle,
-        )
-        test_indices, val_indices = train_test_split(
-            test_val_indices,
-            train_size=val_proportion,
-            random_state=seed,
-            stratify=stratify_array[test_val_indices]
-            if stratify_array is not None
-            else None,
-            shuffle=shuffle,
-        )
-
-        # Convert the indices into boolean masks
-        train_mask = np.isin(np.arange(n), train_indices)
-        test_mask = np.isin(np.arange(n), test_indices)
-        val_mask = np.isin(np.arange(n), val_indices)
-
-        # Get the underlying table
-        table = dataset["train"].data
-
-        # Create separate tables for each split
-        train_table = table.filter(train_mask)
-        test_table = table.filter(test_mask)
-        val_table = table.filter(val_mask)
-
-        separate_dataset_dict = DatasetDict(
-            {
-                "train": Dataset(train_table),
-                "test": Dataset(test_table),
-                "validation": Dataset(val_table),
-            }
-        )
-
-        dataset = to_dashai_dataset(
-            separate_dataset_dict, inputs_columns, outputs_columns
-        )
-        return dataset
-
-
-def to_dashai_dataset(
-    dataset: DatasetDict,
-    inputs_columns: List[str],
-    outputs_columns: List[str],
-) -> DatasetDict:
-    """
-    Convert all datasets within the DatasetDict to DashAIDataset.
-
-    Returns
-    -------
-    DatasetDict:
-        Datasetdict with datasets converted to DashAIDataset.
-    """
-    for key in dataset:
-        dataset[key] = DashAIDataset(
-            dataset[key].data,
-            inputs_columns,
-            outputs_columns,
-        )
-    return dataset
diff --git a/DashAI/back/dependencies/database/models.py b/DashAI/back/dependencies/database/models.py
index c07811e4..1c78cabd 100644
--- a/DashAI/back/dependencies/database/models.py
+++ b/DashAI/back/dependencies/database/models.py
@@ -18,8 +18,6 @@ class Dataset(Base):
     """
     id: Mapped[int] = mapped_column(primary_key=True)
     name: Mapped[str] = mapped_column(String, unique=True, nullable=False)
-    task_name: Mapped[str] = mapped_column(String, nullable=False)
-    feature_names: Mapped[str] = mapped_column(JSON, nullable=False)
     created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now)
     last_modified: Mapped[DateTime] = mapped_column(
         DateTime,
@@ -39,6 +37,9 @@ class Experiment(Base):
     dataset_id: Mapped[int] = mapped_column(ForeignKey("dataset.id"))
     name: Mapped[str] = mapped_column(String, unique=True, nullable=False)
     task_name: Mapped[str] = mapped_column(String, nullable=False)
+    input_columns: Mapped[str] = mapped_column(JSON, nullable=False)
+    output_columns: Mapped[str] = mapped_column(JSON, nullable=False)
+    splits: Mapped[str] = mapped_column(JSON, nullable=False)
     created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now)
     last_modified: Mapped[DateTime] = mapped_column(
         DateTime,
diff --git a/DashAI/back/dependencies/registry/component_registry.py b/DashAI/back/dependencies/registry/component_registry.py
index 3396e732..963b3bf9 100644
--- a/DashAI/back/dependencies/registry/component_registry.py
+++ b/DashAI/back/dependencies/registry/component_registry.py
@@ -20,6 +20,7 @@ class ComponentRegistry:
         "class": Component1,          # Component class.
         "configurable_object": False, # True if the object is a configurable one.
         "schema": {...},              # Configurable object schema if applies.
+        "metadata": {...},            # Component metadata if applies.
         "description": "...",         # An object description.
     }
     ```
@@ -188,6 +189,9 @@ def register_component(self, new_component: Type) -> None:
             "class": new_component,
             "configurable_object": is_configurable_object,
             "schema": new_component.get_schema() if is_configurable_object else None,
+            "metadata": new_component.get_metadata()
+            if hasattr(new_component, "metadata")
+            else None,
             "description": getattr(new_component, "DESCRIPTION", None),
         }
 
diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py
index 1ef5825c..679c6656 100644
--- a/DashAI/back/job/model_job.py
+++ b/DashAI/back/job/model_job.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 from typing import List
@@ -6,7 +7,12 @@
 from sqlalchemy import exc
 from sqlalchemy.orm import Session
 
-from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset, load_dataset
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    DashAIDataset,
+    load_dataset,
+    select_columns,
+    update_dataset_splits,
+)
 from DashAI.back.dependencies.database.models import Dataset, Experiment, Run
 from DashAI.back.job.base_job import BaseJob, JobError
 from DashAI.back.metrics import BaseMetric
@@ -115,7 +121,24 @@ def run(
                 ) from e
 
             try:
-                prepared_dataset = task.prepare_for_task(loaded_dataset)
+                splits = json.loads(experiment.splits)
+                if splits["has_changed"]:
+                    new_splits = {
+                        "train": splits["train"],
+                        "test": splits["test"],
+                        "validation": splits["validation"],
+                    }
+                    loaded_dataset = update_dataset_splits(
+                        loaded_dataset, new_splits, splits["is_random"]
+                    )
+                prepared_dataset = task.prepare_for_task(
+                    loaded_dataset, experiment.output_columns
+                )
+                x, y = select_columns(
+                    prepared_dataset,
+                    experiment.input_columns,
+                    experiment.output_columns,
+                )
             except Exception as e:
                 log.exception(e)
                 raise JobError(
@@ -134,7 +157,7 @@ def run(
 
             try:
                 # Training
-                model.fit(prepared_dataset["train"])
+                model.fit(x["train"], y["train"])
             except Exception as e:
                 log.exception(e)
                 raise JobError(
@@ -154,8 +177,8 @@ def run(
                 model_metrics = {
                     split: {
                         metric.__name__: metric.score(
-                            prepared_dataset[split],
-                            model.predict(prepared_dataset[split]),
+                            y[split],
+                            model.predict(x[split]),
                         )
                         for metric in metrics
                     }
diff --git a/DashAI/back/metrics/classification_metric.py b/DashAI/back/metrics/classification_metric.py
index f03779df..722d7212 100644
--- a/DashAI/back/metrics/classification_metric.py
+++ b/DashAI/back/metrics/classification_metric.py
@@ -35,15 +35,15 @@ def validate_inputs(true_labels: np.ndarray, pred_labels: np.ndarray) -> None:
 
 
 def prepare_to_metric(
-    dataset: DashAIDataset,
+    y: DashAIDataset,
     probs_pred_labels: np.ndarray,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """Prepare true and prediced labels to be used later in metrics.
 
     Parameters
     ----------
-    dataset : DashAIDataset
-        A dashai dataset with labels included.
+    y : DashAIDataset
+        A DashAIDataset with the output columns of the data.
     probs_pred_labels : np.ndarray
         A two-dimensional matrix in which each column represents a class and the row
         values represent the probability that an example belongs to the class
@@ -54,8 +54,8 @@ def prepare_to_metric(
     Tuple[np.ndarray, np.ndarray]
         A tuple with the true and predicted labels in numpy format.
     """
-    output_column = dataset.outputs_columns[0]
-    true_labels = np.array(dataset[output_column])
+    column_name = y.column_names[0]
+    true_labels = np.array(y[column_name])
     validate_inputs(true_labels, probs_pred_labels)
     pred_labels = np.argmax(probs_pred_labels, axis=1)
     return true_labels, pred_labels
diff --git a/DashAI/back/metrics/translation/bleu.py b/DashAI/back/metrics/translation/bleu.py
index 5c14c4e9..92ff0b7f 100644
--- a/DashAI/back/metrics/translation/bleu.py
+++ b/DashAI/back/metrics/translation/bleu.py
@@ -1,5 +1,6 @@
 """BLEU (bilingual evaluation understudy) metric implementation for DashAI."""
 import evaluate
+import numpy as np
 
 from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.metrics.translation_metric import TranslationMetric, prepare_to_metric
@@ -17,14 +18,14 @@ class Bleu(TranslationMetric):
     """
 
     @staticmethod
-    def score(source_sentences: DashAIDataset, target_sentences: list):
+    def score(source_sentences: DashAIDataset, target_sentences: np.ndarray):
         """Calculate the BLEU score between source and target sentences.
 
         Parameters
         ----------
         source_sentences : DashAIDataset
             Sentences in the original language.
-        target_sentences : list
+        target_sentences : ndarray
             Sentences in the target language.
 
         Returns
diff --git a/DashAI/back/metrics/translation/ter.py b/DashAI/back/metrics/translation/ter.py
index e3a38e7b..588e41ba 100644
--- a/DashAI/back/metrics/translation/ter.py
+++ b/DashAI/back/metrics/translation/ter.py
@@ -1,5 +1,6 @@
 """TER (Translation Edit Rate) metric implementation for DashAI."""
 import evaluate
+import numpy as np
 
 from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.metrics.translation_metric import TranslationMetric, prepare_to_metric
@@ -17,14 +18,14 @@ class Ter(TranslationMetric):
     """
 
     @staticmethod
-    def score(source_sentences: DashAIDataset, target_sentences: list):
+    def score(source_sentences: DashAIDataset, target_sentences: np.ndarray):
         """Calculate the TER score between source and target sentences.
 
         Parameters
         ----------
         source_sentences : DashAIDataset
             Sentences in the original language.
-        target_sentences : list
+        target_sentences : ndarray
             Sentences in the target language.
 
         Returns
diff --git a/DashAI/back/metrics/translation_metric.py b/DashAI/back/metrics/translation_metric.py
index bf4c741b..f4f565f0 100644
--- a/DashAI/back/metrics/translation_metric.py
+++ b/DashAI/back/metrics/translation_metric.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.metrics.base_metric import BaseMetric
 
@@ -22,13 +24,13 @@ def validate_inputs(true_labels: list, pred_labels: list) -> None:
         raise ValueError("The length of the true and predicted labels must be equal.")
 
 
-def prepare_to_metric(source_sentences: DashAIDataset, target_sentences: list):
+def prepare_to_metric(y: DashAIDataset, target_sentences: list):
     """Format labels to be used in metrics.
 
     Parameters
     ----------
-    source_sentences : DashAIDataset
-        True sentences.
+    y : DashAIDataset
+        True sentences of the dataset.
     target_sentences : list
         Target sentences.
 
@@ -37,7 +39,7 @@ def prepare_to_metric(source_sentences: DashAIDataset, target_sentences: list):
     tuple
         A tuple with the true and predicted sentences.
     """
-    output_column = source_sentences.outputs_columns[0]
-    source_sentences = [[example[output_column]] for example in source_sentences]
+    column_name = y.column_names[0]
+    source_sentences = np.array(y[column_name])
     validate_inputs(source_sentences, target_sentences)
     return source_sentences, target_sentences
diff --git a/DashAI/back/models/hugging_face/distilbert_transformer.py b/DashAI/back/models/hugging_face/distilbert_transformer.py
index cc092977..3a50d5df 100644
--- a/DashAI/back/models/hugging_face/distilbert_transformer.py
+++ b/DashAI/back/models/hugging_face/distilbert_transformer.py
@@ -1,8 +1,9 @@
 """DashAI implementation of DistilBERT model for english classification."""
 import shutil
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional
 
 import numpy as np
+from datasets import Dataset
 from sklearn.exceptions import NotFittedError
 from transformers import (
     DistilBertForSequenceClassification,
@@ -11,7 +12,6 @@
     TrainingArguments,
 )
 
-from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.models.text_classification_model import TextClassificationModel
 
 
@@ -48,14 +48,18 @@ def __init__(self, model=None, **kwargs):
             self.batch_size = kwargs.pop("batch_size")
             self.device = kwargs.pop("device")
 
-    def get_tokenizer(self, input_column: str, output_column: str) -> Callable:
+    def get_tokenizer(
+        self,
+        input_column: str,
+        output_column: Optional[str] = None,
+    ) -> Callable:
         """Tokenize input and output.
 
         Parameters
         ----------
         input_column : str
             name the input column to be tokenized.
-        output_column : str
+        output_column : Optional[str]
             name the output column to be tokenized.
 
         Returns
@@ -65,7 +69,7 @@ def get_tokenizer(self, input_column: str, output_column: str) -> Callable:
         """
 
         def _tokenize(batch) -> Dict[str, Any]:
-            return {
+            tokenized_batch = {
                 "input_ids": self.tokenizer(
                     batch[input_column],
                     padding="max_length",
@@ -78,12 +82,14 @@ def _tokenize(batch) -> Dict[str, Any]:
                     truncation=True,
                     max_length=512,
                 )["attention_mask"],
-                "labels": batch[output_column],
             }
+            if output_column:
+                tokenized_batch["labels"] = batch[output_column]
+            return tokenized_batch
 
         return _tokenize
 
-    def fit(self, dataset: DashAIDataset):
+    def fit(self, x: Dataset, y: Dataset):
         """Fine-tune the pre-trained model.
 
         Parameters
@@ -92,8 +98,9 @@ def fit(self, dataset: DashAIDataset):
             DashAIDataset with training data.
 
         """
-        input_column = dataset.inputs_columns[0]
-        output_column = dataset.outputs_columns[0]
+        input_column = x.column_names[0]
+        output_column = y.column_names[0]
+        dataset = x.add_column(output_column, y[output_column])
 
         tokenizer_func = self.get_tokenizer(input_column, output_column)
         dataset = dataset.map(tokenizer_func, batched=True, batch_size=self.batch_size)
@@ -125,13 +132,13 @@ def fit(self, dataset: DashAIDataset):
 
         return self
 
-    def predict(self, dataset: DashAIDataset) -> np.array:
+    def predict(self, x: Dataset) -> np.array:
         """Make a prediction with the fine-tuned model.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            DashAIDataset with text data.
+        x : Dataset
+            Dataset with text data.
 
         Returns
         -------
@@ -145,18 +152,15 @@ def predict(self, dataset: DashAIDataset) -> np.array:
                 "estimator."
             )
 
-        input_column = dataset.inputs_columns[0]
-        output_column = dataset.outputs_columns[0]
-        tokenizer_func = self.get_tokenizer(input_column, output_column)
-        dataset = dataset.map(tokenizer_func, batched=True, batch_size=self.batch_size)
-        dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
+        input_column = x.column_names[0]
+        tokenizer_func = self.get_tokenizer(input_column)
+        x = x.map(tokenizer_func, batched=True, batch_size=self.batch_size)
+        x.set_format("torch", columns=["input_ids", "attention_mask"])
 
         probabilities = []
 
         # Calculate the number of batches
-        num_batches = len(dataset) // self.batch_size + (
-            len(dataset) % self.batch_size > 0
-        )
+        num_batches = len(x) // self.batch_size + (len(x) % self.batch_size > 0)
 
         # Iterate over each batch
         for i in range(num_batches):
@@ -165,9 +169,8 @@ def predict(self, dataset: DashAIDataset) -> np.array:
 
             # Extract the batch from the dataset
             batch = {
-                "input_ids": dataset["input_ids"][start_idx:end_idx],
-                "attention_mask": dataset["attention_mask"][start_idx:end_idx],
-                "labels": dataset["labels"][start_idx:end_idx],
+                "input_ids": x["input_ids"][start_idx:end_idx],
+                "attention_mask": x["attention_mask"][start_idx:end_idx],
             }
 
             # Make sure that the tensors are in the correct device.
diff --git a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py
index 01f560ee..8a0fd2cb 100644
--- a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py
+++ b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py
@@ -1,7 +1,8 @@
 """OpusMtEnESTransformer model for english-spanish translation DashAI implementation."""
 import shutil
-from typing import List
+from typing import List, Optional
 
+from datasets import Dataset
 from sklearn.exceptions import NotFittedError
 from transformers import (
     AutoModelForSeq2SeqLM,
@@ -10,7 +11,6 @@
     Seq2SeqTrainingArguments,
 )
 
-from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.models.translation_model import TranslationModel
 
 
@@ -30,8 +30,8 @@ def __init__(self, model=None, **kwargs):
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         if model is None:
             self.training_args = kwargs
-            self.batch_size = kwargs.pop("batch_size")
-            self.device = kwargs.pop("device")
+            self.batch_size = kwargs.pop("batch_size", 16)
+            self.device = kwargs.pop("device", "gpu")
         self.model = (
             model
             if model is not None
@@ -39,35 +39,70 @@ def __init__(self, model=None, **kwargs):
         )
         self.fitted = model is not None
 
-    def fit(self, dataset: DashAIDataset):
-        """Fine-tune the pre-trained model.
+    def tokenize_data(self, x: Dataset, y: Optional[Dataset] = None) -> Dataset:
+        """Tokenize input and output.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            DashAIDataset with training data.
+        x: Dataset
+            Dataset with the input data to preprocess.
+        y: Optional Dataset
+            Dataset with the output data to preprocess.
 
+        Returns
+        -------
+        Dataset
+            Dataset with the processed data.
         """
-        input_column = dataset.inputs_columns[0]
-        output_column = dataset.outputs_columns[0]
-
-        def _tokenize(examples):
-            inputs = self.tokenizer(
-                examples[input_column],
+        is_y = bool(y)
+        if not y:
+            y = Dataset.from_list([{"foo": 0}] * len(x))
+        # Initialize useful variables
+        dataset = []
+        input_column_name = x.column_names[0]
+        output_column_name = y.column_names[0]
+
+        # Preprocess both datasets
+        for input_sample, output_sample in zip(x, y):  # noqa
+            tokenized_input = self.tokenizer(
+                input_sample[input_column_name],
                 truncation=True,
                 padding="max_length",
                 max_length=512,
             )
-            outputs = self.tokenizer(
-                examples[output_column],
-                truncation=True,
-                padding="max_length",
-                max_length=512,
+            tokenized_output = (
+                self.tokenizer(
+                    output_sample[output_column_name],
+                    truncation=True,
+                    padding="max_length",
+                    max_length=512,
+                )
+                if is_y
+                else None
             )
-            inputs["labels"] = outputs["input_ids"]
-            return inputs
+            sample = {
+                "input_ids": tokenized_input["input_ids"],
+                "attention_mask": tokenized_input["attention_mask"],
+                "labels": tokenized_output["input_ids"]
+                if is_y
+                else y[output_column_name],
+            }
+            dataset.append(sample)
+        return Dataset.from_list(dataset)
 
-        dataset = dataset.map(_tokenize, batched=True)
+    def fit(self, x_train: Dataset, y_train: Dataset):
+        """Fine-tune the pre-trained model.
+
+        Parameters
+        ----------
+        x_train : Dataset
+            Dataset with input training data.
+        y_train : Dataset
+            Dataset with output training data.
+
+        """
+
+        dataset = self.tokenize_data(x_train, y_train)
         dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
 
         # Arguments for fine-tuning
@@ -95,13 +130,13 @@ def _tokenize(examples):
         )
         return self
 
-    def predict(self, dataset: DashAIDataset) -> List:
+    def predict(self, x_pred: Dataset) -> List:
         """Predict with the fine-tuned model.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            DashAIDataset with text data.
+        x_pred : Dataset
+            Dataset with text data.
 
         Returns
         -------
@@ -115,17 +150,7 @@ def predict(self, dataset: DashAIDataset) -> List:
                 "estimator."
             )
 
-        input_column = dataset.inputs_columns[0]
-
-        def encode(examples):
-            return self.tokenizer(
-                examples[input_column],
-                truncation=True,
-                padding="max_length",
-                max_length=512,
-            )
-
-        dataset = dataset.map(encode, batched=True)
+        dataset = self.tokenize_data(x_pred)
         dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])
 
         translations = []
diff --git a/DashAI/back/models/hugging_face/vit_transformer.py b/DashAI/back/models/hugging_face/vit_transformer.py
index 59983b77..a2722c28 100644
--- a/DashAI/back/models/hugging_face/vit_transformer.py
+++ b/DashAI/back/models/hugging_face/vit_transformer.py
@@ -1,7 +1,9 @@
 """DashAI implementation of DistilBERT model for image classification."""
 import shutil
+from typing import Optional
 
 import numpy as np
+from datasets import Dataset
 from sklearn.exceptions import NotFittedError
 from transformers import (
     Trainer,
@@ -10,7 +12,6 @@
     ViTForImageClassification,
 )
 
-from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.models.image_classification_model import ImageClassificationModel
 
 
@@ -42,49 +43,65 @@ def __init__(self, model=None, **kwargs):
         self.fitted = model is not None
         if model is None:
             self.training_args = kwargs
-            self.batch_size = kwargs.pop("batch_size")
-            self.device = kwargs.pop("device")
+            self.batch_size = kwargs.pop("batch_size", 8)
+            self.device = kwargs.pop("device", "gpu")
 
-    def get_preprocess_images(self, input_column: str, output_column: str):
+    def preprocess_images(self, x: Dataset, y: Optional[Dataset] = None):
         """Preprocess images for model input.
 
         Parameters
         ----------
-        input_column : str
-            name of the column containing the images to be preprocessed.
-        output_column : str
-            name of the column containing the output labels for the images.
+        x: Dataset
+            Dataset with the input data to preprocess.
+        y: Optional Dataset
+            Dataset with the output data to preprocess.
 
         Returns
         -------
-        Function
-            a function that preprocesses images and outputs a dictionary
-            containing processed images and corresponding labels.
+        Dataset
+            Dataset with the processed data.
         """
 
-        def _preprocess_images(examples):
-            inputs = self.feature_extractor(
-                images=examples[input_column], return_tensors="pt", size=224
+        # If the output datset is not given, create an empty dataset
+        if not y:
+            y = Dataset.from_list([{"foo": 0}] * len(x))
+        # Initialize useful variables
+        dataset = []
+        input_column_name = x.column_names[0]
+        output_column_name = y.column_names[0]
+
+        # Preprocess both datasets
+        for input_sample, output_sample in zip(x, y):  # noqa
+            preprocessed_input = self.feature_extractor(
+                images=input_sample[input_column_name], return_tensors="pt", size=224
             )
-            inputs["labels"] = examples[output_column]
-            return inputs
-
-        return _preprocess_images
+            reshaped_image = preprocessed_input["pixel_values"].reshape(
+                (
+                    preprocessed_input["pixel_values"].shape[1],
+                    preprocessed_input["pixel_values"].shape[2],
+                    preprocessed_input["pixel_values"].shape[3],
+                )
+            )
+            dataset.append(
+                {
+                    "pixel_values": reshaped_image,
+                    "labels": output_sample[output_column_name],
+                }
+            )
+        return Dataset.from_list(dataset)
 
-    def fit(self, dataset: DashAIDataset):
+    def fit(self, x_train: Dataset, y_train: Dataset):
         """Fine-tune the pre-trained model.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            DashAiDataset with training data.
+        x_train : Dataset
+            Dataset with input training data.
+        y_train: Dataset
+            Dataset with output training data.
 
         """
-        input_column = dataset.inputs_columns[0]
-        output_column = dataset.outputs_columns[0]
-
-        feature_extractor_func = self.get_preprocess_images(input_column, output_column)
-        dataset = dataset.map(feature_extractor_func, batched=True)
+        dataset = self.preprocess_images(x_train, y_train)
 
         # Arguments for fine-tuning
         training_args = TrainingArguments(
@@ -109,15 +126,14 @@ def fit(self, dataset: DashAIDataset):
         shutil.rmtree(
             "DashAI/back/user_models/temp_checkpoints_vit", ignore_errors=True
         )
-        return self
 
-    def predict(self, dataset: DashAIDataset) -> np.array:
+    def predict(self, x_pred: Dataset) -> np.array:
         """Make a prediction with the fine-tuned model.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            DashAIDataset with image data.
+        x_pred : Dataset
+            Dataset with image data.
 
         Returns
         -------
@@ -130,11 +146,7 @@ def predict(self, dataset: DashAIDataset) -> np.array:
                 " with appropriate arguments before using this estimator."
             )
 
-        input_column = dataset.inputs_columns[0]
-        output_column = dataset.outputs_columns[0]
-        preprocess_images = self.get_preprocess_images(input_column, output_column)
-
-        dataset = dataset.map(preprocess_images, batched=True)
+        dataset = self.preprocess_images(x_pred)
         dataset.set_format("torch", columns=["pixel_values", "labels"])
 
         probabilities = []
diff --git a/DashAI/back/models/scikit_learn/sklearn_like_model.py b/DashAI/back/models/scikit_learn/sklearn_like_model.py
index a25516d3..f8c48561 100644
--- a/DashAI/back/models/scikit_learn/sklearn_like_model.py
+++ b/DashAI/back/models/scikit_learn/sklearn_like_model.py
@@ -1,7 +1,6 @@
-from typing import Tuple, Type
+from typing import Type
 
 import joblib
-import pandas as pd
 
 from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 from DashAI.back.models.base_model import BaseModel
@@ -22,41 +21,38 @@ def load(filename: str) -> None:
 
     # --- Methods for process the data for sklearn models ---
 
-    def format_data(self, dataset: DashAIDataset) -> Tuple[pd.DataFrame, pd.Series]:
-        """Load and prepare the dataset into dataframes to use in Sklearn Models.
+    def fit(
+        self, x_train: DashAIDataset, y_train: DashAIDataset
+    ) -> Type["SklearnLikeModel"]:
+        """Fit the estimator.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            Dataset to format
+        x_train : pd.DataFrame
+            Dataframe with the input data.
+        y_train : pd.DataFrame
+            Dataframe with the output data.
 
         Returns
         -------
-        Dataframe
-            Dataframe with the data to use in experiments.
+        self
+            The fitted estimator object.
         """
-        data_in_pandas = dataset.to_pandas()
-        x = data_in_pandas.loc[:, dataset.inputs_columns]
-        y = data_in_pandas[dataset.outputs_columns]
-
-        return x, y
+        x_pandas = x_train.to_pandas()
+        y_pandas = y_train.to_pandas()
+        return super().fit(x_pandas, y_pandas)
 
-    def fit(self, dataset: DashAIDataset) -> Type["SklearnLikeModel"]:
-        """Fit the estimator.
+    def predict(self, x_pred: DashAIDataset):
+        """Make a prediction with the model.
 
         Parameters
         ----------
-        dataset : DashAIDataset
-            The training dataset.
+        x_pred : DashAIDataset
+            Dataset with the input data columns.
 
         Returns
         -------
-        self
-            The fitted estimator object.
+        array-like
+            Array with the predicted target values for x_pred
         """
-        x, y = self.format_data(dataset)
-        return super().fit(x, y)
-
-    def predict(self, dataset: DashAIDataset):
-        x, y = self.format_data(dataset)
-        return super().predict_proba(x)
+        return super().predict_proba(x_pred.to_pandas())
diff --git a/DashAI/back/tasks/base_task.py b/DashAI/back/tasks/base_task.py
index 0a412da8..86380a37 100644
--- a/DashAI/back/tasks/base_task.py
+++ b/DashAI/back/tasks/base_task.py
@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Any, Dict, Final
+from typing import Any, Dict, Final, List
 
 from datasets import DatasetDict
 
@@ -14,8 +14,36 @@ class BaseTask:
     def schema(self) -> Dict[str, Any]:
         raise NotImplementedError
 
+    @classmethod
+    def get_metadata(cls) -> Dict[str, Any]:
+        """Get metadata values for the current task
+
+        Returns:
+            Dict[str, Any]: Dictionary with the metadata containing the input and output
+             types/cardinality.
+        """
+        metadata = cls.metadata
+
+        # Extract class names
+        inputs_types = [input_type.__name__ for input_type in metadata["inputs_types"]]
+        outputs_types = [
+            output_type.__name__ for output_type in metadata["outputs_types"]
+        ]
+
+        parsed_metadata: dict = {
+            "inputs_types": inputs_types,
+            "outputs_types": outputs_types,
+            "inputs_cardinality": metadata["inputs_cardinality"],
+            "outputs_cardinality": metadata["outputs_cardinality"],
+        }
+        return parsed_metadata
+
     def validate_dataset_for_task(
-        self, dataset: DatasetDict, dataset_name: str
+        self,
+        dataset: DatasetDict,
+        dataset_name: str,
+        input_columns: List[str],
+        output_columns: List[str],
     ) -> None:
         """Validate a dataset for the current task.
 
@@ -27,14 +55,14 @@ def validate_dataset_for_task(
             Dataset name
         """
         for split in dataset:
-            schema = self.schema
-            allowed_input_types = tuple(schema["inputs_types"])
-            allowed_output_types = tuple(schema["outputs_types"])
-            inputs_cardinality = schema["inputs_cardinality"]
-            outputs_cardinality = schema["outputs_cardinality"]
+            metadata = self.metadata
+            allowed_input_types = tuple(metadata["inputs_types"])
+            allowed_output_types = tuple(metadata["outputs_types"])
+            inputs_cardinality = metadata["inputs_cardinality"]
+            outputs_cardinality = metadata["outputs_cardinality"]
 
             # Check input types
-            for input_col in dataset[split].inputs_columns:
+            for input_col in input_columns:
                 input_col_type = dataset[split].features[input_col]
                 if not isinstance(input_col_type, allowed_input_types):
                     raise TypeError(
@@ -43,39 +71,38 @@ def validate_dataset_for_task(
                     )
 
             # Check output types
-            for output_col in dataset[split].outputs_columns:
+            for output_col in output_columns:
                 output_col_type = dataset[split].features[output_col]
                 if not isinstance(output_col_type, allowed_output_types):
                     raise TypeError(
                         f"Error in split {split} of dataset {dataset_name}. "
-                        f"{output_col_type} is not an allowed type for output columns. "
+                        f"{output_col_type} is not an allowed type for output columns."
                     )
 
             # Check input cardinality
-            if (
-                inputs_cardinality != "n"
-                and len(dataset[split].inputs_columns) != inputs_cardinality
-            ):
+            if inputs_cardinality != "n" and len(input_columns) != inputs_cardinality:
                 raise ValueError(
                     f"Error in split {split} of dataset {dataset_name}. "
-                    f"Input cardinality ({len(dataset[split].inputs_columns)}) does not"
+                    f"Input cardinality ({len(input_columns)}) does not"
                     f" match task cardinality ({inputs_cardinality})"
                 )
 
             # Check output cardinality
             if (
                 outputs_cardinality != "n"
-                and len(dataset[split].outputs_columns) != outputs_cardinality
+                and len(output_columns) != outputs_cardinality
             ):
                 raise ValueError(
                     f"Error in split {split} of dataset {dataset_name}. "
-                    f"Output cardinality ({len(dataset[split].outputs_columns)})"
+                    f"Output cardinality ({len(output_columns)})"
                     f" does not "
                     f"match task cardinality ({outputs_cardinality})"
                 )
 
     @abstractmethod
-    def prepare_for_task(self, dataset: DatasetDict) -> DatasetDict:
+    def prepare_for_task(
+        self, dataset: DatasetDict, outputs_columns: List[str]
+    ) -> DatasetDict:
         """Change column types to suit the task requirements.
 
         Parameters
diff --git a/DashAI/back/tasks/image_classification_task.py b/DashAI/back/tasks/image_classification_task.py
index fcba4949..3f5444c5 100644
--- a/DashAI/back/tasks/image_classification_task.py
+++ b/DashAI/back/tasks/image_classification_task.py
@@ -1,3 +1,5 @@
+from typing import List
+
 from datasets import ClassLabel, DatasetDict, Image
 
 from DashAI.back.tasks.base_task import BaseTask
@@ -9,14 +11,16 @@ class ImageClassificationTask(BaseTask):
     Here you can change the methods provided by class Task.
     """
 
-    schema: dict = {
+    metadata: dict = {
         "inputs_types": [Image],
         "outputs_types": [ClassLabel],
         "inputs_cardinality": 1,
         "outputs_cardinality": 1,
     }
 
-    def prepare_for_task(self, datasetdict: DatasetDict):
+    def prepare_for_task(
+        self, datasetdict: DatasetDict, outputs_columns: List[str]
+    ) -> DatasetDict:
         """Change the column types to suit the tabular classification task.
 
         A copy of the dataset is created.
@@ -31,8 +35,8 @@ def prepare_for_task(self, datasetdict: DatasetDict):
         DatasetDict
             Dataset with the new types
         """
-        outputs_columns = datasetdict["train"].outputs_columns
-        types = {outputs_columns[0]: "Categorical"}
+
+        types = {column: "Categorical" for column in outputs_columns}
         for split in datasetdict:
             datasetdict[split] = datasetdict[split].change_columns_type(types)
         return datasetdict
diff --git a/DashAI/back/tasks/tabular_classification_task.py b/DashAI/back/tasks/tabular_classification_task.py
index f0dcb8c9..ae7571a8 100644
--- a/DashAI/back/tasks/tabular_classification_task.py
+++ b/DashAI/back/tasks/tabular_classification_task.py
@@ -1,3 +1,5 @@
+from typing import List
+
 from datasets import ClassLabel, DatasetDict, Value
 
 from DashAI.back.tasks.base_task import BaseTask
@@ -14,21 +16,23 @@ class TabularClassificationTask(BaseTask):
     labels for structured data organized in tabular form (rows and columns).
     Models are trained to learn patterns and relationships in the data, enabling
     accurate classification of new instances."""
-    schema: dict = {
+    metadata: dict = {
         "inputs_types": [ClassLabel, Value],
         "outputs_types": [ClassLabel],
         "inputs_cardinality": "n",
         "outputs_cardinality": 1,
     }
 
-    def prepare_for_task(self, datasetdict: DatasetDict) -> DatasetDict:
+    def prepare_for_task(
+        self, datasetdict: DatasetDict, outputs_columns: List[str]
+    ) -> DatasetDict:
         """Change the column types to suit the tabular classification task.
 
         A copy of the dataset is created.
 
         Parameters
         ----------
-        dataset : DatasetDict
+        datasetdict : DatasetDict
             Dataset to be changed
 
         Returns
@@ -36,8 +40,7 @@ def prepare_for_task(self, datasetdict: DatasetDict) -> DatasetDict:
         DatasetDict
             Dataset with the new types
         """
-        outputs_columns = datasetdict["train"].outputs_columns
-        types = {outputs_columns[0]: "Categorical"}
+        types = {column: "Categorical" for column in outputs_columns}
         for split in datasetdict:
             datasetdict[split] = datasetdict[split].change_columns_type(types)
         return datasetdict
diff --git a/DashAI/back/tasks/text_classification_task.py b/DashAI/back/tasks/text_classification_task.py
index 2ea11f09..424e4b26 100644
--- a/DashAI/back/tasks/text_classification_task.py
+++ b/DashAI/back/tasks/text_classification_task.py
@@ -1,3 +1,5 @@
+from typing import List
+
 from datasets import ClassLabel, DatasetDict, Value
 
 from DashAI.back.tasks.base_task import BaseTask
@@ -6,7 +8,7 @@
 class TextClassificationTask(BaseTask):
     """Base class for Text Classification Task."""
 
-    schema: dict = {
+    metadata: dict = {
         "inputs_types": [Value],
         "outputs_types": [ClassLabel],
         "inputs_cardinality": 1,
@@ -20,23 +22,25 @@ class TextClassificationTask(BaseTask):
     analysis, spam filtering, topic classification, and document categorization.
     """
 
-    def prepare_for_task(self, datasetdict: DatasetDict) -> DatasetDict:
+    def prepare_for_task(
+        self, datasetdict: DatasetDict, outputs_columns: List[str]
+    ) -> DatasetDict:
         """Change the column types to suit the tabular classification task.
 
-        Note: A copy of the dataset is created.
+        A copy of the dataset is created.
 
         Parameters
         ----------
-        dataset : DatasetDict
-            Dataset to be configured
+        datasetdict : DatasetDict
+            Dataset to be changed
 
         Returns
         -------
         DatasetDict
             Dataset with the new types
         """
-        outputs_columns = datasetdict["train"].outputs_columns
-        types = {outputs_columns[0]: "Categorical"}
+        types = {column: "Categorical" for column in outputs_columns}
+
         for split in datasetdict:
             datasetdict[split] = datasetdict[split].change_columns_type(types)
         return datasetdict
diff --git a/DashAI/back/tasks/translation_task.py b/DashAI/back/tasks/translation_task.py
index 69e548a4..671a28a1 100644
--- a/DashAI/back/tasks/translation_task.py
+++ b/DashAI/back/tasks/translation_task.py
@@ -1,4 +1,6 @@
 """DashAI Translation Task."""
+from typing import List
+
 from datasets import DatasetDict, Sequence, Value
 
 from DashAI.back.tasks.base_task import BaseTask
@@ -7,7 +9,7 @@
 class TranslationTask(BaseTask):
     """Base class for translation task."""
 
-    schema: dict = {
+    metadata: dict = {
         "inputs_types": [Value, Sequence],
         "outputs_types": [Value, Sequence],
         "inputs_cardinality": 1,
@@ -20,7 +22,9 @@ class TranslationTask(BaseTask):
     preserving the meaning and context.
     """
 
-    def prepare_for_task(self, datasetdict: DatasetDict):
+    def prepare_for_task(
+        self, datasetdict: DatasetDict, outputs_columns: List[str]
+    ) -> DatasetDict:
         """Change the column types to suit the tabular classification task.
 
         A copy of the dataset is created.
diff --git a/DashAI/front/src/api/datasets.ts b/DashAI/front/src/api/datasets.ts
index 30b12605..d6ae60d2 100644
--- a/DashAI/front/src/api/datasets.ts
+++ b/DashAI/front/src/api/datasets.ts
@@ -13,24 +13,26 @@ export const getDatasets = async (): Promise<IDataset[]> => {
   return response.data;
 };
 
-export const updateDataset = async (
-  id: number,
-  name: string | undefined = undefined,
-  taskName: string | undefined = undefined,
-): Promise<IDataset> => {
-  let params = {};
+export const getDatasetSample = async (id: number): Promise<object> => {
+  const response = await api.get<object>(`${datasetEndpoint}/${id}/sample`);
+  return response.data;
+};
 
-  if (name !== undefined) {
-    params = { ...params, name };
-  }
+export const getDatasetTypes = async (id: number): Promise<object> => {
+  const response = await api.get<object>(`${datasetEndpoint}/${id}/types`);
+  return response.data;
+};
 
-  if (taskName !== undefined) {
-    params = { ...params, task_name: taskName };
-  }
+export const getDatasetInfo = async (id: number): Promise<object> => {
+  const response = await api.get<object>(`${datasetEndpoint}/${id}/info`);
+  return response.data;
+};
 
-  const response = await api.patch<IDataset>(`${datasetEndpoint}/${id}`, null, {
-    params,
-  });
+export const updateDataset = async (
+  id: number,
+  formData: object,
+): Promise<IDataset> => {
+  const response = await api.patch(`${datasetEndpoint}/${id}`, { ...formData });
   return response.data;
 };
 
diff --git a/DashAI/front/src/api/experiment.ts b/DashAI/front/src/api/experiment.ts
index e152a23f..1a559f3d 100644
--- a/DashAI/front/src/api/experiment.ts
+++ b/DashAI/front/src/api/experiment.ts
@@ -16,12 +16,18 @@ export const getExperimentById = async (id: string): Promise<IExperiment> => {
 export const createExperiment = async (
   datasetId: number,
   taskName: string,
-  name: string,
+  expName: string,
+  inputColumns: number[],
+  outputColumns: number[],
+  splitsValue: JSON,
 ): Promise<IExperiment> => {
   const data = {
     dataset_id: datasetId,
     task_name: taskName,
-    name,
+    name: expName,
+    input_columns: inputColumns,
+    output_columns: outputColumns,
+    splits: splitsValue,
   };
 
   const response = await api.post<IExperiment>("/v1/experiment/", data);
@@ -43,3 +49,22 @@ export const deleteExperiment = async (id: string): Promise<object> => {
   const response = await api.delete(`/v1/experiment/${id}`);
   return response.data;
 };
+
+export const validateColumns = async (
+  taskName: string,
+  datasetId: number,
+  inputColumns: number[],
+  outputColumns: number[],
+): Promise<object> => {
+  const formData = {
+    task_name: taskName,
+    dataset_id: datasetId,
+    inputs_columns: inputColumns,
+    outputs_columns: outputColumns,
+  };
+  const response = await api.post<object>(
+    "/v1/experiment/validation",
+    formData,
+  );
+  return response.data;
+};
diff --git a/DashAI/front/src/components/custom/SelectTypeCell.jsx b/DashAI/front/src/components/custom/SelectTypeCell.jsx
new file mode 100644
index 00000000..ffadb69f
--- /dev/null
+++ b/DashAI/front/src/components/custom/SelectTypeCell.jsx
@@ -0,0 +1,37 @@
+import React, { useState } from "react";
+import { Select } from "@mui/material";
+import PropTypes, { string } from "prop-types";
+const SelectTypeCell = ({ id, value, field, options, updateValue }) => {
+  const [selectedValue, setSelectedValue] = useState(value || "");
+
+  const handleChange = async (event) => {
+    const newValue = event.target.value;
+    setSelectedValue(newValue);
+    updateValue(id, field, newValue);
+  };
+
+  return (
+    <Select
+      native
+      value={selectedValue}
+      onChange={handleChange}
+      size="small"
+      sx={{ height: 1 }}
+      autoFocus
+    >
+      {options.map((option) => (
+        <option key={option} value={option}>
+          {option}
+        </option>
+      ))}
+    </Select>
+  );
+};
+SelectTypeCell.propTypes = {
+  id: PropTypes.oneOfType([PropTypes.string, PropTypes.number]).isRequired,
+  value: PropTypes.string.isRequired,
+  field: PropTypes.string.isRequired,
+  options: PropTypes.arrayOf(string).isRequired,
+  updateValue: PropTypes.func.isRequired,
+};
+export default SelectTypeCell;
diff --git a/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx b/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx
index 63e46a59..34937383 100644
--- a/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx
+++ b/DashAI/front/src/components/datasets/ConfigureAndUploadDataset.jsx
@@ -94,7 +94,6 @@ function ConfigureAndUploadDataset({
 
 ConfigureAndUploadDataset.propTypes = {
   newDataset: PropTypes.shape({
-    task_name: PropTypes.string,
     dataloader: PropTypes.string,
     file: PropTypes.oneOfType([
       PropTypes.instanceOf(File),
diff --git a/DashAI/front/src/components/datasets/DataloaderConfiguration.jsx b/DashAI/front/src/components/datasets/DataloaderConfiguration.jsx
index f6282cbb..4b7998af 100644
--- a/DashAI/front/src/components/datasets/DataloaderConfiguration.jsx
+++ b/DashAI/front/src/components/datasets/DataloaderConfiguration.jsx
@@ -128,7 +128,6 @@ DataloaderConfiguration.propTypes = {
   ).isRequired,
   onSubmit: PropTypes.func.isRequired,
   newDataset: PropTypes.shape({
-    task_name: PropTypes.string,
     dataloader_name: PropTypes.string,
   }).isRequired,
   formSubmitRef: PropTypes.shape({ current: PropTypes.any }).isRequired,
diff --git a/DashAI/front/src/components/datasets/DatasetModal.jsx b/DashAI/front/src/components/datasets/DatasetModal.jsx
index 02895c3c..07edd429 100644
--- a/DashAI/front/src/components/datasets/DatasetModal.jsx
+++ b/DashAI/front/src/components/datasets/DatasetModal.jsx
@@ -13,20 +13,22 @@ import {
   Typography,
   StepButton,
 } from "@mui/material";
-import SelectTaskStep from "./SelectTaskStep";
 import SelectDataloaderStep from "./SelectDataloaderStep";
 import ConfigureAndUploadDataset from "./ConfigureAndUploadDataset";
 import { useSnackbar } from "notistack";
-import { uploadDataset as uploadDatasetRequest } from "../../api/datasets";
+import {
+  uploadDataset as uploadDatasetRequest,
+  updateDataset as updateDatasetRequest,
+} from "../../api/datasets";
+import DatasetSummaryStep from "./DatasetSummaryStep";
 
 const steps = [
-  { name: "selectTask", label: "Select Task" },
   { name: "selectDataloader", label: "Select a way to upload" },
   { name: "uploadDataset", label: "Configure and upload your dataset" },
+  { name: "datasetSummary", label: "Dataset summary" },
 ];
 
 const defaultNewDataset = {
-  task_name: "",
   dataloader: "",
   file: null,
   url: "",
@@ -44,6 +46,10 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
   const [nextEnabled, setNextEnabled] = useState(false);
   const [newDataset, setNewDataset] = useState(defaultNewDataset);
   const [readyToUpload, setReadyToUpload] = useState(false);
+  const [uploaded, setUploaded] = useState(false);
+  const [requestError, setRequestError] = useState(false);
+  const [uploadedDataset, setUploadedDataset] = useState([]);
+  const [columnsSpec, setColumnsSpec] = useState({});
   const formSubmitRef = useRef(null);
   const { enqueueSnackbar } = useSnackbar();
 
@@ -62,18 +68,40 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
       );
       formData.append("url", ""); // TODO: url handling
       formData.append("file", newDataset.file);
-      await uploadDatasetRequest(formData);
+      const dataset = await uploadDatasetRequest(formData);
+      setUploadedDataset(dataset);
       enqueueSnackbar("Dataset uploaded successfully", { variant: "success" });
       updateDatasets();
     } catch (error) {
       console.error(error);
+      setRequestError(true);
       enqueueSnackbar("Error when trying to upload the dataset.");
+    } finally {
+      setUploaded(true);
+    }
+  };
+
+  const handleUpdateColumnsSpec = async () => {
+    try {
+      await updateDatasetRequest(uploadedDataset.id, { columns: columnsSpec });
+    } catch (error) {
+      enqueueSnackbar(
+        "Error while trying to update the column and data types.",
+      );
+      if (error.response) {
+        console.error("Response error:", error.message);
+      } else if (error.request) {
+        console.error("Request error", error.request);
+      } else {
+        console.error("Unknown Error", error.message);
+      }
     }
   };
 
   const handleCloseDialog = () => {
     setActiveStep(0);
     setNewDataset(defaultNewDataset);
+    setUploaded(false);
     setNextEnabled(false);
     setOpen(false);
   };
@@ -83,12 +111,15 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
   };
 
   const handleNextButton = () => {
+    if (activeStep === 1 && !uploaded) {
+      formSubmitRef.current.handleSubmit();
+    }
     if (activeStep < steps.length - 1) {
       setActiveStep(activeStep + 1);
       setNextEnabled(false);
     } else {
-      // trigger dataloader form submit
-      formSubmitRef.current.handleSubmit();
+      handleUpdateColumnsSpec(); // TODO: update only if the Columns spec have changed
+      handleCloseDialog();
     }
   };
 
@@ -108,10 +139,17 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
       readyToUpload
     ) {
       handleSubmitNewDataset();
-      handleCloseDialog();
+      setReadyToUpload(false);
     }
   }, [newDataset]);
 
+  useEffect(() => {
+    if (requestError) {
+      setActiveStep(1);
+      setNextEnabled(false);
+      setRequestError(false);
+    }
+  }, [requestError]);
   return (
     <Dialog
       open={open}
@@ -159,29 +197,31 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
 
       {/* Main content - steps */}
       <DialogContent dividers>
-        {/* Step 1: select task */}
+        {/* Step 1: select dataloader */}
         {activeStep === 0 && (
-          <SelectTaskStep
+          <SelectDataloaderStep
             newDataset={newDataset}
             setNewDataset={setNewDataset}
             setNextEnabled={setNextEnabled}
           />
         )}
-        {/* Step 2: select dataloader */}
+        {/* Step 2: Configure dataloader and upload file */}
         {activeStep === 1 && (
-          <SelectDataloaderStep
+          <ConfigureAndUploadDataset
             newDataset={newDataset}
             setNewDataset={setNewDataset}
             setNextEnabled={setNextEnabled}
+            formSubmitRef={formSubmitRef}
           />
         )}
-        {/* Step 3: Configure dataloader and upload file */}
+        {/* Step 3: Dataset Summary and cast columns types */}
         {activeStep === 2 && (
-          <ConfigureAndUploadDataset
-            newDataset={newDataset}
-            setNewDataset={setNewDataset}
+          <DatasetSummaryStep
+            datasetId={uploadedDataset.id}
             setNextEnabled={setNextEnabled}
-            formSubmitRef={formSubmitRef}
+            datasetUploaded={uploaded}
+            columnsSpec={columnsSpec}
+            setColumnsSpec={setColumnsSpec}
           />
         )}
       </DialogContent>
@@ -194,7 +234,7 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
           </Button>
           <Button
             onClick={() => {
-              if (activeStep === 2) {
+              if (activeStep === 1) {
                 setReadyToUpload(true);
               }
               handleNextButton();
diff --git a/DashAI/front/src/components/datasets/DatasetSummaryModal.jsx b/DashAI/front/src/components/datasets/DatasetSummaryModal.jsx
new file mode 100644
index 00000000..1d8f6784
--- /dev/null
+++ b/DashAI/front/src/components/datasets/DatasetSummaryModal.jsx
@@ -0,0 +1,71 @@
+import React, { useState } from "react";
+import PropTypes from "prop-types";
+import { GridActionsCellItem } from "@mui/x-data-grid";
+import { Search } from "@mui/icons-material";
+import {
+  Dialog,
+  DialogContent,
+  DialogTitle,
+  Grid,
+  Button,
+  DialogActions,
+} from "@mui/material";
+import DatasetSummaryTable from "./DatasetSummaryTable";
+
+function DatasetSummaryModal({ datasetId }) {
+  const [open, setOpen] = useState(false);
+
+  const handleCloseContent = () => {
+    setOpen(false);
+  };
+  return (
+    <React.Fragment>
+      <GridActionsCellItem
+        key="dataset-summary-button"
+        icon={<Search />}
+        label="Dataset Summary"
+        onClick={() => setOpen(true)}
+        sx={{ color: "warning.main" }}
+      />
+      <Dialog
+        open={open}
+        onClose={() => setOpen(false)}
+        fullWidth
+        maxWidth={"md"}
+      >
+        <DialogTitle>Dataset Summary</DialogTitle>
+        <DialogContent>
+          <Grid
+            container
+            direction="row"
+            justifyContent="space-around"
+            alignItems="stretch"
+            spacing={2}
+            onClick={(event) => event.stopPropagation()}
+          >
+            {/* Dataset Summary Table */}
+            <DatasetSummaryTable datasetId={datasetId} />
+          </Grid>
+        </DialogContent>
+        {/* Actions - Close */}
+        <DialogActions>
+          <Button
+            onClick={handleCloseContent}
+            autoFocus
+            variant="contained"
+            color="primary"
+            disabled={false}
+          >
+            Close
+          </Button>
+        </DialogActions>
+      </Dialog>
+    </React.Fragment>
+  );
+}
+
+DatasetSummaryModal.propTypes = {
+  datasetId: PropTypes.number.isRequired,
+};
+
+export default DatasetSummaryModal;
diff --git a/DashAI/front/src/components/datasets/DatasetSummaryStep.jsx b/DashAI/front/src/components/datasets/DatasetSummaryStep.jsx
new file mode 100644
index 00000000..4de9fe5b
--- /dev/null
+++ b/DashAI/front/src/components/datasets/DatasetSummaryStep.jsx
@@ -0,0 +1,58 @@
+import React, { useEffect } from "react";
+import { Paper, Grid, Typography, CircularProgress } from "@mui/material";
+import PropTypes from "prop-types";
+import DatasetSummaryTable from "./DatasetSummaryTable";
+function DatasetSummaryStep({
+  datasetId,
+  setNextEnabled,
+  datasetUploaded,
+  columnsSpec,
+  setColumnsSpec,
+}) {
+  useEffect(() => {
+    if (datasetUploaded) {
+      setNextEnabled(true);
+    }
+  }, [datasetUploaded]);
+  return (
+    <Paper
+      variant="outlined"
+      sx={{ p: 4, maxHeight: "55vh", overflow: "auto" }}
+    >
+      <Grid container direction={"column"} alignItems={"center"}>
+        <Grid item>
+          <Typography variant="subtitle1">Dataset Summary</Typography>
+          <Typography
+            item
+            variant="caption"
+            component="h3"
+            sx={{ mb: 2, color: "grey" }}
+          >
+            Summary of the recently uploaded dataset with predefined column
+            types. You can modify the type by selecting a different value.
+          </Typography>
+        </Grid>
+        <Grid item>
+          {datasetUploaded ? (
+            <DatasetSummaryTable
+              datasetId={datasetId}
+              isEditable={true}
+              columnsSpec={columnsSpec}
+              setColumnsSpec={setColumnsSpec}
+            />
+          ) : (
+            <CircularProgress />
+          )}
+        </Grid>
+      </Grid>
+    </Paper>
+  );
+}
+DatasetSummaryStep.propTypes = {
+  datasetId: PropTypes.number,
+  setNextEnabled: PropTypes.func.isRequired,
+  datasetUploaded: PropTypes.bool,
+  columnsSpec: PropTypes.object.isRequired,
+  setColumnsSpec: PropTypes.func.isRequired,
+};
+export default DatasetSummaryStep;
diff --git a/DashAI/front/src/components/datasets/DatasetSummaryTable.jsx b/DashAI/front/src/components/datasets/DatasetSummaryTable.jsx
new file mode 100644
index 00000000..b1479bea
--- /dev/null
+++ b/DashAI/front/src/components/datasets/DatasetSummaryTable.jsx
@@ -0,0 +1,153 @@
+// DatasetSummaryTable.js
+import React, { useEffect, useMemo, useState } from "react";
+import { useSnackbar } from "notistack";
+import PropTypes from "prop-types";
+import { DataGrid, useGridApiRef } from "@mui/x-data-grid";
+import {
+  getDatasetSample as getDatasetSampleRequest,
+  getDatasetTypes as getDatasetTypesRequest,
+} from "../../api/datasets";
+import { dataTypesList, columnTypesList } from "../../utils/typesLists";
+import SelectTypeCell from "../custom/SelectTypeCell";
+
+function DatasetSummaryTable({
+  datasetId,
+  isEditable,
+  columnsSpec,
+  setColumnsSpec,
+}) {
+  const [loading, setLoading] = useState(true);
+  const { enqueueSnackbar } = useSnackbar();
+  const [rows, setRows] = useState([]);
+  const apiRef = useGridApiRef();
+
+  const getDatasetInfo = async () => {
+    setLoading(true);
+    try {
+      const dataset = await getDatasetSampleRequest(datasetId);
+      const types = await getDatasetTypesRequest(datasetId);
+      const rowsArray = Object.keys(dataset).map((name, idx) => {
+        return {
+          id: idx,
+          columnName: name,
+          example: dataset[name][0],
+          columnType: types[name].type,
+          dataType: types[name].dtype,
+        };
+      });
+      setRows(rowsArray);
+      if (isEditable) {
+        setColumnsSpec(types);
+      }
+    } catch (error) {
+      enqueueSnackbar("Error while trying to obtain the dataset.");
+      if (error.response) {
+        console.error("Response error:", error.message);
+      } else if (error.request) {
+        console.error("Request error", error.request);
+      } else {
+        console.error("Unknown Error", error.message);
+      }
+    } finally {
+      setLoading(false);
+    }
+  };
+
+  const updateCellValue = async (id, field, newValue) => {
+    await apiRef.current.setEditCellValue({ id, field, value: newValue });
+    apiRef.current.stopCellEditMode({ id, field });
+    setRows((prevRows) =>
+      prevRows.map((row) =>
+        row.id === id ? { ...row, [field]: newValue } : row,
+      ),
+    );
+
+    const columnName = rows.find((row) => row.id === id)?.columnName;
+    const updateColumns = { ...columnsSpec };
+
+    if (field === "dataType") {
+      updateColumns[columnName].dtype = newValue;
+    } else if (field === "columnType") {
+      updateColumns[columnName].type = newValue;
+    }
+
+    setColumnsSpec(updateColumns);
+  };
+
+  const renderSelectCell = (params, options) => {
+    return (
+      <SelectTypeCell
+        id={params.id}
+        value={params.value}
+        field={params.field}
+        options={options}
+        updateValue={(id, field, newValue) =>
+          updateCellValue(id, field, newValue)
+        }
+      />
+    );
+  };
+
+  const columns = useMemo(() => [
+    {
+      field: "columnName",
+      headerName: "Column name",
+      minWidth: 200,
+      editable: false,
+    },
+    {
+      field: "example",
+      headerName: "Example",
+      minWidth: 200,
+      editable: false,
+    },
+    {
+      field: "columnType",
+      headerName: "Column type",
+      renderEditCell: (params) =>
+        isEditable && renderSelectCell(params, columnTypesList),
+      minWidth: 200,
+      editable: isEditable,
+    },
+    {
+      field: "dataType",
+      headerName: "Data type",
+      renderEditCell: (params) =>
+        isEditable && renderSelectCell(params, dataTypesList),
+      minWidth: 200,
+      editable: isEditable,
+    },
+  ]);
+
+  useEffect(() => {
+    getDatasetInfo();
+  }, []);
+
+  return (
+    <DataGrid
+      rows={rows}
+      columns={columns}
+      initialState={{
+        pagination: {
+          paginationModel: {
+            pageSize: 4,
+          },
+        },
+      }}
+      pageSize={4}
+      pageSizeOptions={[4, 5, 10]}
+      loading={loading}
+      apiRef={apiRef}
+      autoHeight
+    />
+  );
+}
+
+DatasetSummaryTable.propTypes = {
+  datasetId: PropTypes.number,
+  isEditable: PropTypes.bool,
+  columnsSpec: PropTypes.object,
+  setColumnsSpec: PropTypes.func,
+};
+
+export default DatasetSummaryTable;
diff --git a/DashAI/front/src/components/datasets/DatasetsTable.jsx b/DashAI/front/src/components/datasets/DatasetsTable.jsx
index a352278b..42d79a46 100644
--- a/DashAI/front/src/components/datasets/DatasetsTable.jsx
+++ b/DashAI/front/src/components/datasets/DatasetsTable.jsx
@@ -8,6 +8,7 @@ import {
 import { Button, Grid, Paper, Typography } from "@mui/material";
 import DeleteItemModal from "../custom/DeleteItemModal";
 import EditDatasetModal from "./EditDatasetModal";
+import DatasetSummaryModal from "./DatasetSummaryModal";
 import {
   getDatasets as getDatasetsRequest,
   deleteDataset as deleteDatasetRequest,
@@ -96,12 +97,6 @@ function DatasetsTable({
         minWidth: 250,
         editable: false,
       },
-      {
-        field: "task_name",
-        headerName: "Task",
-        minWidth: 200,
-        editable: false,
-      },
       {
         field: "created",
         headerName: "Created",
@@ -125,7 +120,6 @@ function DatasetsTable({
           <EditDatasetModal
             key="edit-component"
             name={params.row.name}
-            taskName={params.row.task_name}
             datasetId={params.id}
             updateDatasets={() => setUpdateTableFlag(true)}
           />,
@@ -133,6 +127,10 @@ function DatasetsTable({
             key="delete-component"
             deleteFromTable={createDeleteHandler(params.id)}
           />,
+          <DatasetSummaryModal
+            key="dataset-summary-component"
+            datasetId={params.id}
+          />,
         ],
       },
     ],
diff --git a/DashAI/front/src/components/datasets/EditDatasetModal.jsx b/DashAI/front/src/components/datasets/EditDatasetModal.jsx
index 314b405d..19cca357 100644
--- a/DashAI/front/src/components/datasets/EditDatasetModal.jsx
+++ b/DashAI/front/src/components/datasets/EditDatasetModal.jsx
@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from "react";
+import React, { useState } from "react";
 import PropTypes from "prop-types";
 import { GridActionsCellItem } from "@mui/x-data-grid";
 import EditIcon from "@mui/icons-material/Edit";
@@ -12,22 +12,17 @@ import {
   TextField,
   Typography,
 } from "@mui/material";
-import ItemSelectorWithInfo from "../custom/ItemSelectorWithInfo";
 import { updateDataset as updateDatasetRequest } from "../../api/datasets";
-import { getComponents as getComponentsRequest } from "../../api/component";
 import { useSnackbar } from "notistack";
 
-function EditDatasetModal({ datasetId, name, taskName, updateDatasets }) {
+function EditDatasetModal({ datasetId, name, updateDatasets }) {
   const { enqueueSnackbar } = useSnackbar();
   const [open, setOpen] = useState(false);
   const [datasetName, setDatasetName] = useState(name);
-  const [tasks, setTasks] = useState([]);
-  const [selectedTask, setSelectedTask] = useState({});
-  const [loading, setLoading] = useState(true);
 
   const editDataset = async () => {
     try {
-      await updateDatasetRequest(datasetId, datasetName);
+      await updateDatasetRequest(datasetId, { name: datasetName });
       enqueueSnackbar("Dataset updated successfully", {
         variant: "success",
       });
@@ -43,39 +38,12 @@ function EditDatasetModal({ datasetId, name, taskName, updateDatasets }) {
     }
   };
 
-  const getTasks = async () => {
-    setLoading(true);
-    try {
-      const tasks = await getComponentsRequest({ selectTypes: ["Task"] });
-      setTasks(tasks);
-      const previouslySelectedTask = tasks.find(
-        (task) => task.name === taskName,
-      );
-      setSelectedTask(previouslySelectedTask);
-    } catch (error) {
-      enqueueSnackbar("Error while trying to obtain available tasks");
-      if (error.response) {
-        console.error("Response error:", error.message);
-      } else if (error.request) {
-        console.error("Request error", error.request);
-      } else {
-        console.error("Unknown Error", error.message);
-      }
-    } finally {
-      setLoading(false);
-    }
-  };
-
   const handleSaveConfig = () => {
     editDataset();
     updateDatasets();
     setOpen(false);
   };
 
-  // fetch available tasks
-  useEffect(() => {
-    getTasks();
-  }, []);
   return (
     <React.Fragment>
       <GridActionsCellItem
@@ -103,7 +71,7 @@ function EditDatasetModal({ datasetId, name, taskName, updateDatasets }) {
             {/* New name field */}
             <Grid item xs={12}>
               <Typography variant="subtitle1" component="h3" sx={{ mb: 3 }}>
-                Enter a new name and select a new task for your dataset
+                Enter a new name for your dataset
               </Typography>
 
               <TextField
@@ -116,18 +84,6 @@ function EditDatasetModal({ datasetId, name, taskName, updateDatasets }) {
                 sx={{ mb: 2 }}
               />
             </Grid>
-
-            {/* Select a new task */}
-            <Grid item xs={12}>
-              {!loading && (
-                <ItemSelectorWithInfo
-                  itemsList={tasks}
-                  selectedItem={selectedTask}
-                  setSelectedItem={setSelectedTask}
-                  disabled
-                />
-              )}
-            </Grid>
           </Grid>
         </DialogContent>
 
@@ -151,7 +107,6 @@ function EditDatasetModal({ datasetId, name, taskName, updateDatasets }) {
 EditDatasetModal.propTypes = {
   datasetId: PropTypes.number.isRequired,
   name: PropTypes.string.isRequired,
-  taskName: PropTypes.string.isRequired,
   updateDatasets: PropTypes.func.isRequired,
 };
 
diff --git a/DashAI/front/src/components/datasets/SelectDataloaderStep.jsx b/DashAI/front/src/components/datasets/SelectDataloaderStep.jsx
index ccdc3716..79f588b2 100644
--- a/DashAI/front/src/components/datasets/SelectDataloaderStep.jsx
+++ b/DashAI/front/src/components/datasets/SelectDataloaderStep.jsx
@@ -23,7 +23,6 @@ function SelectDataloaderStep({ newDataset, setNewDataset, setNextEnabled }) {
     try {
       const dataloaders = await getComponentsRequest({
         selectTypes: ["DataLoader"],
-        relatedComponent: newDataset.task_name,
       });
       setDataloaders(dataloaders);
       if (newDataset.dataloader !== "") {
@@ -91,7 +90,6 @@ function SelectDataloaderStep({ newDataset, setNewDataset, setNextEnabled }) {
 
 SelectDataloaderStep.propTypes = {
   newDataset: PropTypes.shape({
-    task_name: PropTypes.string,
     dataloader: PropTypes.string,
   }).isRequired,
   setNewDataset: PropTypes.func.isRequired,
diff --git a/DashAI/front/src/components/datasets/SelectTaskStep.jsx b/DashAI/front/src/components/datasets/SelectTaskStep.jsx
deleted file mode 100644
index 0a250f85..00000000
--- a/DashAI/front/src/components/datasets/SelectTaskStep.jsx
+++ /dev/null
@@ -1,96 +0,0 @@
-import React, { useEffect, useState } from "react";
-import PropTypes from "prop-types";
-import { useSnackbar } from "notistack";
-import { getComponents as getComponentsRequest } from "../../api/component";
-import ItemSelectorWithInfo from "../custom/ItemSelectorWithInfo";
-import { DialogContentText, Grid } from "@mui/material";
-
-/**
- * This component renders a list of tasks and allows the user to select one.
- * @param {object} newDataset An object that stores all the important states for the dataset modal.
- * @param {function} setNewDataset function that modifies newDataset state
- * @param {function} setNextEnabled function to enable or disable the "Next" button in the dataset modal.
- */
-function SelectTaskStep({ newDataset, setNewDataset, setNextEnabled }) {
-  const { enqueueSnackbar } = useSnackbar();
-
-  const [tasks, setTasks] = useState([]);
-  const [selectedTask, setSelectedTask] = useState({});
-  const [loading, setLoading] = useState(true);
-
-  const getTasks = async () => {
-    setLoading(true);
-    try {
-      const tasks = await getComponentsRequest({ selectTypes: ["Task"] });
-      setTasks(tasks);
-      // If there was a previously selected task, it is used as the initial value for the task.
-      if (newDataset.task_name !== "") {
-        const previouslySelectedTask =
-          tasks.find((task) => task.name === newDataset.task_name) || {};
-        setSelectedTask(previouslySelectedTask);
-      }
-    } catch (error) {
-      enqueueSnackbar("Error while trying to obtain available tasks");
-      if (error.response) {
-        console.error("Response error:", error.message);
-      } else if (error.request) {
-        console.error("Request error", error.request);
-      } else {
-        console.error("Unknown Error", error.message);
-      }
-    } finally {
-      setLoading(false);
-    }
-  };
-
-  // updates the modal state with the name of the task that is selected by the user
-  useEffect(() => {
-    if (selectedTask && Object.keys(selectedTask).length === 0) {
-      setNewDataset({ ...newDataset, task_name: "" });
-    } else if (selectedTask && "name" in selectedTask) {
-      setNewDataset({ ...newDataset, task_name: selectedTask.name });
-      setNextEnabled(true);
-    }
-  }, [selectedTask]);
-
-  // fetches the available tasks
-  useEffect(() => {
-    getTasks();
-  }, []);
-
-  return (
-    <Grid
-      container
-      direction="column"
-      justifyContent="space-around"
-      alignItems="stretch"
-      spacing={2}
-    >
-      {/* Title */}
-      <Grid item>
-        <DialogContentText sx={{ mb: 3 }}>{`Select a task`}</DialogContentText>
-      </Grid>
-      {/* List of tasks  */}
-      <Grid item>
-        {!loading && (
-          <ItemSelectorWithInfo
-            itemsList={tasks}
-            selectedItem={selectedTask}
-            setSelectedItem={setSelectedTask}
-          />
-        )}
-      </Grid>
-    </Grid>
-  );
-}
-
-SelectTaskStep.propTypes = {
-  newDataset: PropTypes.shape({
-    task_name: PropTypes.string,
-    dataloader_name: PropTypes.string,
-  }).isRequired,
-  setNewDataset: PropTypes.func.isRequired,
-  setNextEnabled: PropTypes.func.isRequired,
-};
-
-export default SelectTaskStep;
diff --git a/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx b/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx
index c9595ad5..02333d98 100644
--- a/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx
+++ b/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx
@@ -157,6 +157,15 @@ ConfigureModelsStep.propTypes = {
     name: PropTypes.string,
     dataset: PropTypes.object,
     task_name: PropTypes.string,
+    input_columns: PropTypes.arrayOf(PropTypes.number),
+    output_columns: PropTypes.arrayOf(PropTypes.number),
+    splits: PropTypes.shape({
+      has_changed: PropTypes.bool,
+      is_random: PropTypes.bool,
+      training: PropTypes.number,
+      validation: PropTypes.number,
+      testing: PropTypes.number,
+    }),
     step: PropTypes.string,
     created: PropTypes.instanceOf(Date),
     last_modified: PropTypes.instanceOf(Date),
diff --git a/DashAI/front/src/components/experiments/DivideDatasetColumns.jsx b/DashAI/front/src/components/experiments/DivideDatasetColumns.jsx
new file mode 100644
index 00000000..4883c24b
--- /dev/null
+++ b/DashAI/front/src/components/experiments/DivideDatasetColumns.jsx
@@ -0,0 +1,132 @@
+import React, { useEffect, useState } from "react";
+import PropTypes from "prop-types";
+import { parseRangeToIndex } from "../../utils/parseRange";
+import { Grid, TextField, Typography } from "@mui/material";
+
+function DivideDatasetColumns({
+  datasetInfo,
+  inputColumns,
+  setInputColumns,
+  outputColumns,
+  setOutputColumns,
+  setColumnsReady,
+}) {
+  // TODO: column and row numbers should be minor to the maximum on the dataset
+  const totalColumns = datasetInfo.total_columns;
+
+  // handle columns numbers change state
+  const [parseInputColumnsError, setParseInputColumnsError] = useState(false);
+  const [parseOutputColumnsError, setParseOutputColumnsError] = useState(false);
+  const [parseInputColumnsErrorText, setParseInputColumnsErrorText] =
+    useState("");
+  const [parseOutputColumnsErrorText, setParseOutputColumnsErrorText] =
+    useState("");
+
+  const handleInputColumnsChange = (event) => {
+    const input = event.target.value.replace(/ /g, ""); // TODO: dont accept spaces between numbers
+    try {
+      const columnIndex = parseRangeToIndex(input, totalColumns);
+      setParseInputColumnsError(false);
+      setInputColumns(columnIndex);
+    } catch (error) {
+      setParseInputColumnsErrorText(error.message);
+      setParseInputColumnsError(true);
+    }
+  };
+  const handleOutputColumnsChange = (event) => {
+    const input = event.target.value.replace(/ /g, "");
+    try {
+      const columnIndex = parseRangeToIndex(input, totalColumns);
+      setParseOutputColumnsError(false);
+      setOutputColumns(columnIndex);
+    } catch (error) {
+      setParseOutputColumnsErrorText(error.message);
+      setParseOutputColumnsError(true);
+    }
+  };
+  useEffect(() => {
+    // check if input and output columns are not empty
+    if (
+      !parseInputColumnsError &&
+      inputColumns.length >= 1 &&
+      !parseOutputColumnsError &&
+      outputColumns.length >= 1
+    ) {
+      setColumnsReady(true);
+    } else {
+      setColumnsReady(false);
+    }
+  }, [
+    inputColumns,
+    outputColumns,
+    parseInputColumnsError,
+    parseOutputColumnsError,
+  ]);
+
+  useEffect(() => {
+    // set default input and output columns
+    setInputColumns(parseRangeToIndex(`1-${totalColumns - 1}`, totalColumns));
+    setOutputColumns(parseRangeToIndex(`${totalColumns}`, totalColumns));
+  }, []);
+  return (
+    <React.Fragment>
+      <Grid item xs={12}>
+        <Typography item variant="subtitle1" component="h3" sx={{ mb: 0 }}>
+          Indicate which columns of the dataset will be used as input and
+          output.
+        </Typography>
+      </Grid>
+      <Grid item xs={12}>
+        <Typography
+          item
+          variant="caption"
+          component="h3"
+          sx={{ mb: 2, color: "grey" }}
+        >
+          The notation is based on ranges. For example: 1-6, 23-108
+        </Typography>
+      </Grid>
+
+      <TextField
+        required
+        id="dataset-input-columns"
+        label="Input"
+        fullWidth
+        autoComplete="off"
+        onChange={handleInputColumnsChange}
+        error={parseInputColumnsError}
+        helperText={parseInputColumnsError ? parseInputColumnsErrorText : ""}
+        defaultValue={`1-${totalColumns - 1}`}
+        sx={{ mb: 2 }}
+      />
+      <TextField
+        required
+        id="dataset-output-columns"
+        label="Output"
+        fullWidth
+        autoComplete="off"
+        onChange={handleOutputColumnsChange}
+        error={parseOutputColumnsError}
+        helperText={parseOutputColumnsError ? parseOutputColumnsErrorText : ""}
+        defaultValue={totalColumns}
+        sx={{ mb: 2 }}
+      />
+    </React.Fragment>
+  );
+}
+
+DivideDatasetColumns.propTypes = {
+  datasetInfo: PropTypes.shape({
+    test_size: PropTypes.number,
+    total_columns: PropTypes.number,
+    total_rows: PropTypes.number,
+    train_size: PropTypes.number,
+    val_size: PropTypes.number,
+  }),
+  inputColumns: PropTypes.arrayOf(PropTypes.number),
+  setInputColumns: PropTypes.func.isRequired,
+  outputColumns: PropTypes.arrayOf(PropTypes.number),
+  setOutputColumns: PropTypes.func.isRequired,
+  setColumnsReady: PropTypes.func.isRequired,
+};
+export default DivideDatasetColumns;
diff --git a/DashAI/front/src/components/experiments/ModelsTable.jsx b/DashAI/front/src/components/experiments/ModelsTable.jsx
index adf06d96..3e7b789e 100644
--- a/DashAI/front/src/components/experiments/ModelsTable.jsx
+++ b/DashAI/front/src/components/experiments/ModelsTable.jsx
@@ -109,6 +109,13 @@ ModelsTable.propTypes = {
     name: PropTypes.string,
     dataset: PropTypes.object,
     task_name: PropTypes.string,
+    input_columns: PropTypes.arrayOf(PropTypes.number),
+    output_columns: PropTypes.arrayOf(PropTypes.number),
+    splits: PropTypes.shape({
+      training: PropTypes.number,
+      validation: PropTypes.number,
+      testing: PropTypes.number,
+    }),
     step: PropTypes.string,
     created: PropTypes.instanceOf(Date),
     last_modified: PropTypes.instanceOf(Date),
diff --git a/DashAI/front/src/components/experiments/NewExperimentModal.jsx b/DashAI/front/src/components/experiments/NewExperimentModal.jsx
index b0bed8cc..623bfd72 100644
--- a/DashAI/front/src/components/experiments/NewExperimentModal.jsx
+++ b/DashAI/front/src/components/experiments/NewExperimentModal.jsx
@@ -24,6 +24,7 @@ import { createRun as createRunRequest } from "../../api/run";
 
 import SetNameAndTaskStep from "./SetNameAndTaskStep";
 import SelectDatasetStep from "./SelectDatasetStep";
+import PrepareDatasetStep from "./PrepareDatasetStep";
 import ConfigureModelsStep from "./ConfigureModelsStep";
 
 import { useSnackbar } from "notistack";
@@ -31,6 +32,7 @@ import { useSnackbar } from "notistack";
 const steps = [
   { name: "selectTask", label: "Set name and task" },
   { name: "selectDataset", label: "Select dataset" },
+  { name: "prepareDataset", label: "Prepare dataset" },
   { name: "configureModels", label: "Configure models" },
 ];
 
@@ -39,6 +41,15 @@ const defaultNewExp = {
   name: "",
   dataset: null,
   task_name: "",
+  input_columns: [],
+  output_columns: [],
+  splits: {
+    has_changed: false,
+    is_random: true,
+    train: 0.6,
+    validation: 0.2,
+    test: 0.2,
+  },
   step: "SET_NAME",
   created: null,
   last_modified: null,
@@ -94,6 +105,9 @@ export default function NewExperimentModal({
         newExp.dataset.id,
         newExp.task_name,
         newExp.name,
+        newExp.input_columns,
+        newExp.output_columns,
+        JSON.stringify(newExp.splits),
       );
       const experimentId = response.id;
       await uploadRuns(experimentId);
@@ -227,6 +241,13 @@ export default function NewExperimentModal({
           />
         )}
         {activeStep === 2 && (
+          <PrepareDatasetStep
+            newExp={newExp}
+            setNewExp={setNewExp}
+            setNextEnabled={setNextEnabled}
+          />
+        )}
+        {activeStep === 3 && (
           <ConfigureModelsStep
             newExp={newExp}
             setNewExp={setNewExp}
@@ -248,7 +269,7 @@ export default function NewExperimentModal({
             color="primary"
             disabled={!nextEnabled}
           >
-            {activeStep === 2 ? "Save" : "Next"}
+            {activeStep === 3 ? "Save" : "Next"}
           </Button>
         </ButtonGroup>
       </DialogActions>
diff --git a/DashAI/front/src/components/experiments/PrepareDatasetStep.jsx b/DashAI/front/src/components/experiments/PrepareDatasetStep.jsx
new file mode 100644
index 00000000..a88ba849
--- /dev/null
+++ b/DashAI/front/src/components/experiments/PrepareDatasetStep.jsx
@@ -0,0 +1,251 @@
+import React, { useEffect, useState } from "react";
+import PropTypes from "prop-types";
+
+import { Grid, CircularProgress, Box, Alert, AlertTitle } from "@mui/material";
+import DivideDatasetColumns from "./DivideDatasetColumns";
+import SplitDatasetRows from "./SplitDatasetRows";
+import { getDatasetInfo as getDatasetInfoRequest } from "../../api/datasets";
+import { getComponents as getComponentsRequest } from "../../api/component";
+import { validateColumns as validateColumnsRequest } from "../../api/experiment";
+import { useSnackbar } from "notistack";
+/**
+ * Step of the experiment modal: Set the input and output columns to use for clasification
+ * and the splits for training, validation and testing
+ * @param {object} newExp object that contains the Experiment Modal state
+ * @param {function} setNewExp updates the Eperimento Modal state (newExp)
+ * @param {function} setNextEnabled function to enable or disable the "Next" button in the modal
+ */
+function PrepareDatasetStep({ newExp, setNewExp, setNextEnabled }) {
+  // dataset info state
+  const [datasetInfo, setDatasetInfo] = useState({});
+  const { enqueueSnackbar } = useSnackbar();
+  const [infoLoading, setInfoLoading] = useState(true);
+
+  // task requirements state
+  const [taskRequirements, setTaskRequirements] = useState({
+    name: "",
+    metadata: {
+      inputs_types: [],
+      inputs_cardinality: "",
+      outputs_types: [],
+      outputs_cardinality: "",
+    },
+  });
+
+  // columns index state
+  const [inputColumns, setInputColumns] = useState([]);
+  const [outputColumns, setOutputColumns] = useState([]);
+  const [columnsReady, setColumnsReady] = useState(true);
+  const [columnsAreValid, setColumnsAreValid] = useState(false);
+
+  // rows index state
+  const defaultParitionsIndex = {
+    train: [],
+    validation: [],
+    test: [],
+  };
+  const defaultPartitionsPercentage = {
+    train: 0.6,
+    validation: 0.2,
+    test: 0.2,
+  };
+
+  const [rowsPartitionsIndex, setRowsPartitionsIndex] = useState(
+    defaultParitionsIndex,
+  );
+  const [rowsPartitionsPercentage, setRowsPartitionsPercentage] = useState(
+    defaultPartitionsPercentage,
+  );
+  const [isRandom, setIsRandom] = useState(true);
+  const [splitsReady, setSplitsReady] = useState(false);
+
+  const getDatasetInfo = async () => {
+    setInfoLoading(true);
+    try {
+      const datasetInfo = await getDatasetInfoRequest(newExp.dataset.id);
+      setDatasetInfo(datasetInfo);
+    } catch (error) {
+      enqueueSnackbar("Error while trying to obtain the dataset info.");
+      if (error.response) {
+        console.error("Response error:", error.message);
+      } else if (error.request) {
+        console.error("Request error", error.request);
+      } else {
+        console.error("Unknown Error", error.message);
+      }
+    } finally {
+      setInfoLoading(false);
+    }
+  };
+
+  const getTaskRequirements = async () => {
+    try {
+      const task = await getComponentsRequest({
+        selectTypes: ["Task"],
+      });
+
+      setTaskRequirements(
+        task.filter((task) => task.name === newExp.task_name)[0],
+      );
+    } catch (error) {
+      enqueueSnackbar("Error while trying to obtain the task requirements.");
+      if (error.response) {
+        console.error("Response error:", error.message);
+      } else if (error.request) {
+        console.error("Request error", error.request);
+      } else {
+        console.error("Unknown Error", error.message);
+      }
+    }
+  };
+
+  const checkIfSplitsHasChanged = () => {
+    if (
+      rowsPartitionsPercentage.train !== defaultPartitionsPercentage.train ||
+      rowsPartitionsPercentage.test !== defaultPartitionsPercentage.test ||
+      rowsPartitionsPercentage.validation !==
+        defaultPartitionsPercentage.validation
+    ) {
+      return true;
+    }
+    return false;
+  };
+
+  const validateColumns = async () => {
+    try {
+      const validation = await validateColumnsRequest(
+        newExp.task_name,
+        newExp.dataset.id,
+        inputColumns,
+        outputColumns,
+      );
+      setColumnsAreValid(validation.dataset_status === "valid");
+    } catch (error) {
+      enqueueSnackbar("Error while trying to obtain the columns validation.");
+      if (error.response) {
+        console.error("Response error:", error.message);
+      } else if (error.request) {
+        console.error("Request error", error.request);
+      } else {
+        console.error("Unknown Error", error.message);
+      }
+    }
+  };
+
+  useEffect(() => {
+    if (columnsReady && splitsReady) {
+      validateColumns();
+    }
+  }, [columnsReady]);
+
+  useEffect(() => {
+    if (columnsAreValid) {
+      setNewExp({
+        ...newExp,
+        input_columns: inputColumns,
+        output_columns: outputColumns,
+        splits: !isRandom
+          ? { ...rowsPartitionsIndex, is_random: false, has_changed: true }
+          : {
+              ...rowsPartitionsPercentage,
+              is_random: true,
+              has_changed: checkIfSplitsHasChanged(),
+            },
+      }); // splits should depend on preference
+      setNextEnabled(true);
+    } else {
+      setNextEnabled(false);
+    }
+  }, [columnsReady, splitsReady, columnsAreValid]);
+
+  useEffect(() => {
+    getDatasetInfo();
+    getTaskRequirements();
+  }, []);
+
+  const parseListOfStrings = (stringsList) => {
+    return stringsList.join(" or ");
+  };
+  return (
+    <React.Fragment>
+      <Alert severity={columnsAreValid ? "success" : "error"} sx={{ mb: 1 }}>
+        <AlertTitle>
+          {columnsAreValid
+            ? "Current Input and Output columns match"
+            : "Current Input and Output columns doesn't match"}{" "}
+          {taskRequirements.name} requirements
+        </AlertTitle>
+        <Grid container spacing={2}>
+          <Grid item xs={12}>
+            The input columns must be of the types{" "}
+            {taskRequirements
+              ? parseListOfStrings(taskRequirements.metadata.inputs_types)
+              : null}
+            , and they should have a cardinality of{" "}
+            {taskRequirements.metadata.inputs_cardinality}.
+          </Grid>
+          <Grid item xs={12}>
+            The output columns must be of the types{" "}
+            {taskRequirements
+              ? parseListOfStrings(taskRequirements.metadata.outputs_types)
+              : null}
+            , and they should have a cardinality of{" "}
+            {taskRequirements.metadata.outputs_cardinality}.
+          </Grid>
+        </Grid>
+      </Alert>
+
+      {!infoLoading ? (
+        <Grid container spacing={1}>
+          <DivideDatasetColumns
+            datasetInfo={datasetInfo}
+            inputColumns={inputColumns}
+            setInputColumns={setInputColumns}
+            outputColumns={outputColumns}
+            setOutputColumns={setOutputColumns}
+            setColumnsReady={setColumnsReady}
+          />
+          <SplitDatasetRows
+            datasetInfo={datasetInfo}
+            rowsPartitionsIndex={rowsPartitionsIndex}
+            setRowsPartitionsIndex={setRowsPartitionsIndex}
+            rowsPartitionsPercentage={rowsPartitionsPercentage}
+            setRowsPartitionsPercentage={setRowsPartitionsPercentage}
+            setSplitsReady={setSplitsReady}
+            isRandom={isRandom}
+            setIsRandom={setIsRandom}
+          />
+        </Grid>
+      ) : (
+        <Box sx={{ display: "flex" }}>
+          <CircularProgress />
+        </Box>
+      )}
+    </React.Fragment>
+  );
+}
+
+PrepareDatasetStep.propTypes = {
+  newExp: PropTypes.shape({
+    id: PropTypes.string,
+    name: PropTypes.string,
+    dataset: PropTypes.object,
+    task_name: PropTypes.string,
+    input_columns: PropTypes.arrayOf(PropTypes.number),
+    output_columns: PropTypes.arrayOf(PropTypes.number),
+    splits: PropTypes.shape({
+      has_changed: PropTypes.bool,
+      is_random: PropTypes.bool,
+      training: PropTypes.number,
+      validation: PropTypes.number,
+      testing: PropTypes.number,
+    }),
+    step: PropTypes.string,
+    created: PropTypes.instanceOf(Date),
+    last_modified: PropTypes.instanceOf(Date),
+    runs: PropTypes.array,
+  }),
+  setNewExp: PropTypes.func.isRequired,
+  setNextEnabled: PropTypes.func.isRequired,
+};
+export default PrepareDatasetStep;
diff --git a/DashAI/front/src/components/experiments/SelectDatasetStep.jsx b/DashAI/front/src/components/experiments/SelectDatasetStep.jsx
index 221d06c3..39b01d71 100644
--- a/DashAI/front/src/components/experiments/SelectDatasetStep.jsx
+++ b/DashAI/front/src/components/experiments/SelectDatasetStep.jsx
@@ -54,12 +54,7 @@ function SelectDatasetStep({ newExp, setNewExp, setNextEnabled }) {
     setLoading(true);
     try {
       const datasets = await getDatasetsRequest();
-
-      const expTaskName = newExp.task_name;
-      const filteredDatasets = datasets.filter(
-        (dataset) => dataset.task_name === expTaskName,
-      );
-      setDatasets(filteredDatasets);
+      setDatasets(datasets);
     } catch (error) {
       enqueueSnackbar("Error while trying to obtain the datasets list.");
       setRequestError(true);
@@ -129,9 +124,7 @@ function SelectDatasetStep({ newExp, setNewExp, setNextEnabled }) {
       {datasets.length === 0 && !loading && !requestError && (
         <React.Fragment>
           <Alert severity="warning" sx={{ mb: 2 }}>
-            <AlertTitle>
-              There is no datasets associated to the selected task.
-            </AlertTitle>
+            <AlertTitle>There is no datasets available.</AlertTitle>
             Go to{" "}
             <Link component={RouterLink} to="/app/data">
               data tab
@@ -173,6 +166,15 @@ SelectDatasetStep.propTypes = {
     name: PropTypes.string,
     dataset: PropTypes.object,
     task_name: PropTypes.string,
+    input_columns: PropTypes.arrayOf(PropTypes.number),
+    output_columns: PropTypes.arrayOf(PropTypes.number),
+    splits: PropTypes.shape({
+      has_changed: PropTypes.bool,
+      is_random: PropTypes.bool,
+      training: PropTypes.number,
+      validation: PropTypes.number,
+      testing: PropTypes.number,
+    }),
     step: PropTypes.string,
     created: PropTypes.instanceOf(Date),
     last_modified: PropTypes.instanceOf(Date),
diff --git a/DashAI/front/src/components/experiments/SetNameAndTaskStep.jsx b/DashAI/front/src/components/experiments/SetNameAndTaskStep.jsx
index 782d8d57..6f4a0466 100644
--- a/DashAI/front/src/components/experiments/SetNameAndTaskStep.jsx
+++ b/DashAI/front/src/components/experiments/SetNameAndTaskStep.jsx
@@ -150,6 +150,15 @@ SetNameAndTaskStep.propTypes = {
     name: PropTypes.string,
     dataset: PropTypes.object,
     task_name: PropTypes.string,
+    input_columns: PropTypes.arrayOf(PropTypes.number),
+    output_columns: PropTypes.arrayOf(PropTypes.number),
+    splits: PropTypes.shape({
+      has_changed: PropTypes.bool,
+      is_random: PropTypes.bool,
+      training: PropTypes.number,
+      validation: PropTypes.number,
+      testing: PropTypes.number,
+    }),
     step: PropTypes.string,
     created: PropTypes.instanceOf(Date),
     last_modified: PropTypes.instanceOf(Date),
diff --git a/DashAI/front/src/components/experiments/SplitDatasetRows.jsx b/DashAI/front/src/components/experiments/SplitDatasetRows.jsx
new file mode 100644
index 00000000..e85c1bfd
--- /dev/null
+++ b/DashAI/front/src/components/experiments/SplitDatasetRows.jsx
@@ -0,0 +1,272 @@
+import React, { useEffect, useState } from "react";
+import PropTypes from "prop-types";
+import { parseRangeToIndex } from "../../utils/parseRange";
+import {
+  Grid,
+  TextField,
+  Typography,
+  FormControlLabel,
+  Radio,
+  RadioGroup,
+  FormHelperText,
+} from "@mui/material";
+
+function SplitDatasetRows({
+  datasetInfo,
+  rowsPartitionsIndex,
+  setRowsPartitionsIndex,
+  rowsPartitionsPercentage,
+  setRowsPartitionsPercentage,
+  setSplitsReady,
+  isRandom,
+  setIsRandom,
+}) {
+  const totalRows = datasetInfo.total_rows;
+
+  // handle rows numbers change state
+  const [rowsPartitionsError, setRowsPartitionsError] = useState(false);
+  const [rowsPartitionsErrorText, setRowsPartitionsErrorText] = useState("");
+
+  const checkSplit = (train, validation, test) => {
+    return train + validation + test === 1;
+  };
+
+  const handleRowsPreferenceChange = (event) => {
+    if (event.target.value === "splitByIndex") {
+      setIsRandom(false);
+      setRowsPartitionsPercentage({ train: 0.6, test: 0.2, validation: 0.2 });
+    } else {
+      setIsRandom(true);
+      setRowsPartitionsIndex({ train: [], test: [], validation: [] });
+    }
+    setRowsPartitionsError(false);
+    setRowsPartitionsErrorText("");
+  };
+
+  const handleRowsChange = (event) => {
+    const value = event.target.value;
+    const id = event.target.id; // TODO: check that the training, validation and testing rows dont overlap
+    if (isRandom === false) {
+      try {
+        const rowsIndex = parseRangeToIndex(value, totalRows);
+        switch (id) {
+          case "train":
+            setRowsPartitionsIndex({
+              ...rowsPartitionsIndex,
+              train: rowsIndex,
+            });
+            break;
+          case "validation":
+            setRowsPartitionsIndex({
+              ...rowsPartitionsIndex,
+              validation: rowsIndex,
+            });
+            break;
+          case "test":
+            setRowsPartitionsIndex({
+              ...rowsPartitionsIndex,
+              test: rowsIndex,
+            });
+            break;
+        }
+        setRowsPartitionsError(false);
+      } catch (error) {
+        setRowsPartitionsErrorText(error.message);
+        setRowsPartitionsError(true);
+      }
+    } else {
+      let newSplit = rowsPartitionsPercentage;
+      switch (id) {
+        case "train":
+          newSplit = { ...newSplit, train: parseFloat(value) };
+          break;
+        case "validation":
+          newSplit = { ...newSplit, validation: parseFloat(value) };
+          break;
+        case "test":
+          newSplit = { ...newSplit, test: parseFloat(value) };
+          break;
+      }
+      setRowsPartitionsPercentage(newSplit);
+      if (!checkSplit(newSplit.train, newSplit.validation, newSplit.test)) {
+        setRowsPartitionsErrorText(
+          "Splits should be numbers between 0 and 1 and should add 1 in total",
+        );
+        setRowsPartitionsError(true);
+      } else {
+        setRowsPartitionsError(false);
+      }
+    }
+  };
+
+  useEffect(() => {
+    // check if splits doesnt have errors and arent empty
+    if (
+      isRandom === false &&
+      !rowsPartitionsError &&
+      rowsPartitionsIndex.train.length >= 1 &&
+      rowsPartitionsIndex.validation.length >= 1 &&
+      rowsPartitionsIndex.test.length >= 1
+    ) {
+      setSplitsReady(true);
+    } else if (
+      isRandom === true &&
+      !rowsPartitionsError &&
+      rowsPartitionsPercentage.train > 0 &&
+      rowsPartitionsPercentage.validation > 0 &&
+      rowsPartitionsPercentage.test > 0
+    ) {
+      setSplitsReady(true);
+    } else {
+      setSplitsReady(false);
+    }
+  }, [rowsPartitionsIndex, rowsPartitionsPercentage, rowsPartitionsError]);
+  return (
+    <React.Fragment>
+      <Grid container spacing={1}>
+        <Grid item xs={12}>
+          <Typography variant="subtitle1" component="h3" sx={{ mb: 2 }}>
+            Select how to divide the dataset into training, validation and test
+            subsets.
+          </Typography>
+        </Grid>
+      </Grid>
+      <RadioGroup
+        defaultValue={"random"}
+        name="radio-buttons-group"
+        onChange={handleRowsPreferenceChange}
+      >
+        <FormControlLabel
+          value="random"
+          control={<Radio />}
+          label="Use random rows by specifying wich portion of the dataset you want to use for each subset"
+          sx={{ my: 1 }}
+        />
+        {isRandom === true ? (
+          <React.Fragment>
+            <Grid container direction="row" spacing={4}>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="train"
+                  label="Train"
+                  autoComplete="off"
+                  type="number"
+                  size="small"
+                  error={rowsPartitionsError}
+                  defaultValue={rowsPartitionsPercentage.train}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="validation"
+                  label="Validation"
+                  autoComplete="off"
+                  type="number"
+                  size="small"
+                  error={rowsPartitionsError}
+                  defaultValue={rowsPartitionsPercentage.validation}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="test"
+                  label="Test"
+                  type="number"
+                  size="small"
+                  autoComplete="off"
+                  error={rowsPartitionsError}
+                  defaultValue={rowsPartitionsPercentage.test}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+            </Grid>
+            {rowsPartitionsError ? (
+              <FormHelperText>{rowsPartitionsErrorText}</FormHelperText>
+            ) : (
+              <React.Fragment />
+            )}
+          </React.Fragment>
+        ) : (
+          <React.Fragment />
+        )}
+        <FormControlLabel
+          value="splitByIndex"
+          control={<Radio />}
+          label="Use manual splitting by specifying the row indexes of each subset"
+          sx={{ my: 1 }}
+        />
+        {isRandom === false ? (
+          <React.Fragment>
+            <Grid container direction="row" spacing={4}>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="train"
+                  label="Train"
+                  autoComplete="off"
+                  size="small"
+                  error={rowsPartitionsError}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="validation"
+                  label="Validation"
+                  autoComplete="off"
+                  size="small"
+                  error={rowsPartitionsError}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="test"
+                  label="Test"
+                  autoComplete="off"
+                  size="small"
+                  error={rowsPartitionsError}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+            </Grid>
+            {rowsPartitionsError ? (
+              <FormHelperText>{rowsPartitionsErrorText}</FormHelperText>
+            ) : (
+              <React.Fragment />
+            )}
+          </React.Fragment>
+        ) : (
+          <React.Fragment />
+        )}
+      </RadioGroup>
+    </React.Fragment>
+  );
+}
+
+SplitDatasetRows.propTypes = {
+  datasetInfo: PropTypes.shape({
+    test_size: PropTypes.number,
+    total_columns: PropTypes.number,
+    total_rows: PropTypes.number,
+    train_size: PropTypes.number,
+    val_size: PropTypes.number,
+  }),
+  rowsPartitionsIndex: PropTypes.shape({
+    train: PropTypes.arrayOf(PropTypes.number),
+    validation: PropTypes.arrayOf(PropTypes.number),
+    test: PropTypes.arrayOf(PropTypes.number),
+  }),
+  setRowsPartitionsIndex: PropTypes.func.isRequired,
+  rowsPartitionsPercentage: PropTypes.shape({
+    train: PropTypes.number,
+    validation: PropTypes.number,
+    test: PropTypes.number,
+  }),
+  setRowsPartitionsPercentage: PropTypes.func.isRequired,
+  setSplitsReady: PropTypes.func.isRequired,
+  isRandom: PropTypes.bool,
+  setIsRandom: PropTypes.func.isRequired,
+};
+export default SplitDatasetRows;
diff --git a/DashAI/front/src/types/component.ts b/DashAI/front/src/types/component.ts
index f6a59a7a..6234fe10 100644
--- a/DashAI/front/src/types/component.ts
+++ b/DashAI/front/src/types/component.ts
@@ -1,9 +1,11 @@
 import type { IParameterJsonSchema } from "./configurableObject";
+import type { ITaskMetadataParameters } from "./task";
 
 export interface IComponent {
   name: string;
   type: string;
   configurable_object: boolean;
   schema: IParameterJsonSchema;
+  metadata: ITaskMetadataParameters;
   description: string;
 }
diff --git a/DashAI/front/src/types/dataset.ts b/DashAI/front/src/types/dataset.ts
index e1f61920..7bc99c90 100644
--- a/DashAI/front/src/types/dataset.ts
+++ b/DashAI/front/src/types/dataset.ts
@@ -1,7 +1,6 @@
 export interface IDataset {
   id: string;
   name: string;
-  task_name: string;
   created: Date;
   last_modified: Date;
 }
diff --git a/DashAI/front/src/types/experiment.ts b/DashAI/front/src/types/experiment.ts
index 8208a405..56c8c8b6 100644
--- a/DashAI/front/src/types/experiment.ts
+++ b/DashAI/front/src/types/experiment.ts
@@ -5,6 +5,9 @@ export interface IExperiment {
   id: string;
   dataset: IDataset;
   task_name: string;
+  input_columns: string;
+  output_columns: string;
+  splits: string;
   step: string;
   created: Date;
   last_modified: Date;
diff --git a/DashAI/front/src/types/task.ts b/DashAI/front/src/types/task.ts
index f99c7827..3bca2b15 100644
--- a/DashAI/front/src/types/task.ts
+++ b/DashAI/front/src/types/task.ts
@@ -5,3 +5,10 @@ export interface ITask {
   description: string;
   type: string;
 }
+
+export interface ITaskMetadataParameters {
+  inputs_columns: string[];
+  outputs_columns: string[];
+  inputs_cardinality: "n" | number;
+  outputs_cardinality: "n" | number;
+}
diff --git a/DashAI/front/src/utils/parseRange.js b/DashAI/front/src/utils/parseRange.js
new file mode 100644
index 00000000..12d7b3ed
--- /dev/null
+++ b/DashAI/front/src/utils/parseRange.js
@@ -0,0 +1,31 @@
+const rangeRegex = /^(\d+)(-(\d+))*(,(\d+)(-(\d+))*)*$/;
+
+export function parseRangeToIndex(range, maxValue) {
+  const indexArray = [];
+  if (!rangeRegex.test(range)) {
+    throw new Error("The entered text doesn't fit the example format");
+  }
+  const ranges = range.split(",");
+  ranges.forEach((range) => {
+    const [min, max] = range.split("-");
+    if (!range.includes("-") && parseInt(range) <= maxValue) {
+      indexArray.push(parseInt(range));
+    } else if (
+      (!range.includes("-") && parseInt(range) > maxValue) ||
+      parseInt(max) > maxValue
+    ) {
+      throw new Error(
+        "The indexes should be minor than the total of columns or rows",
+      );
+    } else if (parseInt(min) > parseInt(max)) {
+      throw new Error(
+        "The second number of a range must be greater than the first",
+      );
+    } else {
+      for (let i = parseInt(min); i <= parseInt(max); i++) {
+        indexArray.push(i);
+      }
+    }
+  });
+  return indexArray;
+}
diff --git a/DashAI/front/src/utils/typesLists.js b/DashAI/front/src/utils/typesLists.js
new file mode 100644
index 00000000..055d9c69
--- /dev/null
+++ b/DashAI/front/src/utils/typesLists.js
@@ -0,0 +1,30 @@
+export const dataTypesList = [
+  "null",
+  "bool",
+  "int8",
+  "int16",
+  "int32",
+  "int64",
+  "uint8",
+  "uint16",
+  "uint32",
+  "uint64",
+  "float16",
+  "float32",
+  "float64",
+  "time32[(s|ms)]",
+  "time64[(us|ns)]",
+  "timestamp[(s|ms|us|ns)]",
+  "timestamp[(s|ms|us|ns), tz=(tzstring)]",
+  "date32",
+  "date64",
+  "duration[(s|ms|us|ns)]",
+  "decimal128",
+  "decimal256",
+  "binary",
+  "large_binary",
+  "string",
+  "large_string",
+];
+
+export const columnTypesList = ["Value", "ClassLabel"];
diff --git a/alembic.ini b/alembic.ini
new file mode 100644
index 00000000..8ee69ae3
--- /dev/null
+++ b/alembic.ini
@@ -0,0 +1,110 @@
+# A generic, single database configuration.
+
+[alembic]
+# path to migration scripts
+script_location = alembic
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python-dateutil library that can be
+# installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the
+# "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to alembic/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "version_path_separator" below.
+# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions
+
+# version path separator; As mentioned above, this is the character used to split
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
+# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
+# Valid values for version_path_separator are:
+#
+# version_path_separator = :
+# version_path_separator = ;
+# version_path_separator = space
+version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+sqlalchemy.url = sqlite:///DashAI/back/database/DashAI.sqlite
+
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S
diff --git a/alembic/env.py b/alembic/env.py
new file mode 100644
index 00000000..7907b308
--- /dev/null
+++ b/alembic/env.py
@@ -0,0 +1,67 @@
+from logging.config import fileConfig
+
+from sqlalchemy import engine_from_config, pool
+
+from alembic import context
+from DashAI.back.database.models import Base
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+target_metadata = Base.metadata
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    with connectable.connect() as connection:
+        context.configure(connection=connection, target_metadata=target_metadata)
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()
diff --git a/alembic/script.py.mako b/alembic/script.py.mako
new file mode 100644
index 00000000..55df2863
--- /dev/null
+++ b/alembic/script.py.mako
@@ -0,0 +1,24 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}
diff --git a/alembic/versions/2f997776b9a1_add_splits_metadata_to_experiment_model_.py b/alembic/versions/2f997776b9a1_add_splits_metadata_to_experiment_model_.py
new file mode 100644
index 00000000..cb96a15f
--- /dev/null
+++ b/alembic/versions/2f997776b9a1_add_splits_metadata_to_experiment_model_.py
@@ -0,0 +1,25 @@
+"""Add splits metadata to Experiment model on DB
+
+Revision ID: 2f997776b9a1
+Revises: 7ba93a3d42a8
+Create Date: 2023-10-09 12:34:12.786244
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "2f997776b9a1"
+down_revision = "7ba93a3d42a8"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###
diff --git a/alembic/versions/4bce47006c88_add_splits_metadata_to_experiment_model_.py b/alembic/versions/4bce47006c88_add_splits_metadata_to_experiment_model_.py
new file mode 100644
index 00000000..c8c7a000
--- /dev/null
+++ b/alembic/versions/4bce47006c88_add_splits_metadata_to_experiment_model_.py
@@ -0,0 +1,28 @@
+"""Add splits metadata to Experiment model on DB
+
+Revision ID: 4bce47006c88
+Revises: 2f997776b9a1
+Create Date: 2023-10-09 12:57:13.075246
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "4bce47006c88"
+down_revision = "2f997776b9a1"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("experiment", sa.Column("splits", sa.JSON(), nullable=False))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("experiment", "splits")
+    # ### end Alembic commands ###
diff --git a/alembic/versions/649c500185b8_delete_task_name_and_adds_dataloader_.py b/alembic/versions/649c500185b8_delete_task_name_and_adds_dataloader_.py
new file mode 100644
index 00000000..dd781e15
--- /dev/null
+++ b/alembic/versions/649c500185b8_delete_task_name_and_adds_dataloader_.py
@@ -0,0 +1,32 @@
+"""Delete task_name and adds dataloader name to Dataset model
+
+Revision ID: 649c500185b8
+Revises: 80faf11605a2
+Create Date: 2023-08-09 08:43:29.112965
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "649c500185b8"
+down_revision = "80faf11605a2"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("dataset", sa.Column("feature_names", sa.JSON(), nullable=False))
+    op.add_column("dataset", sa.Column("dataset_name", sa.String(), nullable=False))
+    op.drop_column("dataset", "task_name")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("dataset", sa.Column("task_name", sa.VARCHAR(), nullable=False))
+    op.drop_column("dataset", "dataset_name")
+    op.drop_column("dataset", "feature_names")
+    # ### end Alembic commands ###
diff --git a/alembic/versions/7ba93a3d42a8_delete_feature_names_to_experiment_.py b/alembic/versions/7ba93a3d42a8_delete_feature_names_to_experiment_.py
new file mode 100644
index 00000000..7ddac412
--- /dev/null
+++ b/alembic/versions/7ba93a3d42a8_delete_feature_names_to_experiment_.py
@@ -0,0 +1,31 @@
+"""Delete feature_names to Experiment model on DB
+
+Revision ID: 7ba93a3d42a8
+Revises: 7e990d4a75ec
+Create Date: 2023-09-25 10:03:29.526284
+
+"""
+import sqlalchemy as sa
+from sqlalchemy.dialects import sqlite
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "7ba93a3d42a8"
+down_revision = "7e990d4a75ec"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("experiment", "feature_names")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "experiment", sa.Column("feature_names", sqlite.JSON(), nullable=False)
+    )
+    # ### end Alembic commands ###
diff --git a/alembic/versions/7e990d4a75ec_add_metadata_to_experiment_model_on_db.py b/alembic/versions/7e990d4a75ec_add_metadata_to_experiment_model_on_db.py
new file mode 100644
index 00000000..216084ee
--- /dev/null
+++ b/alembic/versions/7e990d4a75ec_add_metadata_to_experiment_model_on_db.py
@@ -0,0 +1,30 @@
+"""Add metadata to Experiment model on DB
+
+Revision ID: 7e990d4a75ec
+Revises: 8307d4c14824
+Create Date: 2023-09-22 22:45:34.448464
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "7e990d4a75ec"
+down_revision = "8307d4c14824"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("experiment", sa.Column("input_columns", sa.JSON(), nullable=False))
+    op.add_column("experiment", sa.Column("output_columns", sa.JSON(), nullable=False))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("experiment", "output_columns")
+    op.drop_column("experiment", "input_columns")
+    # ### end Alembic commands ###
diff --git a/alembic/versions/80faf11605a2_initial_migration.py b/alembic/versions/80faf11605a2_initial_migration.py
new file mode 100644
index 00000000..d3da1324
--- /dev/null
+++ b/alembic/versions/80faf11605a2_initial_migration.py
@@ -0,0 +1,30 @@
+"""Initial migration
+
+Revision ID: 80faf11605a2
+Revises:
+Create Date: 2023-08-08 09:13:58.794747
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "80faf11605a2"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("experiment", "step")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "experiment", sa.Column("step", sa.VARCHAR(length=19), nullable=False)
+    )
+    # ### end Alembic commands ###
diff --git a/alembic/versions/8307d4c14824_add_metadata_to_experiment_model_on_db.py b/alembic/versions/8307d4c14824_add_metadata_to_experiment_model_on_db.py
new file mode 100644
index 00000000..39c1b41d
--- /dev/null
+++ b/alembic/versions/8307d4c14824_add_metadata_to_experiment_model_on_db.py
@@ -0,0 +1,28 @@
+"""Add metadata to Experiment model on DB
+
+Revision ID: 8307d4c14824
+Revises: d267eb740e0f
+Create Date: 2023-09-22 21:46:37.521272
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8307d4c14824"
+down_revision = "d267eb740e0f"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("experiment", sa.Column("feature_names", sa.JSON(), nullable=False))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("experiment", "feature_names")
+    # ### end Alembic commands ###
diff --git a/alembic/versions/d267eb740e0f_delete_dataloader_name_in_dataset_model.py b/alembic/versions/d267eb740e0f_delete_dataloader_name_in_dataset_model.py
new file mode 100644
index 00000000..13ef6cc7
--- /dev/null
+++ b/alembic/versions/d267eb740e0f_delete_dataloader_name_in_dataset_model.py
@@ -0,0 +1,28 @@
+"""Delete dataloader name in Dataset model
+
+Revision ID: d267eb740e0f
+Revises: 649c500185b8
+Create Date: 2023-08-09 20:45:43.229983
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "d267eb740e0f"
+down_revision = "649c500185b8"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("dataset", "dataset_name")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("dataset", sa.Column("dataset_name", sa.VARCHAR(), nullable=False))
+    # ### end Alembic commands ###
diff --git a/alembic/versions/e40a1eeee0ef_delete_feature_names_from_datasets.py b/alembic/versions/e40a1eeee0ef_delete_feature_names_from_datasets.py
new file mode 100644
index 00000000..ac89ee98
--- /dev/null
+++ b/alembic/versions/e40a1eeee0ef_delete_feature_names_from_datasets.py
@@ -0,0 +1,29 @@
+"""Delete feature_names from datasets
+
+Revision ID: e40a1eeee0ef
+Revises: 4bce47006c88
+Create Date: 2023-11-28 15:58:18.913263
+
+"""
+import sqlalchemy as sa
+from sqlalchemy.dialects import sqlite
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "e40a1eeee0ef"
+down_revision = "4bce47006c88"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("dataset", "feature_names")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("dataset", sa.Column("feature_names", sqlite.JSON(), nullable=False))
+    # ### end Alembic commands ###
diff --git a/requirements.txt b/requirements.txt
index 2a14bd19..2c6bf98b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 dependency-injector>=4.0,<5.0
 fastapi[all]>=0.96
 SQLAlchemy>=2.0
+alembic==1.11.1
 numpy>=1.17.3
 joblib>=1.2.0
 pydantic>=2.0.2


Test patch:
diff --git a/tests/back/api/conftest.py b/tests/back/api/conftest.py
index 7269a213..4965ec6c 100644
--- a/tests/back/api/conftest.py
+++ b/tests/back/api/conftest.py
@@ -1,5 +1,6 @@
 import pathlib
 import shutil
+import time
 
 import pytest
 from fastapi.testclient import TestClient
@@ -9,11 +10,24 @@
 TEST_PATH = "tmp"
 
 
+def remove_dir_with_retry(directory, max_attempts=5, sleep_seconds=1):
+    for attempt in range(max_attempts):
+        try:
+            shutil.rmtree(directory)
+            print(f"Successfully removed directory: {directory}")
+            break
+        except PermissionError as e:
+            print(f"Attempt {attempt + 1} failed with error: {e}")
+            time.sleep(sleep_seconds)
+    else:
+        print(f"Failed to remove directory after {max_attempts} attempts.")
+
+
 @pytest.fixture(scope="module", autouse=True)
 def client():
-    app = create_app(local_path=pathlib.Path(TEST_PATH), logging_level="DEBUG")
+    app = create_app(local_path=pathlib.Path(TEST_PATH), logging_level="ERROR")
 
     yield TestClient(app)
 
     app.container.db().dispose_engine()
-    shutil.rmtree(app.container.config.provided()["LOCAL_PATH"])
+    remove_dir_with_retry(app.container.config.provided()["LOCAL_PATH"])
diff --git a/tests/back/api/test_components_api.py b/tests/back/api/test_components_api.py
index 4d814487..74db79e1 100644
--- a/tests/back/api/test_components_api.py
+++ b/tests/back/api/test_components_api.py
@@ -1,5 +1,6 @@
 """Tests module for components API."""
 import pytest
+from datasets import ClassLabel, Image, Value
 from fastapi.testclient import TestClient
 
 from DashAI.back.dataloaders import BaseDataLoader
@@ -33,6 +34,13 @@
 class TestTask1(BaseTask):
     DESCRIPTION = "Task 1."
 
+    metadata = {
+        "inputs_types": [ClassLabel, Value],
+        "outputs_types": [ClassLabel],
+        "inputs_cardinality": "n",
+        "outputs_cardinality": 1,
+    }
+
     @classmethod
     def get_schema(cls) -> dict:
         return {"class": "TestTask1"}
@@ -40,6 +48,12 @@ def get_schema(cls) -> dict:
 
 class TestTask2(BaseTask):
     DESCRIPTION = "Task 2."
+    metadata = {
+        "inputs_types": [Image],
+        "outputs_types": [ClassLabel],
+        "inputs_cardinality": 1,
+        "outputs_cardinality": 1,
+    }
 
     @classmethod
     def get_schema(cls) -> dict:
@@ -136,6 +150,12 @@ def test_get_component_by_id(client: TestClient):
         "type": "Task",
         "configurable_object": False,
         "schema": None,
+        "metadata": {
+            "inputs_types": ["ClassLabel", "Value"],
+            "outputs_types": ["ClassLabel"],
+            "inputs_cardinality": "n",
+            "outputs_cardinality": 1,
+        },
         "description": "Task 1.",
     }
 
@@ -146,6 +166,12 @@ def test_get_component_by_id(client: TestClient):
         "type": "Task",
         "configurable_object": False,
         "schema": None,
+        "metadata": {
+            "inputs_types": ["Image"],
+            "outputs_types": ["ClassLabel"],
+            "inputs_cardinality": 1,
+            "outputs_cardinality": 1,
+        },
         "description": "Task 2.",
     }
 
@@ -156,6 +182,7 @@ def test_get_component_by_id(client: TestClient):
         "type": "DataLoader",
         "configurable_object": True,
         "schema": {},
+        "metadata": None,
         "description": None,
     }
 
@@ -197,6 +224,12 @@ def test_get_all_components(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["ClassLabel", "Value"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": "n",
+                "outputs_cardinality": 1,
+            },
             "description": "Task 1.",
         },
         {
@@ -204,6 +237,12 @@ def test_get_all_components(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["Image"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": 1,
+                "outputs_cardinality": 1,
+            },
             "description": "Task 2.",
         },
         {
@@ -211,6 +250,7 @@ def test_get_all_components(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -218,6 +258,7 @@ def test_get_all_components(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -225,6 +266,7 @@ def test_get_all_components(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -232,6 +274,7 @@ def test_get_all_components(client: TestClient):
             "type": "Model",
             "configurable_object": True,
             "schema": {"properties": {"parameter_1": {"type": "number"}}},
+            "metadata": None,
             "description": None,
         },
         {
@@ -241,6 +284,7 @@ def test_get_all_components(client: TestClient):
             "schema": {
                 "properties": {"parameter_2": {"type": "string", "enum": ["a", "b"]}}
             },
+            "metadata": None,
             "description": None,
         },
     ]
@@ -262,6 +306,12 @@ def test_get_components_select_only_tasks(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["ClassLabel", "Value"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": "n",
+                "outputs_cardinality": 1,
+            },
             "description": "Task 1.",
         },
         {
@@ -269,6 +319,12 @@ def test_get_components_select_only_tasks(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["Image"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": 1,
+                "outputs_cardinality": 1,
+            },
             "description": "Task 2.",
         },
     ]
@@ -285,6 +341,7 @@ def test_get_components_select_only_dataloaders(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -292,6 +349,7 @@ def test_get_components_select_only_dataloaders(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -299,6 +357,7 @@ def test_get_components_select_only_dataloaders(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -315,6 +374,7 @@ def test_get_components_select_tasks_and_models(client: TestClient):
             "type": "Model",
             "configurable_object": True,
             "schema": {"properties": {"parameter_1": {"type": "number"}}},
+            "metadata": None,
             "description": None,
         },
         {
@@ -324,6 +384,7 @@ def test_get_components_select_tasks_and_models(client: TestClient):
             "schema": {
                 "properties": {"parameter_2": {"type": "string", "enum": ["a", "b"]}}
             },
+            "metadata": None,
             "description": None,
         },
         {
@@ -331,6 +392,12 @@ def test_get_components_select_tasks_and_models(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["ClassLabel", "Value"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": "n",
+                "outputs_cardinality": 1,
+            },
             "description": "Task 1.",
         },
         {
@@ -338,6 +405,12 @@ def test_get_components_select_tasks_and_models(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["Image"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": 1,
+                "outputs_cardinality": 1,
+            },
             "description": "Task 2.",
         },
     ]
@@ -379,6 +452,12 @@ def test_get_components_ignore_models(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["ClassLabel", "Value"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": "n",
+                "outputs_cardinality": 1,
+            },
             "description": "Task 1.",
         },
         {
@@ -386,6 +465,12 @@ def test_get_components_ignore_models(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["Image"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": 1,
+                "outputs_cardinality": 1,
+            },
             "description": "Task 2.",
         },
         {
@@ -393,6 +478,7 @@ def test_get_components_ignore_models(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -400,6 +486,7 @@ def test_get_components_ignore_models(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -407,6 +494,7 @@ def test_get_components_ignore_models(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -422,6 +510,7 @@ def test_get_components_ignore_tasks_and_models(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -429,6 +518,7 @@ def test_get_components_ignore_tasks_and_models(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -436,6 +526,7 @@ def test_get_components_ignore_tasks_and_models(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -476,6 +567,7 @@ def test_get_components_related_with_some_task(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -483,6 +575,7 @@ def test_get_components_related_with_some_task(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -490,6 +583,7 @@ def test_get_components_related_with_some_task(client: TestClient):
             "type": "Model",
             "configurable_object": True,
             "schema": {"properties": {"parameter_1": {"type": "number"}}},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -504,6 +598,12 @@ def test_get_components_related_inverse_relation(client: TestClient):
             "type": "Task",
             "configurable_object": False,
             "schema": None,
+            "metadata": {
+                "inputs_types": ["ClassLabel", "Value"],
+                "outputs_types": ["ClassLabel"],
+                "inputs_cardinality": "n",
+                "outputs_cardinality": 1,
+            },
             "description": "Task 1.",
         }
     ]
@@ -542,6 +642,7 @@ def test_get_components_dataloader_component_parent(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -549,6 +650,7 @@ def test_get_components_dataloader_component_parent(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -581,6 +683,7 @@ def test_get_components_by_type_and_task(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -588,6 +691,7 @@ def test_get_components_by_type_and_task(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -605,6 +709,7 @@ def test_get_components_by_type_and_task_2(client: TestClient):
             "type": "Model",
             "configurable_object": True,
             "schema": {"properties": {"parameter_1": {"type": "number"}}},
+            "metadata": None,
             "description": None,
         }
     ]
@@ -622,6 +727,7 @@ def test_get_components_select_and_ignore_by_type(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -629,6 +735,7 @@ def test_get_components_select_and_ignore_by_type(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -636,6 +743,7 @@ def test_get_components_select_and_ignore_by_type(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
@@ -654,6 +762,7 @@ def test_get_components_select_type_and_parent(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
         {
@@ -661,6 +770,7 @@ def test_get_components_select_type_and_parent(client: TestClient):
             "type": "DataLoader",
             "configurable_object": True,
             "schema": {},
+            "metadata": None,
             "description": None,
         },
     ]
diff --git a/tests/back/api/test_dataset_api.py b/tests/back/api/test_dataset_api.py
index 79c72947..e0ce2292 100644
--- a/tests/back/api/test_dataset_api.py
+++ b/tests/back/api/test_dataset_api.py
@@ -1,4 +1,3 @@
-import json
 import os
 
 import pytest
@@ -14,10 +13,8 @@ def create_dataset_1(client):
         response = client.post(
             "/api/v1/dataset/",
             data={
-                "params": """{  "task_name": "TabularClassificationTask",
-                                    "dataloader": "CSVDataLoader",
+                "params": """{  "dataloader": "CSVDataLoader",
                                     "dataset_name": "test_csv",
-                                    "outputs_columns": [],
                                     "splits_in_folders": false,
                                     "splits": {
                                         "train_size": 0.8,
@@ -47,10 +44,8 @@ def create_dataset_2(client):
         response = client.post(
             "/api/v1/dataset/",
             data={
-                "params": """{  "task_name": "TabularClassificationTask",
-                                    "dataloader": "CSVDataLoader",
+                "params": """{  "dataloader": "CSVDataLoader",
                                     "dataset_name": "test_csv2",
-                                    "outputs_columns": [],
                                     "splits_in_folders": false,
                                     "splits": {
                                         "train_size": 0.5,
@@ -78,16 +73,6 @@ def test_create_csv_dataset(client: TestClient, response_1, response_2) -> None:
     assert response_1.status_code == 200, response_1.text
     data = response_1.json()
     assert data["name"] == "test_csv"
-    assert data["task_name"] == "TabularClassificationTask"
-    assert data["feature_names"] == json.dumps(
-        [
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ]
-    )
-
     response_2 = client.get("/api/v1/dataset/2")
     assert response_2.status_code == 200, response_2.text
     data = response_2.json()
@@ -99,7 +84,6 @@ def test_get_all_datasets(client: TestClient):
     assert response.status_code == 200, response.text
     data = response.json()
     assert data[0]["name"] == "test_csv"
-    assert data[0]["task_name"] == "TabularClassificationTask"
     assert data[1]["name"] == "test_csv2"
 
 
@@ -109,17 +93,28 @@ def test_get_unexistant_dataset(client: TestClient):
     assert response.text == '{"detail":"Dataset not found"}'
 
 
-def test_modify_dataset(client: TestClient):
+def test_get_types(client: TestClient):
+    response = client.get("/api/v1/dataset/2/types")
+    data = response.json()
+    assert data == {
+        "SepalLengthCm": {"type": "Value", "dtype": "float64"},
+        "SepalWidthCm": {"type": "Value", "dtype": "float64"},
+        "PetalLengthCm": {"type": "Value", "dtype": "float64"},
+        "PetalWidthCm": {"type": "Value", "dtype": "float64"},
+        "Species": {"type": "Value", "dtype": "string"},
+    }
+
+
+def test_modify_dataset_name(client: TestClient):
     response = client.patch(
         "/api/v1/dataset/2",
-        params={"name": "test_modify_name", "task_name": "UnknownTask"},
+        json={"name": "test_modify_name"},
     )
     assert response.status_code == 200, response.text
     response = client.get("/api/v1/dataset/2")
     assert response.status_code == 200, response.text
     data = response.json()
     assert data["name"] == "test_modify_name"
-    assert data["task_name"] == "UnknownTask"
 
 
 def test_delete_dataset(client: TestClient):
@@ -128,43 +123,3 @@ def test_delete_dataset(client: TestClient):
 
     response = client.delete("/api/v1/dataset/2")
     assert response.status_code == 204, response.text
-
-
-def test_dataset_without_feature_names(client: TestClient):
-    script_dir = os.path.dirname(__file__)
-    test_dataset = "iris_no_header_names.csv"
-    abs_file_path = os.path.join(script_dir, test_dataset)
-    with open(abs_file_path, "rb") as csv:
-        response = client.post(
-            "/api/v1/dataset/",
-            data={
-                "params": """{  "task_name": "TabularClassificationTask",
-                                    "dataloader": "CSVDataLoader",
-                                    "dataset_name": "iris_no_header_test_csv",
-                                    "outputs_columns": [],
-                                    "splits_in_folders": false,
-                                    "splits": {
-                                        "train_size": 0.8,
-                                        "test_size": 0.1,
-                                        "val_size": 0.1,
-                                        "seed": 42,
-                                        "shuffle": true,
-                                        "stratify": false
-                                    },
-                                    "dataloader_params": {
-                                        "separator": ","
-                                    }
-                                }""",
-                "url": "",
-            },
-            files={"file": ("filename", csv, "text/csv")},
-        )
-    assert response.status_code == 201, response.text
-    response = client.get("/api/v1/dataset/3")
-    assert response.status_code == 200, response.text
-    data = response.json()
-    assert data["name"] == "iris_no_header_test_csv"
-    assert data["task_name"] == "TabularClassificationTask"
-    assert data["feature_names"] == json.dumps(
-        ["Unnamed: 0", "Unnamed: 1", "Unnamed: 2", "Unnamed: 3", "Unnamed: 4"]
-    )
diff --git a/tests/back/api/test_experiments_api.py b/tests/back/api/test_experiments_api.py
index 9ee1a192..bb5a408e 100644
--- a/tests/back/api/test_experiments_api.py
+++ b/tests/back/api/test_experiments_api.py
@@ -1,38 +1,58 @@
 import json
+import os
 
 import pytest
 from fastapi.testclient import TestClient
 
-from DashAI.back.dependencies.database.models import Dataset
+input_columns_1 = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]
+input_columns_2 = ["SepalLengthCm", "PetalWidthCm"]
+output_columns = ["Species"]
+splits = json.dumps(
+    {
+        "train": 0.5,
+        "test": 0.2,
+        "validation": 0.3,
+        "is_random": True,
+        "has_changed": True,
+        "seed": 42,
+        "shuffle": True,
+        "stratify": False,
+    }
+)
 
 
 @pytest.fixture(scope="module", name="dataset_id")
-def create_dummy_dataset(client: TestClient):
-    """Create a dummy dataset for the experiments."""
-    container = client.app.container
-    session = container.db.provided().session
-
-    with session() as db:
-        # Create Dummy Dataset
-        dummy_dataset = Dataset(
-            name="DummyDataset",
-            task_name="TabularClassificationTask",
-            file_path="dummy.csv",
-            feature_names=json.dumps([]),
+def create_dataset(client):
+    """Create testing dataset 1."""
+    abs_file_path = os.path.join(os.path.dirname(__file__), "iris.csv")
+
+    with open(abs_file_path, "rb") as csv:
+        response = client.post(
+            "/api/v1/dataset/",
+            data={
+                "params": """{  "dataloader": "CSVDataLoader",
+                                    "dataset_name": "DummyDataset",
+                                    "splits_in_folders": false,
+                                    "splits": {
+                                        "train_size": 0.8,
+                                        "test_size": 0.1,
+                                        "val_size": 0.1,
+                                        "seed": 42,
+                                        "shuffle": true
+                                    },
+                                    "dataloader_params": {
+                                        "separator": ","
+                                    }
+                                }""",
+                "url": "",
+            },
+            files={"file": ("filename", csv, "text/csv")},
         )
-        db.add(dummy_dataset)
-        db.commit()
-        db.refresh(dummy_dataset)
-
-        yield dummy_dataset.id
-
-        # Delete the dataset
-        db.delete(dummy_dataset)
-        db.commit()
+    return response.json()["id"]
 
 
 @pytest.fixture(scope="module", name="response_1")
-def create_experiment_1(client: TestClient, dataset_id: int):
+def create_experiment_1(client: TestClient, dataset_id):
     """Create experiment 1."""
     return client.post(
         "/api/v1/experiment/",
@@ -40,19 +60,25 @@ def create_experiment_1(client: TestClient, dataset_id: int):
             "dataset_id": dataset_id,
             "task_name": "TabularClassificationTask",
             "name": "ExperimentA",
+            "input_columns": [1, 2, 3, 4],
+            "output_columns": [5],
+            "splits": splits,
         },
     )
 
 
 @pytest.fixture(scope="module", name="response_2")
-def create_experiment_2(client: TestClient, dataset_id: int):
+def create_experiment_2(client: TestClient, dataset_id):
     """Create experiment 2."""
     return client.post(
         "/api/v1/experiment/",
         json={
             "dataset_id": dataset_id,
             "task_name": "TabularClassificationTask",
-            "name": "Experiment2",
+            "name": "ExperimentB",
+            "input_columns": [1, 4],
+            "output_columns": [5],
+            "splits": splits,
         },
     )
 
@@ -71,6 +97,9 @@ def test_create_and_get_experiment(
     assert data["dataset_id"] == dataset_id
     assert data["task_name"] == "TabularClassificationTask"
     assert data["name"] == "ExperimentA"
+    assert data["input_columns"] == input_columns_1
+    assert data["output_columns"] == output_columns
+    assert data["splits"] == splits
 
     # test get experiment by id 2.
     response = client.get("/api/v1/experiment/2")
@@ -78,10 +107,13 @@ def test_create_and_get_experiment(
     data = response.json()
     assert data["dataset_id"] == dataset_id
     assert data["task_name"] == "TabularClassificationTask"
-    assert data["name"] == "Experiment2"
+    assert data["name"] == "ExperimentB"
+    assert data["input_columns"] == input_columns_2
+    assert data["output_columns"] == output_columns
+    assert data["splits"] == splits
 
 
-def test_get_all_experiments(client: TestClient, dataset_id: int):
+def test_get_all_experiments(client: TestClient, dataset_id):
     """Test that all experiments can be retrieved."""
     response = client.get("/api/v1/experiment")
 
@@ -137,3 +169,63 @@ def test_delete_experiment(client: TestClient):
 
     response = client.delete("/api/v1/experiment/2")
     assert response.status_code == 204, response.text
+
+
+def test_get_columns_validation_valid(client: TestClient, dataset_id: int):
+    response = client.post(
+        "/api/v1/experiment/validation",
+        json={
+            "task_name": "TabularClassificationTask",
+            "dataset_id": dataset_id,
+            "inputs_columns": [1, 2, 3, 4],
+            "outputs_columns": [5],
+        },
+    )
+    assert response.status_code == 200, response.text
+    json = response.json()
+    assert json["dataset_status"] == "valid"
+
+
+def test_get_columns_validation_invalid(client: TestClient, dataset_id: int):
+    response = client.post(
+        "/api/v1/experiment/validation",
+        json={
+            "task_name": "ImageClassificationTask",
+            "dataset_id": dataset_id,
+            "inputs_columns": [1, 2, 3, 4],
+            "outputs_columns": [5],
+        },
+    )
+    assert response.status_code == 200, response.text
+    json = response.json()
+    assert json["dataset_status"] == "invalid"
+
+
+def test_get_columns_validation_wrong_task_name(client: TestClient, dataset_id: int):
+    response = client.post(
+        "/api/v1/experiment/validation",
+        json={
+            "task_name": "TabularClassTask",
+            "dataset_id": dataset_id,
+            "inputs_columns": [1, 2, 3, 4],
+            "outputs_columns": [5],
+        },
+    )
+    assert response.status_code == 404, response.text
+    assert (
+        response.text == '{"detail":"Task TabularClassTask not found in the registry."}'
+    )
+
+
+def test_get_columns_validation_wrong_dataset(client: TestClient):
+    response = client.post(
+        "/api/v1/experiment/validation",
+        json={
+            "task_name": "TabularClassificationTask",
+            "dataset_id": 127,
+            "inputs_columns": [1, 2, 3, 4],
+            "outputs_columns": [5],
+        },
+    )
+    assert response.status_code == 404, response.text
+    assert response.text == '{"detail":"Dataset not found"}'
diff --git a/tests/back/api/test_jobs.py b/tests/back/api/test_jobs.py
index 6c681891..2ecc5c86 100644
--- a/tests/back/api/test_jobs.py
+++ b/tests/back/api/test_jobs.py
@@ -1,3 +1,4 @@
+import json
 import os
 
 import joblib
@@ -16,12 +17,8 @@
 class DummyTask(BaseTask):
     name: str = "DummyTask"
 
-    def prepare_for_task(self, dataset):
-        return {
-            "train": {"input": [], "output": []},
-            "validation": {"input": [], "output": []},
-            "test": {"input": [], "output": []},
-        }
+    def prepare_for_task(self, dataset, output_columns):
+        return dataset
 
 
 class DummyModel(BaseModel):
@@ -37,10 +34,10 @@ def save(self, filename):
     def load(self, filename):
         return
 
-    def predict(self, data):
+    def predict(self, x):
         return {}
 
-    def fit(self, data):
+    def fit(self, x, y):
         return
 
 
@@ -57,10 +54,10 @@ def save(self, filename):
     def load(self, filename):
         return
 
-    def predict(self, data):
+    def predict(self, x):
         return {}
 
-    def fit(self, data):
+    def fit(self, x, y):
         raise Exception("Always fails")
 
 
@@ -101,18 +98,15 @@ def fixture_dataset_id(client: TestClient):
         response = client.post(
             "/api/v1/dataset/",
             data={
-                "params": """{  "task_name": "TabularClassificationTask",
-                                    "dataloader": "CSVDataLoader",
-                                    "dataset_name": "test_csv2",
-                                    "outputs_columns": [],
+                "params": """{ "dataloader": "CSVDataLoader",
+                                    "dataset_name": "test_csv3",
                                     "splits_in_folders": false,
                                     "splits": {
                                         "train_size": 0.5,
                                         "test_size": 0.2,
                                         "val_size": 0.3,
                                         "seed": 42,
-                                        "shuffle": true,
-                                        "stratify": false
+                                        "shuffle": true
                                     },
                                     "dataloader_params": {
                                         "separator": ","
@@ -141,6 +135,20 @@ def create_experiment(client: TestClient, dataset_id: int):
             dataset_id=dataset_id,
             name="DummyExperiment",
             task_name="DummyTask",
+            input_columns=[],
+            output_columns=[],
+            splits=json.dumps(
+                {
+                    "train": 0.5,
+                    "test": 0.2,
+                    "validation": 0.3,
+                    "is_random": True,
+                    "has_changed": True,
+                    "seed": 42,
+                    "shuffle": True,
+                    "stratify": False,
+                }
+            ),
         )
         db.add(experiment)
         db.commit()
@@ -292,6 +300,7 @@ def test_execute_jobs(client: TestClient, run_id: int, failed_run_id: int):
 
     response = client.get(f"/api/v1/run/{run_id}")
     data = response.json()
+    print(data["status"])
     assert data["status"] == 3
     assert isinstance(data["train_metrics"], dict)
     assert "DummyMetric" in data["train_metrics"]
diff --git a/tests/back/api/test_predict.py b/tests/back/api/test_predict.py
index 4dcc2d91..4c7dfacb 100644
--- a/tests/back/api/test_predict.py
+++ b/tests/back/api/test_predict.py
@@ -1,3 +1,4 @@
+import json
 import os
 
 import pytest
@@ -17,16 +18,14 @@ def create_dataset(client: TestClient):
             data={
                 "params": """{  "task_name": "TabularClassificationTask",
                                     "dataloader": "JSONDataLoader",
-                                    "dataset_name": "test_json_2",
-                                    "outputs_columns": ["class"],
+                                    "dataset_name": "test_json",
                                     "splits_in_folders": false,
                                     "splits": {
                                         "train_size": 0.5,
                                         "test_size": 0.2,
                                         "val_size": 0.3,
                                         "seed": 42,
-                                        "shuffle": false,
-                                        "stratify": false
+                                        "shuffle": false
                                     },
                                     "dataloader_params": {
                                         "data_key": "data"
@@ -54,6 +53,20 @@ def create_experiment(client: TestClient, dataset_id: int):
             dataset_id=dataset_id,
             name="Experiment",
             task_name="TabularClassificationTask",
+            input_columns=["feature_0", "feature_1", "feature_2", "feature_3"],
+            output_columns=["class"],
+            splits=json.dumps(
+                {
+                    "train": 0.5,
+                    "test": 0.2,
+                    "validation": 0.3,
+                    "is_random": True,
+                    "has_changed": True,
+                    "seed": 42,
+                    "shuffle": True,
+                    "stratify": False,
+                }
+            ),
         )
         db.add(experiment)
         db.commit()
diff --git a/tests/back/api/test_runs_api.py b/tests/back/api/test_runs_api.py
index 64bcae10..1ea987b1 100644
--- a/tests/back/api/test_runs_api.py
+++ b/tests/back/api/test_runs_api.py
@@ -1,44 +1,66 @@
 import json
+import os
 
 import pytest
 from fastapi.testclient import TestClient
 
-from DashAI.back.dependencies.database.models import Dataset, Experiment
 
-
-@pytest.fixture(scope="module", name="experiment_id")
-def create_experiment(client: TestClient):
-    session = client.app.container.db.provided().session
-
-    with session() as db:
-        # create Dummy Dataset
-        dummy_dataset = Dataset(
-            name="DummyDataset",
-            task_name="TabularClassificationTask",
-            file_path="dummy.csv",
-            feature_names=json.dumps([]),
-        )
-        db.add(dummy_dataset)
-        db.commit()
-        db.refresh(dummy_dataset)
-
-        # Create Dummy Experiment
-        dummy_experiment = Experiment(
-            dataset_id=dummy_dataset.id,
-            task_name="TabularClassificationTask",
-            name="Test Experiment",
+@pytest.fixture(scope="module", name="dataset_id")
+def create_dataset(client):
+    """Create testing dataset 1."""
+    abs_file_path = os.path.join(os.path.dirname(__file__), "iris.csv")
+
+    with open(abs_file_path, "rb") as csv:
+        response = client.post(
+            "/api/v1/dataset/",
+            data={
+                "params": """{  "dataloader": "CSVDataLoader",
+                                    "dataset_name": "DummyDataset2",
+                                    "splits_in_folders": false,
+                                    "splits": {
+                                        "train_size": 0.8,
+                                        "test_size": 0.1,
+                                        "val_size": 0.1,
+                                        "seed": 42,
+                                        "shuffle": true
+                                    },
+                                    "dataloader_params": {
+                                        "separator": ","
+                                    }
+                                }""",
+                "url": "",
+            },
+            files={"file": ("filename", csv, "text/csv")},
         )
+    return response.json()["id"]
 
-        db.add(dummy_experiment)
-        db.commit()
-        db.refresh(dummy_experiment)
-
-        yield dummy_experiment.id
 
-        # Delete the dataset and experiment
-        db.delete(dummy_experiment)
-        db.delete(dummy_dataset)
-        db.commit()
+@pytest.fixture(scope="module", name="experiment_id")
+def create_experiment(client: TestClient, dataset_id):
+    """Create experiment 1."""
+    response = client.post(
+        "/api/v1/experiment/",
+        json={
+            "dataset_id": dataset_id,
+            "task_name": "TabularClassificationTask",
+            "name": "Test Experiment",
+            "input_columns": [1, 2, 3, 4],
+            "output_columns": [5],
+            "splits": json.dumps(
+                {
+                    "train": 0.5,
+                    "test": 0.2,
+                    "validation": 0.3,
+                    "is_random": True,
+                    "has_changed": True,
+                    "seed": 42,
+                    "shuffle": True,
+                    "stratify": False,
+                }
+            ),
+        },
+    )
+    return response.json()["id"]
 
 
 def test_create_run(client: TestClient, experiment_id: int):
diff --git a/tests/back/converters/test_drop_column_converters.py b/tests/back/converters/test_drop_column_converters.py
index 074ead4d..bf7db8a9 100644
--- a/tests/back/converters/test_drop_column_converters.py
+++ b/tests/back/converters/test_drop_column_converters.py
@@ -7,8 +7,10 @@
 from DashAI.back.converters.column_dropper_by_index import ColumnDropperByIndex
 from DashAI.back.converters.column_dropper_by_name import ColumnDropperByName
 from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader
-from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    DashAIDataset,
+    to_dashai_dataset,
+)
 
 
 @pytest.fixture(name="iris_dataset")
@@ -26,18 +28,7 @@ def prepare_iris_dataset():
         params={"separator": ","},
     )
 
-    inputs_columns = [
-        "SepalLengthCm",
-        "SepalWidthCm",
-        "PetalLengthCm",
-        "PetalWidthCm",
-    ]
-
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns,
-        outputs_columns=["Species"],
-    )
+    datasetdict = to_dashai_dataset(datasetdict)
 
     return datasetdict
 
@@ -57,13 +48,7 @@ def prepare_iris_petal_width_dropped_dataset():
         params={"separator": ","},
     )
 
-    inputs_columns = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm"]
-
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns,
-        outputs_columns=["Species"],
-    )
+    datasetdict = to_dashai_dataset(datasetdict)
 
     return datasetdict
 
@@ -83,13 +68,7 @@ def prepare_iris_dataset_petal_cols_dropped():
         params={"separator": ","},
     )
 
-    inputs_columns = ["SepalLengthCm", "SepalWidthCm"]
-
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns,
-        outputs_columns=["Species"],
-    )
+    datasetdict = to_dashai_dataset(datasetdict)
 
     return datasetdict
 
@@ -98,7 +77,6 @@ def test_remove_input_column_with_column_name(
     iris_dataset: DatasetDict, iris_dataset_petal_width_dropped: DatasetDict
 ):
     dropper = ColumnDropperByName(column_names="PetalWidthCm")
-    print(iris_dataset)
     dataset_obtained = dropper.transform(iris_dataset)
     assert set(dataset_obtained.keys()) == set(iris_dataset_petal_width_dropped.keys())
     for split in dataset_obtained:
diff --git a/tests/back/dataloaders/iris_species_twice.csv b/tests/back/dataloaders/iris_species_twice.csv
new file mode 100644
index 00000000..2c286fed
--- /dev/null
+++ b/tests/back/dataloaders/iris_species_twice.csv
@@ -0,0 +1,151 @@
+SepalLengthCm,SepalWidthCm,PetalLengthCm,PetalWidthCm,Species,Species-2
+5.1,3.5,1.4,0.2,Iris-setosa,setosa
+4.9,3,1.4,0.2,Iris-setosa,setosa
+4.7,3.2,1.3,0.2,Iris-setosa,setosa
+4.6,3.1,1.5,0.2,Iris-setosa,setosa
+5,3.6,1.4,0.2,Iris-setosa,setosa
+5.4,3.9,1.7,0.4,Iris-setosa,setosa
+4.6,3.4,1.4,0.3,Iris-setosa,setosa
+5,3.4,1.5,0.2,Iris-setosa,setosa
+4.4,2.9,1.4,0.2,Iris-setosa,setosa
+4.9,3.1,1.5,0.1,Iris-setosa,setosa
+5.4,3.7,1.5,0.2,Iris-setosa,setosa
+4.8,3.4,1.6,0.2,Iris-setosa,setosa
+4.8,3,1.4,0.1,Iris-setosa,setosa
+4.3,3,1.1,0.1,Iris-setosa,setosa
+5.8,4,1.2,0.2,Iris-setosa,setosa
+5.7,4.4,1.5,0.4,Iris-setosa,setosa
+5.4,3.9,1.3,0.4,Iris-setosa,setosa
+5.1,3.5,1.4,0.3,Iris-setosa,setosa
+5.7,3.8,1.7,0.3,Iris-setosa,setosa
+5.1,3.8,1.5,0.3,Iris-setosa,setosa
+5.4,3.4,1.7,0.2,Iris-setosa,setosa
+5.1,3.7,1.5,0.4,Iris-setosa,setosa
+4.6,3.6,1,0.2,Iris-setosa,setosa
+5.1,3.3,1.7,0.5,Iris-setosa,setosa
+4.8,3.4,1.9,0.2,Iris-setosa,setosa
+5,3,1.6,0.2,Iris-setosa,setosa
+5,3.4,1.6,0.4,Iris-setosa,setosa
+5.2,3.5,1.5,0.2,Iris-setosa,setosa
+5.2,3.4,1.4,0.2,Iris-setosa,setosa
+4.7,3.2,1.6,0.2,Iris-setosa,setosa
+4.8,3.1,1.6,0.2,Iris-setosa,setosa
+5.4,3.4,1.5,0.4,Iris-setosa,setosa
+5.2,4.1,1.5,0.1,Iris-setosa,setosa
+5.5,4.2,1.4,0.2,Iris-setosa,setosa
+4.9,3.1,1.5,0.1,Iris-setosa,setosa
+5,3.2,1.2,0.2,Iris-setosa,setosa
+5.5,3.5,1.3,0.2,Iris-setosa,setosa
+4.9,3.1,1.5,0.1,Iris-setosa,setosa
+4.4,3,1.3,0.2,Iris-setosa,setosa
+5.1,3.4,1.5,0.2,Iris-setosa,setosa
+5,3.5,1.3,0.3,Iris-setosa,setosa
+4.5,2.3,1.3,0.3,Iris-setosa,setosa
+4.4,3.2,1.3,0.2,Iris-setosa,setosa
+5,3.5,1.6,0.6,Iris-setosa,setosa
+5.1,3.8,1.9,0.4,Iris-setosa,setosa
+4.8,3,1.4,0.3,Iris-setosa,setosa
+5.1,3.8,1.6,0.2,Iris-setosa,setosa
+4.6,3.2,1.4,0.2,Iris-setosa,setosa
+5.3,3.7,1.5,0.2,Iris-setosa,setosa
+5,3.3,1.4,0.2,Iris-setosa,setosa
+7,3.2,4.7,1.4,Iris-versicolor,versicolor
+6.4,3.2,4.5,1.5,Iris-versicolor,versicolor
+6.9,3.1,4.9,1.5,Iris-versicolor,versicolor
+5.5,2.3,4,1.3,Iris-versicolor,versicolor
+6.5,2.8,4.6,1.5,Iris-versicolor,versicolor
+5.7,2.8,4.5,1.3,Iris-versicolor,versicolor
+6.3,3.3,4.7,1.6,Iris-versicolor,versicolor
+4.9,2.4,3.3,1,Iris-versicolor,versicolor
+6.6,2.9,4.6,1.3,Iris-versicolor,versicolor
+5.2,2.7,3.9,1.4,Iris-versicolor,versicolor
+5,2,3.5,1,Iris-versicolor,versicolor
+5.9,3,4.2,1.5,Iris-versicolor,versicolor
+6,2.2,4,1,Iris-versicolor,versicolor
+6.1,2.9,4.7,1.4,Iris-versicolor,versicolor
+5.6,2.9,3.6,1.3,Iris-versicolor,versicolor
+6.7,3.1,4.4,1.4,Iris-versicolor,versicolor
+5.6,3,4.5,1.5,Iris-versicolor,versicolor
+5.8,2.7,4.1,1,Iris-versicolor,versicolor
+6.2,2.2,4.5,1.5,Iris-versicolor,versicolor
+5.6,2.5,3.9,1.1,Iris-versicolor,versicolor
+5.9,3.2,4.8,1.8,Iris-versicolor,versicolor
+6.1,2.8,4,1.3,Iris-versicolor,versicolor
+6.3,2.5,4.9,1.5,Iris-versicolor,versicolor
+6.1,2.8,4.7,1.2,Iris-versicolor,versicolor
+6.4,2.9,4.3,1.3,Iris-versicolor,versicolor
+6.6,3,4.4,1.4,Iris-versicolor,versicolor
+6.8,2.8,4.8,1.4,Iris-versicolor,versicolor
+6.7,3,5,1.7,Iris-versicolor,versicolor
+6,2.9,4.5,1.5,Iris-versicolor,versicolor
+5.7,2.6,3.5,1,Iris-versicolor,versicolor
+5.5,2.4,3.8,1.1,Iris-versicolor,versicolor
+5.5,2.4,3.7,1,Iris-versicolor,versicolor
+5.8,2.7,3.9,1.2,Iris-versicolor,versicolor
+6,2.7,5.1,1.6,Iris-versicolor,versicolor
+5.4,3,4.5,1.5,Iris-versicolor,versicolor
+6,3.4,4.5,1.6,Iris-versicolor,versicolor
+6.7,3.1,4.7,1.5,Iris-versicolor,versicolor
+6.3,2.3,4.4,1.3,Iris-versicolor,versicolor
+5.6,3,4.1,1.3,Iris-versicolor,versicolor
+5.5,2.5,4,1.3,Iris-versicolor,versicolor
+5.5,2.6,4.4,1.2,Iris-versicolor,versicolor
+6.1,3,4.6,1.4,Iris-versicolor,versicolor
+5.8,2.6,4,1.2,Iris-versicolor,versicolor
+5,2.3,3.3,1,Iris-versicolor,versicolor
+5.6,2.7,4.2,1.3,Iris-versicolor,versicolor
+5.7,3,4.2,1.2,Iris-versicolor,versicolor
+5.7,2.9,4.2,1.3,Iris-versicolor,versicolor
+6.2,2.9,4.3,1.3,Iris-versicolor,versicolor
+5.1,2.5,3,1.1,Iris-versicolor,versicolor
+5.7,2.8,4.1,1.3,Iris-versicolor,versicolor
+6.3,3.3,6,2.5,Iris-virginica,virginica
+5.8,2.7,5.1,1.9,Iris-virginica,virginica
+7.1,3,5.9,2.1,Iris-virginica,virginica
+6.3,2.9,5.6,1.8,Iris-virginica,virginica
+6.5,3,5.8,2.2,Iris-virginica,virginica
+7.6,3,6.6,2.1,Iris-virginica,virginica
+4.9,2.5,4.5,1.7,Iris-virginica,virginica
+7.3,2.9,6.3,1.8,Iris-virginica,virginica
+6.7,2.5,5.8,1.8,Iris-virginica,virginica
+7.2,3.6,6.1,2.5,Iris-virginica,virginica
+6.5,3.2,5.1,2,Iris-virginica,virginica
+6.4,2.7,5.3,1.9,Iris-virginica,virginica
+6.8,3,5.5,2.1,Iris-virginica,virginica
+5.7,2.5,5,2,Iris-virginica,virginica
+5.8,2.8,5.1,2.4,Iris-virginica,virginica
+6.4,3.2,5.3,2.3,Iris-virginica,virginica
+6.5,3,5.5,1.8,Iris-virginica,virginica
+7.7,3.8,6.7,2.2,Iris-virginica,virginica
+7.7,2.6,6.9,2.3,Iris-virginica,virginica
+6,2.2,5,1.5,Iris-virginica,virginica
+6.9,3.2,5.7,2.3,Iris-virginica,virginica
+5.6,2.8,4.9,2,Iris-virginica,virginica
+7.7,2.8,6.7,2,Iris-virginica,virginica
+6.3,2.7,4.9,1.8,Iris-virginica,virginica
+6.7,3.3,5.7,2.1,Iris-virginica,virginica
+7.2,3.2,6,1.8,Iris-virginica,virginica
+6.2,2.8,4.8,1.8,Iris-virginica,virginica
+6.1,3,4.9,1.8,Iris-virginica,virginica
+6.4,2.8,5.6,2.1,Iris-virginica,virginica
+7.2,3,5.8,1.6,Iris-virginica,virginica
+7.4,2.8,6.1,1.9,Iris-virginica,virginica
+7.9,3.8,6.4,2,Iris-virginica,virginica
+6.4,2.8,5.6,2.2,Iris-virginica,virginica
+6.3,2.8,5.1,1.5,Iris-virginica,virginica
+6.1,2.6,5.6,1.4,Iris-virginica,virginica
+7.7,3,6.1,2.3,Iris-virginica,virginica
+6.3,3.4,5.6,2.4,Iris-virginica,virginica
+6.4,3.1,5.5,1.8,Iris-virginica,virginica
+6,3,4.8,1.8,Iris-virginica,virginica
+6.9,3.1,5.4,2.1,Iris-virginica,virginica
+6.7,3.1,5.6,2.4,Iris-virginica,virginica
+6.9,3.1,5.1,2.3,Iris-virginica,virginica
+5.8,2.7,5.1,1.9,Iris-virginica,virginica
+6.8,3.2,5.9,2.3,Iris-virginica,virginica
+6.7,3.3,5.7,2.5,Iris-virginica,virginica
+6.7,3,5.2,2.3,Iris-virginica,virginica
+6.3,2.5,5,1.9,Iris-virginica,virginica
+6.5,3,5.2,2,Iris-virginica,virginica
+6.2,3.4,5.4,2.3,Iris-virginica,virginica
+5.9,3,5.1,1.8,Iris-virginica,virginica
diff --git a/tests/back/dataloaders/test_dashai_dataset.py b/tests/back/dataloaders/test_dashai_dataset.py
index a9f561ec..83ca7c55 100644
--- a/tests/back/dataloaders/test_dashai_dataset.py
+++ b/tests/back/dataloaders/test_dashai_dataset.py
@@ -2,17 +2,33 @@
 import shutil
 
 import pytest
-from datasets import DatasetDict
+from datasets import ClassLabel, DatasetDict
 from pyarrow.lib import ArrowInvalid
 from starlette.datastructures import UploadFile
 
+from DashAI.back.api.api_v1.schemas.datasets_params import ColumnSpecItemParams
 from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader
 from DashAI.back.dataloaders.classes.dashai_dataset import (
     DashAIDataset,
     load_dataset,
+    parse_columns_indices,
     save_dataset,
+    select_columns,
+    split_dataset,
+    split_indexes,
+    to_dashai_dataset,
+    update_columns_spec,
+    update_dataset_splits,
+    validate_inputs_outputs,
 )
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
+
+
+@pytest.fixture(scope="module", autouse=True)
+def clean():
+    shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
+    yield
+    shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
+    return True
 
 
 @pytest.fixture(scope="module", name="dataset_created")
@@ -35,15 +51,18 @@ def fixture_dataset():
     return datasetdict
 
 
-def test_inputs_outputs_columns(dataset_created: DatasetDict):
-    inputs_columns = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]
+def test_validate_empty_inputs_outputs_columns(dataset_created: DatasetDict):
+    inputs_colums = []
     outputs_columns = ["Species"]
-    datasetdict = to_dashai_dataset(dataset_created, inputs_columns, outputs_columns)
-    assert datasetdict["train"].inputs_columns == inputs_columns
-    assert datasetdict["train"].outputs_columns == outputs_columns
 
+    with pytest.raises(
+        ValueError,
+        match="Inputs and outputs columns lists to validate must not be empty",
+    ):
+        validate_inputs_outputs(dataset_created, inputs_colums, outputs_columns)
 
-def test_wrong_size_inputs_outputs_columns(dataset_created: DatasetDict):
+
+def test_validate_wrong_size_inputs_outputs_columns(dataset_created: DatasetDict):
     inputs_columns = [
         "SepalLengthCm",
         "SepalWidthCm",
@@ -58,30 +77,20 @@ def test_wrong_size_inputs_outputs_columns(dataset_created: DatasetDict):
             r"inputs: 4, number of outputs: 2, number of names: 5. "
         ),
     ):
-        to_dashai_dataset(dataset_created, inputs_columns, outputs_columns)
+        validate_inputs_outputs(dataset_created, inputs_columns, outputs_columns)
 
 
-def test_undefined_inputs_outputs_columns(dataset_created: DatasetDict):
-    inputs_columns = ["SepalLengthCm", "SepalWidthCm", "PetalWidthCm"]
+def test_validate_wrong_name_outputs_columns(dataset_created: DatasetDict):
+    inputs_columns = ["Sepal", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]
     outputs_columns = ["Species"]
     with pytest.raises(
         ValueError,
         match=(
-            r"The union of the elements of inputs and outputs list must be equal "
-            r"to elements in the list of names."
+            r"Inputs and outputs can only contain elements that exist in names. "
+            r"Extra elements: Sepal"
         ),
     ):
-        to_dashai_dataset(dataset_created, inputs_columns, outputs_columns)
-
-
-def test_wrong_name_outputs_columns(dataset_created: DatasetDict):
-    inputs_columns = ["Sepal", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]
-    outputs_columns = ["Species"]
-    with pytest.raises(
-        ValueError,
-        match=r"Inputs and outputs can only contain elements that exist in names.",
-    ):
-        to_dashai_dataset(dataset_created, inputs_columns, outputs_columns)
+        validate_inputs_outputs(dataset_created, inputs_columns, outputs_columns)
 
 
 @pytest.fixture(scope="module", name="dashaidataset_created")
@@ -100,20 +109,18 @@ def fixture_dashaidataset():
         temp_path="tests/back/dataloaders",
         params=params,
     )
-    inputs_columns = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]
-    outputs_columns = ["Species"]
 
-    datasetdict = to_dashai_dataset(datasetdict, inputs_columns, outputs_columns)
+    datasetdict = to_dashai_dataset(datasetdict)
 
-    return [datasetdict, csv_dataloader]
+    return datasetdict
 
 
 def test_dashaidataset_sample(dashaidataset_created: list):
     methods = ["head", "tail", "random"]
     n_samples = [1, 10]
 
-    for split in dashaidataset_created[0]:
-        dataset = dashaidataset_created[0][split]
+    for split in dashaidataset_created:
+        dataset = dashaidataset_created[split]
 
         for n in n_samples:
             for method in methods:
@@ -139,7 +146,7 @@ def test_dashaidataset_sample(dashaidataset_created: list):
 def test_wrong_name_column(dashaidataset_created: list):
     col_types = {"Speci": "Categorical"}
 
-    for split in dashaidataset_created[0]:
+    for split in dashaidataset_created:
         with pytest.raises(
             ValueError,
             match=(
@@ -147,7 +154,7 @@ def test_wrong_name_column(dashaidataset_created: list):
                 r"exist in dataset."
             ),
         ):
-            dashaidataset_created[0][split] = dashaidataset_created[0][
+            dashaidataset_created[split] = dashaidataset_created[
                 split
             ].change_columns_type(col_types)
 
@@ -155,44 +162,46 @@ def test_wrong_name_column(dashaidataset_created: list):
 def test_wrong_type_column(dashaidataset_created: list):
     col_types = {"Species": "Numerical"}
 
-    for split in dashaidataset_created[0]:
+    for split in dashaidataset_created:
         with pytest.raises(ArrowInvalid):
-            dashaidataset_created[0][split] = dashaidataset_created[0][
+            dashaidataset_created[split] = dashaidataset_created[
                 split
             ].change_columns_type(col_types)
 
 
 def test_dashaidataset_after_cast(dashaidataset_created: DatasetDict):
-    inputs_columns = dashaidataset_created[0]["train"].inputs_columns
+    features = dashaidataset_created["train"].features.copy()
+    features["Species"] = ClassLabel(
+        names=list(set(dashaidataset_created["train"]["Species"]))
+    )
+
     col_types = {"Species": "Categorical"}
-    for split in dashaidataset_created[0]:
-        dashaidataset_created[0][split] = dashaidataset_created[0][
-            split
-        ].change_columns_type(col_types)
-    assert dashaidataset_created[0]["train"].inputs_columns == inputs_columns
+    for split in dashaidataset_created:
+        dashaidataset_created[split] = dashaidataset_created[split].change_columns_type(
+            col_types
+        )
+    assert dashaidataset_created["train"].features == features
 
 
 def test_split_dataset(dashaidataset_created: list):
-    inputs_columns = dashaidataset_created[0]["train"].inputs_columns
-    outputs_columns = dashaidataset_created[0]["train"].outputs_columns
-    totals_rows = dashaidataset_created[0]["train"].num_rows
-    separate_datasetdict = dashaidataset_created[1].split_dataset(
-        dashaidataset_created[0], 0.7, 0.1, 0.2, class_column=outputs_columns[0]
+    totals_rows = dashaidataset_created["train"].num_rows
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=totals_rows, train_size=0.7, test_size=0.1, val_size=0.2
+    )
+    separate_datasetdict = split_dataset(
+        dashaidataset_created["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
     )
 
     train_rows = separate_datasetdict["train"].num_rows
     test_rows = separate_datasetdict["test"].num_rows
     validation_rows = separate_datasetdict["validation"].num_rows
-    assert separate_datasetdict["train"].inputs_columns == inputs_columns
-    assert separate_datasetdict["test"].inputs_columns == inputs_columns
-    assert separate_datasetdict["validation"].inputs_columns == inputs_columns
-    assert separate_datasetdict["train"].outputs_columns == outputs_columns
-    assert separate_datasetdict["test"].outputs_columns == outputs_columns
-    assert separate_datasetdict["validation"].outputs_columns == outputs_columns
     assert totals_rows == train_rows + test_rows + validation_rows
 
 
-def split_dataset():
+def split_iris_dataset():
     test_dataset_path = "tests/back/dataloaders/iris.csv"
     dataloader_test = CSVDataLoader()
 
@@ -207,72 +216,270 @@ def split_dataset():
         params={"separator": ","},
     )
 
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species"],
+    datasetdict = to_dashai_dataset(datasetdict)
+
+    total_rows = len(datasetdict["train"])
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=total_rows, train_size=0.7, test_size=0.1, val_size=0.2
     )
-    separate_datasetdict = dataloader_test.split_dataset(
-        datasetdict,
-        train_size=0.7,
-        test_size=0.1,
-        val_size=0.2,
-        class_column=datasetdict["train"].outputs_columns[0],
+    separate_datasetdict = split_dataset(
+        datasetdict["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
     )
 
     return separate_datasetdict
 
 
 def test_save_to_disk_and_load():
-    dataset = split_dataset()
-    inputs_columns = dataset["train"].inputs_columns
-    outputs_columns = dataset["train"].outputs_columns
-
+    dataset = split_iris_dataset()
+    feature_names = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+        "Species",
+    ]
     save_dataset(dataset, "tests/back/dataloaders/dashaidataset")
     dashai_datasetdict = load_dataset("tests/back/dataloaders/dashaidataset")
     shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
 
-    assert dashai_datasetdict["train"].inputs_columns == inputs_columns
-    assert dashai_datasetdict["test"].inputs_columns == inputs_columns
-    assert dashai_datasetdict["validation"].inputs_columns == inputs_columns
-    assert dashai_datasetdict["train"].outputs_columns == outputs_columns
-    assert dashai_datasetdict["validation"].outputs_columns == outputs_columns
-    assert dashai_datasetdict["test"].outputs_columns == outputs_columns
+    assert list((dashai_datasetdict["train"].features).keys()) == feature_names
+    assert list((dashai_datasetdict["test"].features).keys()) == feature_names
+    assert list((dashai_datasetdict["validation"].features).keys()) == feature_names
 
 
-@pytest.fixture(name="iris_dataset")
-def prepare_iris_dataset():
-    test_dataset_path = "tests/back/dataloaders/iris.csv"
+def test_parse_columns_indices():
+    input_columns_indices = [1, 3, 4]
+    input_columns_names = ["SepalLengthCm", "PetalLengthCm", "PetalWidthCm"]
+    dataset = split_iris_dataset()
+    feature_names1 = parse_columns_indices(dataset, input_columns_indices)
+
+    assert feature_names1 == input_columns_names
+
+
+def test_parse_columns_indices_wrong_index():
+    input_columns_indices = [1, 3, 6]
+
+    dataset = split_iris_dataset()
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            r"The list of indices can only contain elements within the amount "
+            r"of columns. Index 6 is greater than the total of columns."
+        ),
+    ):
+        parse_columns_indices(dataset, input_columns_indices)
+
+
+def test_select_columns():
+    inputs_columns = [
+        "SepalLengthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+    outputs_columns = ["Species"]
+    dataset = split_iris_dataset()
+
+    train_rows = dataset["train"].num_rows
+    validation_rows = dataset["validation"].num_rows
+    test_rows = dataset["test"].num_rows
+
+    x, y = select_columns(dataset, inputs_columns, outputs_columns)
+
+    assert x["train"].shape == (train_rows, len(inputs_columns))
+    assert x["validation"].shape == (
+        validation_rows,
+        len(inputs_columns),
+    )
+    assert x["test"].shape == (test_rows, len(inputs_columns))
+    assert y["train"].shape == (train_rows, len(outputs_columns))
+    assert y["validation"].shape == (
+        validation_rows,
+        len(outputs_columns),
+    )
+    assert y["test"].shape == (test_rows, len(outputs_columns))
+
+
+def test_update_columns_spec_valid():
+    dataset = split_iris_dataset()
+    modify_data = {
+        "SepalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "SepalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "PetalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "PetalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "Species": ColumnSpecItemParams(type="Value", dtype="string"),
+    }
+
+    save_dataset(dataset, "tests/back/dataloaders/dashaidataset")
+    updated_dataset = update_columns_spec(
+        "tests/back/dataloaders/dashaidataset", modify_data
+    )
+    shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
+
+    new_features = updated_dataset["train"].features
+    assert list(new_features.keys()) == [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+        "Species",
+    ]
+    assert new_features["SepalLengthCm"]._type == "Value"
+    assert new_features["SepalLengthCm"].dtype == "string"
+    assert new_features["SepalWidthCm"]._type == "Value"
+    assert new_features["SepalWidthCm"].dtype == "float64"
+    assert new_features["PetalLengthCm"]._type == "Value"
+    assert new_features["PetalLengthCm"].dtype == "string"
+    assert new_features["PetalWidthCm"]._type == "Value"
+    assert new_features["PetalWidthCm"].dtype == "float64"
+    assert new_features["Species"]._type == "Value"
+    assert new_features["Species"].dtype == "string"
+
+
+def test_update_columns_spec_unsoported_input():
+    dataset = split_iris_dataset()
+    modify_data = {
+        "SepalLengthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "SepalWidthCm": ColumnSpecItemParams(type="Value", dtype="bool"),
+        "PetalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "PetalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "Species": ColumnSpecItemParams(type="Value", dtype="bool"),
+    }
+
+    save_dataset(dataset, "tests/back/dataloaders/dashaidataset")
+    with pytest.raises(
+        ValueError,
+        match=("Error while trying to cast the columns"),
+    ):
+        update_columns_spec("tests/back/dataloaders/dashaidataset", modify_data)
+
+    shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
+
+
+def test_update_columns_spec_one_class_column():
+    dataset = split_iris_dataset()
+    modify_data = {
+        "SepalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "SepalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "PetalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "PetalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "Species": ColumnSpecItemParams(type="ClassLabel", dtype=""),
+    }
+
+    save_dataset(dataset, "tests/back/dataloaders/dashaidataset")
+    updated_dataset = update_columns_spec(
+        "tests/back/dataloaders/dashaidataset", modify_data
+    )
+    shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
+
+    new_features = updated_dataset["train"].features
+    assert list(new_features.keys()) == [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+        "Species",
+    ]
+    assert new_features["SepalLengthCm"]._type == "Value"
+    assert new_features["SepalLengthCm"].dtype == "string"
+    assert new_features["SepalWidthCm"]._type == "Value"
+    assert new_features["SepalWidthCm"].dtype == "float64"
+    assert new_features["PetalLengthCm"]._type == "Value"
+    assert new_features["PetalLengthCm"].dtype == "string"
+    assert new_features["PetalWidthCm"]._type == "Value"
+    assert new_features["PetalWidthCm"].dtype == "float64"
+    assert new_features["Species"]._type == "ClassLabel"
+
+
+def split_dataset_with_two_classes():
+    test_dataset_path = "tests/back/dataloaders/iris_species_twice.csv"
     dataloader_test = CSVDataLoader()
 
     with open(test_dataset_path, "r") as file:
-        csv_binary = io.BytesIO(bytes(file.read(), encoding="utf8"))
+        csv_data = file.read()
+        csv_binary = io.BytesIO(bytes(csv_data, encoding="utf8"))
         file = UploadFile(csv_binary)
 
     datasetdict = dataloader_test.load_data(
         filepath_or_buffer=file,
-        temp_path="tests/back/converters",
+        temp_path="tests/back/dataloaders",
         params={"separator": ","},
     )
 
-    inputs_columns = [
+    datasetdict = to_dashai_dataset(datasetdict)
+
+    total_rows = len(datasetdict["train"])
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=total_rows, train_size=0.7, test_size=0.1, val_size=0.2
+    )
+    separate_datasetdict = split_dataset(
+        datasetdict["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
+    )
+
+    return separate_datasetdict
+
+
+def test_update_columns_spec_multiple_class_columns():
+    dataset = split_dataset_with_two_classes()
+    modify_data = {
+        "SepalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "SepalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "PetalLengthCm": ColumnSpecItemParams(type="Value", dtype="string"),
+        "PetalWidthCm": ColumnSpecItemParams(type="Value", dtype="float64"),
+        "Species": ColumnSpecItemParams(type="ClassLabel", dtype=""),
+        "Species-2": ColumnSpecItemParams(type="ClassLabel", dtype=""),
+    }
+
+    save_dataset(dataset, "tests/back/dataloaders/dashaidataset")
+    updated_dataset = update_columns_spec(
+        "tests/back/dataloaders/dashaidataset", modify_data
+    )
+    shutil.rmtree("tests/back/dataloaders/dashaidataset", ignore_errors=True)
+
+    new_features = updated_dataset["train"].features
+    assert list(new_features.keys()) == [
         "SepalLengthCm",
         "SepalWidthCm",
         "PetalLengthCm",
         "PetalWidthCm",
+        "Species",
+        "Species-2",
     ]
+    assert new_features["SepalLengthCm"]._type == "Value"
+    assert new_features["SepalLengthCm"].dtype == "string"
+    assert new_features["SepalWidthCm"]._type == "Value"
+    assert new_features["SepalWidthCm"].dtype == "float64"
+    assert new_features["PetalLengthCm"]._type == "Value"
+    assert new_features["PetalLengthCm"].dtype == "string"
+    assert new_features["PetalWidthCm"]._type == "Value"
+    assert new_features["PetalWidthCm"].dtype == "float64"
+    assert new_features["Species"]._type == "ClassLabel"
+    assert new_features["Species-2"]._type == "ClassLabel"
+
+
+@pytest.fixture(name="iris_dataset")
+def prepare_iris_dataset():
+    test_dataset_path = "tests/back/dataloaders/iris.csv"
+    dataloader_test = CSVDataLoader()
+
+    with open(test_dataset_path, "r") as file:
+        csv_binary = io.BytesIO(bytes(file.read(), encoding="utf8"))
+        file = UploadFile(csv_binary)
 
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns,
-        outputs_columns=["Species"],
+    datasetdict = dataloader_test.load_data(
+        filepath_or_buffer=file,
+        temp_path="tests/back/converters",
+        params={"separator": ","},
     )
 
+    datasetdict = to_dashai_dataset(datasetdict)
+
     return datasetdict
 
 
@@ -291,13 +498,7 @@ def prepare_iris_petal_width_dropped_dataset():
         params={"separator": ","},
     )
 
-    inputs_columns = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm"]
-
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns,
-        outputs_columns=["Species"],
-    )
+    datasetdict = to_dashai_dataset(datasetdict)
 
     return datasetdict
 
@@ -317,5 +518,41 @@ def test_remove_columns(
     for column_name in train_split.column_names:
         assert train_split[column_name] == train_dropped_split[column_name]
     assert train_split.features == train_dropped_split.features
-    assert train_split.inputs_columns == train_dropped_split.inputs_columns
-    assert train_split.outputs_columns == train_dropped_split.outputs_columns
+
+
+def test_update_splits_by_percentage():
+    iris_dataset = split_iris_dataset()
+    n = (
+        len(iris_dataset["train"])
+        + len(iris_dataset["test"])
+        + len(iris_dataset["validation"])
+    )
+    new_splits = {"train": 0.5, "test": 0.3, "validation": 0.2}
+
+    new_dataset = update_dataset_splits(
+        iris_dataset, new_splits=new_splits, is_random=True
+    )
+
+    assert len(new_dataset["train"]) == n * 0.5
+    assert len(new_dataset["test"]) == n * 0.3
+    assert len(new_dataset["validation"]) == n * 0.2
+
+
+def test_update_splits_by_specific_rows():
+    iris_dataset = split_iris_dataset()
+    train_indexes = list(range(0, 100))
+    test_indexes = list(range(100, 130))
+    val_indexes = list(range(130, 150))
+    new_splits = {
+        "train": train_indexes,
+        "test": test_indexes,
+        "validation": val_indexes,
+    }
+
+    new_dataset = update_dataset_splits(
+        iris_dataset, new_splits=new_splits, is_random=False
+    )
+
+    assert len(new_dataset["train"]) == 100
+    assert len(new_dataset["test"]) == 30
+    assert len(new_dataset["validation"]) == 20
diff --git a/tests/back/metrics/test_classification_metrics.py b/tests/back/metrics/test_classification_metrics.py
index 1997842e..03f66781 100644
--- a/tests/back/metrics/test_classification_metrics.py
+++ b/tests/back/metrics/test_classification_metrics.py
@@ -1,156 +1,82 @@
-import io
-import os
-from typing import Tuple
+from typing import Dict, List
 
+import numpy as np
 import pytest
-from datasets import DatasetDict
-from starlette.datastructures import UploadFile
+from datasets import Dataset
 
-from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
 from DashAI.back.metrics.classification.accuracy import Accuracy
 from DashAI.back.metrics.classification.f1 import F1
 from DashAI.back.metrics.classification.precision import Precision
 from DashAI.back.metrics.classification.recall import Recall
-from DashAI.back.models.scikit_learn.random_forest_classifier import (
-    RandomForestClassifier,
-)
-from DashAI.back.tasks.tabular_classification_task import TabularClassificationTask
-
-
-@pytest.fixture(scope="module")
-def classification_metrics_fixture():
-    test_dataset_path = "tests/back/metrics/iris.csv"
-
-    csv_dataloader = CSVDataLoader()
-    with open(test_dataset_path, "r") as file:
-        csv_data = file.read()
-        csv_binary = io.BytesIO(bytes(csv_data, encoding="utf8"))
-        file = UploadFile(csv_binary)
-
-    dataset_dict = csv_dataloader.load_data(
-        filepath_or_buffer=file,
-        temp_path="tests/back/metrics",
-        params={"separator": ","},
-    )
-    dashai_dataset = to_dashai_dataset(
-        dataset_dict,
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species"],
-    )
-
-    outputs_columns = dashai_dataset["train"].outputs_columns
-    dataset = csv_dataloader.split_dataset(
-        dashai_dataset, 0.7, 0.1, 0.2, class_column=outputs_columns[0]
-    )
-
-    tabular_task = TabularClassificationTask()
 
-    name_datasetdict = "Iris"
-    dataset = tabular_task.prepare_for_task(dataset)
-    tabular_task.validate_dataset_for_task(dataset, name_datasetdict)
 
-    model = RandomForestClassifier()
-    model.fit(dataset["train"])
-    model.save("tests/back/metrics/rf_model")
+@pytest.fixture(scope="module", name="metric_input")
+def fixture_metric_input() -> dict:
+    return {
+        "true_labels": Dataset.from_list([{"foo": 1}, {"foo": 2}, {"foo": 0}]),
+        "pred_labels": np.array([[0.1, 0.1, 0.8], [0.0, 0.8, 0.2], [0.9, 0.0, 0.5]]),
+        "wrong_size_labels": np.array([[0.1, 0.1, 0.8], [0.9, 0.0, 0.5]]),
+    }
 
-    yield dataset, model
 
-    os.remove("tests/back/metrics/rf_model")
-
-
-def test_accuracy(
-    classification_metrics_fixture: Tuple[DatasetDict, RandomForestClassifier]
-):
-    dataset, model = classification_metrics_fixture
-    y_pred = model.predict(dataset["test"])
-    score = Accuracy.score(dataset["test"], y_pred)
+def test_accuracy(metric_input: Dict[str, List[int]]):
+    score = Accuracy.score(metric_input["true_labels"], metric_input["pred_labels"])
 
     assert isinstance(score, float)
     assert score >= 0.0
     assert score <= 1.0
 
 
-def test_precision(
-    classification_metrics_fixture: Tuple[DatasetDict, RandomForestClassifier]
-):
-    dataset, model = classification_metrics_fixture
-    y_pred = model.predict(dataset["test"])
-
-    score = Precision.score(dataset["test"], y_pred)
+def test_precision(metric_input: Dict[str, List[int]]):
+    score = Precision.score(metric_input["true_labels"], metric_input["pred_labels"])
 
     assert isinstance(score, float)
     assert score >= 0.0
     assert score <= 1.0
 
 
-def test_recall(
-    classification_metrics_fixture: Tuple[DatasetDict, RandomForestClassifier]
-):
-    dataset, model = classification_metrics_fixture
-    y_pred = model.predict(dataset["test"])
-    score = Recall.score(dataset["test"], y_pred)
+def test_recall(metric_input: Dict[str, List[int]]):
+    score = Recall.score(metric_input["true_labels"], metric_input["pred_labels"])
 
     assert isinstance(score, float)
     assert score >= 0.0
     assert score <= 1.0
 
 
-def test_f1_score(
-    classification_metrics_fixture: Tuple[DatasetDict, RandomForestClassifier]
-):
-    dataset, model = classification_metrics_fixture
-    y_pred = model.predict(dataset["test"])
-    score = F1.score(dataset["test"], y_pred)
+def test_f1_score(metric_input: Dict[str, List[int]]):
+    score = F1.score(metric_input["true_labels"], metric_input["pred_labels"])
 
     assert isinstance(score, float)
     assert score >= 0.0
     assert score <= 1.0
 
 
-def test_metrics_different_input_sizes(
-    classification_metrics_fixture: Tuple[DatasetDict, RandomForestClassifier]
-):
-    dataset, model = classification_metrics_fixture
-    y_pred = model.predict(dataset["validation"])
+def test_metrics_different_input_sizes(metric_input: Dict[str, List[int]]):
+    error_pattern = (
+        r"The length of the true labels and the predicted labels must be equal, "
+        r"given: len\(true_labels\) = 3 and len\(pred_labels\) = 2\."
+    )
 
     with pytest.raises(
         ValueError,
-        match=(
-            r"The length of the true labels and the predicted labels must be equal, "
-            r"given: len\(true_labels\) = 15 and len\(pred_labels\) = 30\."
-        ),
+        match=error_pattern,
     ):
-        Accuracy.score(dataset["test"], y_pred)
+        Accuracy.score(metric_input["true_labels"], metric_input["wrong_size_labels"])
 
     with pytest.raises(
         ValueError,
-        match=(
-            r"The length of the true labels and the predicted labels must be equal, "
-            r"given: len\(true_labels\) = 15 and len\(pred_labels\) = 30\."
-        ),
+        match=error_pattern,
     ):
-        Precision.score(dataset["test"], y_pred)
+        Precision.score(metric_input["true_labels"], metric_input["wrong_size_labels"])
 
     with pytest.raises(
         ValueError,
-        match=(
-            r"The length of the true labels and the predicted labels must be equal, "
-            r"given: len\(true_labels\) = 15 and len\(pred_labels\) = 30\."
-        ),
+        match=error_pattern,
     ):
-        Recall.score(dataset["test"], y_pred)
+        Recall.score(metric_input["true_labels"], metric_input["wrong_size_labels"])
 
     with pytest.raises(
         ValueError,
-        match=(
-            r"The length of the true labels and the predicted labels must be equal, "
-            r"given: len\(true_labels\) = 15 and len\(pred_labels\) = 30\."
-        ),
+        match=error_pattern,
     ):
-        F1.score(dataset["test"], y_pred)
+        F1.score(metric_input["true_labels"], metric_input["wrong_size_labels"])
diff --git a/tests/back/metrics/test_translation_metrics.py b/tests/back/metrics/test_translation_metrics.py
index d9424190..a7bd25ff 100644
--- a/tests/back/metrics/test_translation_metrics.py
+++ b/tests/back/metrics/test_translation_metrics.py
@@ -1,75 +1,42 @@
 """Translation Metrics Tests."""
-import io
-
 import pytest
-from datasets import DatasetDict
-from starlette.datastructures import UploadFile
+from datasets import Dataset
 
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
-from DashAI.back.dataloaders.classes.json_dataloader import JSONDataLoader
 from DashAI.back.metrics.translation.bleu import Bleu
 from DashAI.back.metrics.translation.ter import Ter
-from DashAI.back.tasks.translation_task import TranslationTask
-
-
-@pytest.fixture(scope="module", name="test_dataset")
-def translation_metrics_tests_fixture():
-    test_dataset_path = "tests/back/metrics/translationEngSpaDatasetSmall.json"
-    json_dataloader = JSONDataLoader()
-
-    with open(test_dataset_path, "r", encoding="utf8") as file:
-        json_data = file.read()
-        json_binary = io.BytesIO(bytes(json_data, encoding="utf8"))
-        file = UploadFile(json_binary)
-
-    dataset_dict = json_dataloader.load_data(
-        filepath_or_buffer=file,
-        temp_path="tests/back/models",
-        params={"data_key": "data"},
-    )
-
-    dashai_dataset = to_dashai_dataset(
-        dataset_dict,
-        inputs_columns=["text"],
-        outputs_columns=["class"],
-    )
 
-    translation_task = TranslationTask()
-    dashai_dataset = translation_task.prepare_for_task(dashai_dataset)
-    translation_task.validate_dataset_for_task(dashai_dataset, "EngSpaDataset")
 
-    split_dataset = json_dataloader.split_dataset(
-        dashai_dataset, 0.7, 0.1, 0.2, seed=42
-    )
-    return split_dataset
+@pytest.fixture(scope="module", name="metric_input")
+def fixture_metric_input() -> dict:
+    return {
+        "true_sentences": Dataset.from_list(
+            [{"foo": "Tuve que hacerlo"}, {"foo": "Quiero que suspendas la pelea."}]
+        ),
+        "pred_sentences": ["Tuve que hacerlo", "Quiero que suspendas la pelea."],
+        "wrong_size_sentences": ["Tuve que hacerlo"],
+    }
 
 
-def test_bleu(test_dataset: DatasetDict):
-    pred_opus_mt_en_es = test_dataset["test"]["class"]
-    score = Bleu.score(test_dataset["test"], pred_opus_mt_en_es)
+def test_bleu(metric_input: dict):
+    score = Bleu.score(metric_input["true_sentences"], metric_input["pred_sentences"])
 
     assert isinstance(score, float)
     assert score >= 0.0
     assert score <= 1.0
 
 
-def test_ter(test_dataset: DatasetDict):
-    pred_opus_mt_en_es = test_dataset["test"]["class"]
-    score = Ter.score(test_dataset["test"], pred_opus_mt_en_es)
+def test_ter(metric_input: dict):
+    score = Ter.score(metric_input["true_sentences"], metric_input["pred_sentences"])
 
     assert isinstance(score, float)
     assert score >= 0.0
     assert score <= 1.0
 
 
-def test_metrics_different_input_sizes(test_dataset: DatasetDict):
-    pred_opus_mt_en_es = test_dataset["test"]["class"]
-    with pytest.raises(
-        ValueError, match="The length of the true and predicted labels must be equal."
-    ):
-        Bleu.score(test_dataset["validation"], pred_opus_mt_en_es)
+def test_metrics_different_input_sizes(metric_input: dict):
+    err_pattern = "The length of the true and predicted labels must be equal."
+    with pytest.raises(ValueError, match=err_pattern):
+        Bleu.score(metric_input["true_sentences"], metric_input["wrong_size_sentences"])
 
-    with pytest.raises(
-        ValueError, match="The length of the true and predicted labels must be equal."
-    ):
-        Ter.score(test_dataset["validation"], pred_opus_mt_en_es)
+    with pytest.raises(ValueError, match=err_pattern):
+        Ter.score(metric_input["true_sentences"], metric_input["wrong_size_sentences"])
diff --git a/tests/back/metrics/translationEngSpaDatasetSmall.json b/tests/back/metrics/translationEngSpaDatasetSmall.json
index 06f658ad..26cf58f3 100644
--- a/tests/back/metrics/translationEngSpaDatasetSmall.json
+++ b/tests/back/metrics/translationEngSpaDatasetSmall.json
@@ -1,1 +1,25 @@
-{"task_info": "TranslationTask", "data": [{"text": "I have no choice at all.", "class": "No tengo otra opci\u00f3n en absoluto."}, {"text": "I had to do it.", "class": "Tuve que hacerlo."}, {"text": "Is this yours?", "class": "\u00bfEs suyo esto?"}, {"text": "The family had a hard time after the war.", "class": "La familia atraves\u00f3 por un mal tiempo despu\u00e9s de la guerra."}, {"text": "I want you to call off the fight.", "class": "Quiero que suspendas la pelea."}, {"text": "His wish is to go to America.", "class": "Su deseo es ir a Am\u00e9rica."}, {"text": "Tom is Mary's neighbor.", "class": "Tom es el vecino de Mary."}, {"text": "Please visit me tomorrow.", "class": "Por favor ven a visitarme ma\u00f1ana."}, {"text": "Is Tom always like this?", "class": "\u00bfEs Tom siempre as\u00ed?"}, {"text": "We have our dinner between seven and seven-thirty.", "class": "Tuvimos nuestra cena entre las siete y siete con treinta"}, {"text": "I've been working with Tom for three months.", "class": "Llevo trabajando con Tom tres meses."}, {"text": "It's late. I have to go.", "class": "Es tarde, me tengo que ir."}, {"text": "Tom has no sense of style.", "class": "Tom no tiene gusto para la moda."}, {"text": "I regret that I wasted the money.", "class": "Me arrepiento de haber gastado el dinero."}, {"text": "Do you miss me?", "class": "\u00bfMe echas de menos?"}, {"text": "My mother is sick in bed.", "class": "Mi madre est\u00e1 en cama enferma."}, {"text": "Of course, to be valuable, an old carpet must be in good condition.", "class": "Por supuesto, para que una vieja alfombra tenga valor debe estar en buenas condiciones."}, {"text": "Can you explain it more clearly?", "class": "\u00bfPuedes explicarlo mejor?"}, {"text": "The lawyer believed in his client's innocence.", "class": "El abogado cre\u00eda en la inocencia de su cliente."}, {"text": "It's cold outside.", "class": "Hace fr\u00edo afuera."}, {"text": "The soldier took shelter in the foxhole.", "class": "El soldado se refugi\u00f3 en la trinchera."}, {"text": "Can you at least pretend you want to be here?", "class": "\u00bfPodes simular que quer\u00e9s estar aqu\u00ed?"}, {"text": "I'm not sure what's wrong. We should have heard from him by now.", "class": "No estoy seguro de qu\u00e9 est\u00e1 mal. Deber\u00edamos haber o\u00eddo de \u00e9l por ahora."}, {"text": "What instruments does Tom play?", "class": "\u00bfQu\u00e9 instrumentos toca Tom?"}, {"text": "I thought Tom would ask Mary to go with him to Boston.", "class": "Pens\u00e9 que Tom le pedir\u00eda a Mar\u00eda ir con \u00e9l a Boston."}, {"text": "Tom is broad-minded.", "class": "Tom es un tipo abierto."}, {"text": "Do me a favor by switching off that radio.", "class": "Hazme un favor y apaga esa radio."}, {"text": "Would you like some more beef?", "class": "\u00bfUn poco m\u00e1s de ternera?"}, {"text": "Is Tom trying to ruin our evening?", "class": "\u00bfEstar\u00e1 Tom tratando de cagar nuestra noche?"}, {"text": "Are you bald?", "class": "\u00bfEst\u00e1s t\u00fa calvo?"}]}
+{
+    "task_info": "TranslationTask",
+    "data": [
+        {
+            "text": "I have no choice at all.",
+            "class": "No tengo otra opci\u00f3n en absoluto."
+        },
+        {
+            "text": "I had to do it.",
+            "class": "Tuve que hacerlo."
+        },
+        {
+            "text": "Is this yours?",
+            "class": "\u00bfEs suyo esto?"
+        },
+        {
+            "text": "The family had a hard time after the war.",
+            "class": "La familia atraves\u00f3 por un mal tiempo despu\u00e9s de la guerra."
+        },
+        {
+            "text": "I want you to call off the fight.",
+            "class": "Quiero que suspendas la pelea."
+        }
+    ]
+}
diff --git a/tests/back/models/test_tabular_class_models.py b/tests/back/models/test_tabular_class_models.py
index 8467ff52..5c3ebe89 100644
--- a/tests/back/models/test_tabular_class_models.py
+++ b/tests/back/models/test_tabular_class_models.py
@@ -1,5 +1,6 @@
 import io
 import os
+from typing import Tuple
 
 import numpy as np
 import pytest
@@ -9,7 +10,12 @@
 from starlette.datastructures import UploadFile
 
 from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    select_columns,
+    split_dataset,
+    split_indexes,
+    to_dashai_dataset,
+)
 from DashAI.back.models.scikit_learn.k_neighbors_classifier import KNeighborsClassifier
 from DashAI.back.models.scikit_learn.random_forest_classifier import (
     RandomForestClassifier,
@@ -18,7 +24,7 @@
 from DashAI.back.models.scikit_learn.svc import SVC
 
 
-@pytest.fixture(scope="module", name="split_dataset")
+@pytest.fixture(scope="module", name="divided_dataset")
 def tabular_model_fixture():
     test_dataset_path = "tests/back/models/iris.csv"
     dataloader_test = CSVDataLoader()
@@ -33,37 +39,42 @@ def tabular_model_fixture():
         params={"separator": ","},
     )
 
-    datasetdict = to_dashai_dataset(
-        datasetdict,
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species"],
-    )
+    datasetdict = to_dashai_dataset(datasetdict)
 
-    split_dataset = dataloader_test.split_dataset(
-        datasetdict,
-        train_size=0.7,
-        test_size=0.1,
-        val_size=0.2,
-        class_column=datasetdict["train"].outputs_columns[0],
+    total_rows = len(datasetdict["train"])
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=total_rows, train_size=0.7, test_size=0.1, val_size=0.2
+    )
+    split_dataset_dict = split_dataset(
+        datasetdict["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
     )
 
-    return split_dataset
+    inputs_columns = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+    outputs_columns = ["Species"]
+    divided_dataset = select_columns(
+        split_dataset_dict, inputs_columns, outputs_columns
+    )
+    return divided_dataset
 
 
-def test_check_is_fitted(split_dataset: DatasetDict):
+def test_check_is_fitted(divided_dataset: Tuple[DatasetDict, DatasetDict]):
     knn_model = KNeighborsClassifier()
-    knn_model.fit(split_dataset["train"])
+
+    knn_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
 
     rf_model = RandomForestClassifier()
-    rf_model.fit(split_dataset["train"])
+    rf_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
 
     svc_model = SVC()
-    svc_model.fit(split_dataset["train"])
+    svc_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
 
     try:
         check_is_fitted(knn_model)
@@ -74,46 +85,46 @@ def test_check_is_fitted(split_dataset: DatasetDict):
         pytest.fail(f"Unexpected error in test_fit_models_tabular: {repr(e)}")
 
 
-def test_predict_tabular_models(split_dataset: DatasetDict):
+def test_predict_tabular_models(divided_dataset: Tuple[DatasetDict, DatasetDict]):
     knn_model = KNeighborsClassifier()
-    knn_model.fit(split_dataset["train"])
-    y_pred_knn = knn_model.predict(split_dataset["test"])
+    knn_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
+    y_pred_knn = knn_model.predict(divided_dataset[0]["test"])
 
     rf_model = RandomForestClassifier()
-    rf_model.fit(split_dataset["train"])
-    y_pred_rf = rf_model.predict(split_dataset["test"])
+    rf_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
+    y_pred_rf = rf_model.predict(divided_dataset[0]["test"])
 
     svc_model = SVC()
-    svc_model.fit(split_dataset["train"])
-    y_pred_svm = svc_model.predict(split_dataset["test"])
+    svc_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
+    y_pred_svm = svc_model.predict(divided_dataset[0]["test"])
 
     assert isinstance(y_pred_knn, np.ndarray)
     assert isinstance(y_pred_rf, np.ndarray)
     assert isinstance(y_pred_svm, np.ndarray)
 
-    assert split_dataset["test"].num_rows == len(y_pred_knn)
-    assert split_dataset["test"].num_rows == len(y_pred_rf)
-    assert split_dataset["test"].num_rows == len(y_pred_svm)
+    assert divided_dataset[0]["test"].num_rows == len(y_pred_knn)
+    assert divided_dataset[0]["test"].num_rows == len(y_pred_rf)
+    assert divided_dataset[0]["test"].num_rows == len(y_pred_svm)
 
 
-def test_not_fitted_model(split_dataset: DatasetDict):
+def test_not_fitted_model(divided_dataset: Tuple[DatasetDict, DatasetDict]):
     rf = RandomForestClassifier()
 
     with pytest.raises(NotFittedError):
-        rf.predict(split_dataset["test"])
+        rf.predict(divided_dataset[0]["test"])
 
 
-def test_save_and_load_model(split_dataset: DatasetDict):
+def test_save_and_load_model(divided_dataset: Tuple[DatasetDict, DatasetDict]):
     svc_model = SVC()
-    svc_model.fit(split_dataset["train"])
+    svc_model.fit(divided_dataset[0]["train"], divided_dataset[1]["train"])
 
     svc_model.save("tests/back/models/svm_model")
     loaded_model = SklearnLikeModel.load("tests/back/models/svm_model")
 
-    y_pred_svm = loaded_model.predict(split_dataset["test"])
+    y_pred_svm = loaded_model.predict(divided_dataset[0]["test"])
 
     assert isinstance(y_pred_svm, np.ndarray)
-    assert split_dataset["test"].num_rows == len(y_pred_svm)
+    assert divided_dataset[0]["test"].num_rows == len(y_pred_svm)
 
     os.remove("tests/back/models/svm_model")
 
diff --git a/tests/back/registries/test_registry.py b/tests/back/registries/test_registry.py
index e48f9f2c..637286a8 100644
--- a/tests/back/registries/test_registry.py
+++ b/tests/back/registries/test_registry.py
@@ -77,6 +77,7 @@ class NoComponent:
     "class": Component1,
     "configurable_object": True,
     "schema": {"properties": {"parameter_1": {"type": "number"}}},
+    "metadata": None,
     "description": None,
 }
 COMPONENT2_DICT = {
@@ -85,6 +86,7 @@ class NoComponent:
     "class": Component2,
     "configurable_object": True,
     "schema": {"properties": {"parameter_2": {"type": "string", "enum": ["a", "b"]}}},
+    "metadata": None,
     "description": None,
 }
 SUBCOMPONENT1_DICT = {
@@ -93,6 +95,7 @@ class NoComponent:
     "class": SubComponent1,
     "configurable_object": True,
     "schema": {"properties": {"parameter_1": {"type": "number"}}},
+    "metadata": None,
     "description": None,
 }
 COMPONENT3_DICT = {
@@ -101,6 +104,7 @@ class NoComponent:
     "class": Component3,
     "configurable_object": False,
     "schema": None,
+    "metadata": None,
     "description": "Some static component",
 }
 RELATED_COMPONENT1_DICT = {
@@ -109,6 +113,7 @@ class NoComponent:
     "class": RelatedComponent1,
     "configurable_object": False,
     "schema": None,
+    "metadata": None,
     "description": None,
 }
 RELATED_COMPONENT2_DICT = {
@@ -117,6 +122,7 @@ class NoComponent:
     "class": RelatedComponent2,
     "configurable_object": False,
     "schema": None,
+    "metadata": None,
     "description": None,
 }
 
diff --git a/tests/back/tasks/test_tasks.py b/tests/back/tasks/test_tasks.py
index dc3a2204..5eb7757c 100644
--- a/tests/back/tasks/test_tasks.py
+++ b/tests/back/tasks/test_tasks.py
@@ -6,7 +6,11 @@
 from starlette.datastructures import Headers, UploadFile
 
 from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader
-from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    split_dataset,
+    split_indexes,
+    to_dashai_dataset,
+)
 from DashAI.back.dataloaders.classes.image_dataloader import ImageDataLoader
 from DashAI.back.dataloaders.classes.json_dataloader import JSONDataLoader
 from DashAI.back.tasks.image_classification_task import ImageClassificationTask
@@ -32,40 +36,33 @@ def load_csv_into_datasetdict(file_name):
 
 
 def test_validate_tabular_task():
-    dataset = to_dashai_dataset(
-        load_csv_into_datasetdict("iris.csv"),
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species"],
-    )
+    dataset = to_dashai_dataset(load_csv_into_datasetdict("iris.csv"))
 
     for split in dataset:
         dataset[split] = dataset[split].change_columns_type(
             column_types={"Species": "Categorical"}
         )
     tabular_task = TabularClassificationTask()
-
+    inputs_columns = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+    outputs_columns = ["Species"]
     try:
-        tabular_task.validate_dataset_for_task(dataset=dataset, dataset_name="Iris")
+        tabular_task.validate_dataset_for_task(
+            dataset=dataset,
+            dataset_name="Iris",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
+        )
     except Exception as e:
         pytest.fail(f"Unexpected error in test_validate_task: {repr(e)}")
 
 
 def test_wrong_type_task():
-    dataset = to_dashai_dataset(
-        load_csv_into_datasetdict("iris_extra_feature.csv"),
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species", "StemCm"],
-    )
+    dataset = to_dashai_dataset(load_csv_into_datasetdict("iris_extra_feature.csv"))
 
     for split in dataset:
         dataset[split] = dataset[split].change_columns_type(
@@ -74,47 +71,73 @@ def test_wrong_type_task():
 
     tabular_task = TabularClassificationTask()
 
+    inputs_columns = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+    outputs_columns = ["Species", "StemCm"]
     with pytest.raises(TypeError):
         tabular_task.validate_dataset_for_task(
             dataset=dataset,
             dataset_name="Iris",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
         )
 
 
 def test_prepare_task():
-    dataset = to_dashai_dataset(
-        load_csv_into_datasetdict("iris.csv"),
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species"],
-    )
+    dataset = to_dashai_dataset(load_csv_into_datasetdict("iris.csv"))
     tabular_task = TabularClassificationTask()
-    dataset = tabular_task.prepare_for_task(dataset)
+    inputs_columns = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+    outputs_columns = ["Species"]
+    dataset = tabular_task.prepare_for_task(dataset, outputs_columns)
     try:
-        tabular_task.validate_dataset_for_task(dataset=dataset, dataset_name="Iris")
+        tabular_task.validate_dataset_for_task(
+            dataset=dataset,
+            dataset_name="Iris",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
+        )
     except Exception as e:
         pytest.fail(f"Unexpected error in test_prepare_task: {repr(e)}")
 
 
 def test_not_prepared_task():
-    dataset = to_dashai_dataset(
-        load_csv_into_datasetdict("iris.csv"),
-        inputs_columns=[
-            "SepalLengthCm",
-            "SepalWidthCm",
-            "PetalLengthCm",
-            "PetalWidthCm",
-        ],
-        outputs_columns=["Species"],
-    )
+    dataset = to_dashai_dataset(load_csv_into_datasetdict("iris.csv"))
     tabular_task = TabularClassificationTask()
+    inputs_columns = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+    outputs_columns = ["Species"]
 
     with pytest.raises(TypeError):
-        tabular_task.validate_dataset_for_task(dataset=dataset, dataset_name="Iris")
+        tabular_task.validate_dataset_for_task(
+            dataset=dataset,
+            dataset_name="Iris",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
+        )
+
+
+def test_get_tabular_class_task_metadata():
+    tabular_class_task = TabularClassificationTask()
+    metadata = tabular_class_task.get_metadata()
+
+    assert len(metadata.keys()) == 4
+    assert metadata["inputs_types"] == ["ClassLabel", "Value"]
+    assert metadata["outputs_types"] == ["ClassLabel"]
+    assert metadata["inputs_cardinality"] == "n"
+    assert metadata["outputs_cardinality"] == 1
 
 
 @pytest.fixture(scope="module", name="text_classification_dataset")
@@ -133,36 +156,51 @@ def text_classification_dataset_fixture():
         params={"data_key": "data"},
     )
 
-    dashai_dataset = to_dashai_dataset(
-        dataset,
-        inputs_columns=["text"],
-        outputs_columns=["class"],
-    )
+    dashai_dataset = to_dashai_dataset(dataset)
 
-    split_dataset = json_dataloader.split_dataset(
-        dashai_dataset,
-        train_size=0.7,
-        test_size=0.1,
-        val_size=0.2,
-        class_column=dataset["train"].outputs_columns[0],
+    total_rows = len(dashai_dataset["train"])
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=total_rows, train_size=0.7, test_size=0.1, val_size=0.2
     )
-    return split_dataset
+    split_datasetdict = split_dataset(
+        dashai_dataset["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
+    )
+
+    return split_datasetdict
 
 
 def test_validate_text_dataset(text_classification_dataset: DatasetDict):
     text_class_task = TextClassificationTask()
+    inputs_columns = ["text"]
+    outputs_columns = ["class"]
     imbd_sentiment_dataset = text_class_task.prepare_for_task(
-        text_classification_dataset
+        text_classification_dataset, outputs_columns
     )
     try:
         text_class_task.validate_dataset_for_task(
             dataset=imbd_sentiment_dataset,
             dataset_name="IMDBDataset",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
         )
     except Exception as e:
         pytest.fail(f"Unexpected error in test_validate_task: {repr(e)}")
 
 
+def test_get_text_class_task_metadata():
+    text_class_task = TextClassificationTask()
+    metadata = text_class_task.get_metadata()
+
+    assert len(metadata.keys()) == 4
+    assert metadata["inputs_types"] == ["Value"]
+    assert metadata["outputs_types"] == ["ClassLabel"]
+    assert metadata["inputs_cardinality"] == 1
+    assert metadata["outputs_cardinality"] == 1
+
+
 @pytest.fixture(scope="module", name="image_classification_dataset")
 def image_classification_dataset_fixture():
     test_dataset_path = "tests/back/tasks/beans_dataset_small.zip"
@@ -180,33 +218,53 @@ def image_classification_dataset_fixture():
             temp_path="tests/back/tasks/beans_dataset",
         )
 
-    dataset = to_dashai_dataset(
-        dataset_dict, inputs_columns=["image"], outputs_columns=["label"]
+    dataset = to_dashai_dataset(dataset_dict)
+
+    total_rows = len(dataset["train"])
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=total_rows, train_size=0.7, test_size=0.1, val_size=0.2
     )
-    split_dataset = image_dataloader.split_dataset(
-        dataset,
-        train_size=0.7,
-        test_size=0.1,
-        val_size=0.2,
-        class_column=dataset["train"].outputs_columns[0],
+    split_datasetdict = split_dataset(
+        dataset["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
     )
 
-    yield split_dataset
+    yield split_datasetdict
     shutil.rmtree("tests/back/tasks/beans_dataset", ignore_errors=True)
 
 
 def test_validate_image_class_task(image_classification_dataset):
     image_class_task = ImageClassificationTask()
-    dataset = image_class_task.prepare_for_task(image_classification_dataset)
+    inputs_columns = ["image"]
+    outputs_columns = ["label"]
+
+    dataset = image_class_task.prepare_for_task(
+        image_classification_dataset, outputs_columns
+    )
     try:
         image_class_task.validate_dataset_for_task(
             dataset=dataset,
             dataset_name="Beans Dataset",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
         )
     except Exception as e:
         pytest.fail(f"Unexpected error in test_validate_task: {repr(e)}")
 
 
+def test_get_image_class_task_metadata():
+    image_class_task = ImageClassificationTask()
+    metadata = image_class_task.get_metadata()
+
+    assert len(metadata.keys()) == 4
+    assert metadata["inputs_types"] == ["Image"]
+    assert metadata["outputs_types"] == ["ClassLabel"]
+    assert metadata["inputs_cardinality"] == 1
+    assert metadata["outputs_cardinality"] == 1
+
+
 @pytest.fixture(scope="module", name="translation_dataset")
 def translation_dataset_fixture():
     test_dataset_path = "tests/back/tasks/translationEngSpaDatasetSmall.json"
@@ -222,26 +280,43 @@ def translation_dataset_fixture():
         params={"data_key": "data"},
     )
 
-    dataset = to_dashai_dataset(
-        dataset, inputs_columns=["text"], outputs_columns=["class"]
-    )
+    dataset = to_dashai_dataset(dataset)
 
-    split_dataset = json_dataloader.split_dataset(
-        dataset,
-        train_size=0.7,
-        test_size=0.1,
-        val_size=0.2,
-        class_column=dataset["train"].outputs_columns[0],
+    total_rows = len(dataset["train"])
+    train_indexes, test_indexes, val_indexes = split_indexes(
+        total_rows=total_rows, train_size=0.7, test_size=0.1, val_size=0.2
+    )
+    split_datasetdict = split_dataset(
+        dataset["train"],
+        train_indexes=train_indexes,
+        test_indexes=test_indexes,
+        val_indexes=val_indexes,
     )
-    return split_dataset
+    return split_datasetdict
 
 
 def test_validate_translation_task(translation_dataset):
     translation_task = TranslationTask()
-    dataset = translation_task.prepare_for_task(translation_dataset)
+    inputs_columns = ["text"]
+    outputs_columns = ["class"]
+    dataset = translation_task.prepare_for_task(translation_dataset, outputs_columns)
     try:
         translation_task.validate_dataset_for_task(
-            dataset=dataset, dataset_name="EngSpaDataset"
+            dataset=dataset,
+            dataset_name="EngSpaDataset",
+            input_columns=inputs_columns,
+            output_columns=outputs_columns,
         )
     except Exception as e:
         pytest.fail(f"Unexpected error in test_validate_task: {repr(e)}")
+
+
+def test_get_translation_task_metadata():
+    translation_task = TranslationTask()
+    metadata = translation_task.get_metadata()
+
+    assert len(metadata.keys()) == 4
+    assert metadata["inputs_types"] == ["Value", "Sequence"]
+    assert metadata["outputs_types"] == ["Value", "Sequence"]
+    assert metadata["inputs_cardinality"] == 1
+    assert metadata["outputs_cardinality"] == 1