Repository: DashAISoftware/DashAI (language: python)

Issue:
Command Line Interface for DashAI
# Summary

This pull request adds a CLI (Command Line Interface) for DashAI.

The motivation for this is to be able to have the option to configure the execution of DashAI through the command line when initializing the platform.

In this pr only the  `dev_mode` option is added, which currently does nothing, but will be useful for the following iterations.


## Type of change

- Back end new feature.

## Changes

- A CLI was added to run DashAI based on Typer (library from the same creator of FastAPI, with a very similar interface to make CLI). Now the application must be initialized as: `python -m DashAI` (see next section for more information).
- The contents of `__init__.py` were moved to `__main__.py` to tell python to directly execute the contents of the file when the DashAI module is executed through the past command.

## How to Test

Due to the type of change executed, no tests were implemented.


First, install the requirements:

```bash
$ pip install -U -r requirements.txt
$ pip install -U -r requirements-dev.txt
```

To start the application, run: 

```bash
$ python -m DashAI
```

To view the CLI help, run:

```bash
$ python -m DashAI --help
```

To set the app startup in dev mode (it does nothing yet, it is only demonstrative), run: 

```bash
$ python -m DashAI --dev-mode
```

## Screenshots

<img width="1681" alt="image" src="https://github.com/DashAISoftware/DashAI/assets/22802485/69b56825-3e51-4c0a-b353-897fbacecdca">



Discussion / hints:

Fix patch:
diff --git a/DashAI/__init__.py b/DashAI/__main__.py
similarity index 64%
rename from DashAI/__init__.py
rename to DashAI/__main__.py
index 8d2b743e..c7045164 100644
--- a/DashAI/__init__.py
+++ b/DashAI/__main__.py
@@ -1,19 +1,20 @@
-# Este archivo corre al hacer import DashAI
 import logging
 import sys
 import threading
 import webbrowser
 
+import typer
 import uvicorn
 from sqlalchemy.exc import DBAPIError, SQLAlchemyError
 from sqlalchemy.sql import text
+from typing_extensions import Annotated
 
 from DashAI.back.database.models import Base
 from DashAI.back.database.session import SessionLocal, engine
 from DashAI.back.main import app
 
 logging.basicConfig(level=logging.DEBUG)
-log = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
 
 
 def open_browser():
@@ -21,14 +22,26 @@ def open_browser():
     webbrowser.open(url, new=0, autoraise=True)
 
 
-def run():
+def main(
+    dev_mode: Annotated[
+        bool, typer.Option(help="Run DashAI in development mode.")
+    ] = True,
+):
+    if dev_mode:
+        logging.info("DashAI was set to development mode.")
+
     db = SessionLocal()
     Base.metadata.create_all(engine)
     timer = threading.Timer(1, open_browser)
     timer.start()
+
     try:
         db.execute(text("SELECT 1"))
     except (SQLAlchemyError, DBAPIError):
-        log.error("There was an error checking database health")
+        _logger.error("There was an error checking database health")
         sys.exit(1)
     uvicorn.run(app, host="127.0.0.1", port=8000)
+
+
+if __name__ == "__main__":
+    typer.run(main)
diff --git a/DashAI/back/api/api_v1/api.py b/DashAI/back/api/api_v1/api.py
index 1615f0e0..585588af 100644
--- a/DashAI/back/api/api_v1/api.py
+++ b/DashAI/back/api/api_v1/api.py
@@ -5,6 +5,7 @@
     datasets,
     experiments,
     jobs,
+    predict,
     runs,
 )
 
@@ -14,3 +15,4 @@
 api_router_v1.include_router(experiments.router, prefix="/experiment")
 api_router_v1.include_router(jobs.router, prefix="/job")
 api_router_v1.include_router(runs.router, prefix="/run")
+api_router_v1.include_router(predict.router, prefix="/predict")
diff --git a/DashAI/back/api/api_v1/endpoints/datasets.py b/DashAI/back/api/api_v1/endpoints/datasets.py
index 20b2a199..5b707e4a 100644
--- a/DashAI/back/api/api_v1/endpoints/datasets.py
+++ b/DashAI/back/api/api_v1/endpoints/datasets.py
@@ -14,7 +14,10 @@
 from DashAI.back.api.utils import parse_params
 from DashAI.back.core.config import component_registry, settings
 from DashAI.back.database.models import Dataset
-from DashAI.back.dataloaders.classes.dashai_dataset import save_dataset
+from DashAI.back.dataloaders.classes.dashai_dataset import (
+    get_dataset_info,
+    save_dataset,
+)
 from DashAI.back.dataloaders.classes.dataloader import to_dashai_dataset
 
 logging.basicConfig(level=logging.DEBUG)
@@ -77,6 +80,39 @@ async def get_dataset(dataset_id: int, db: Session = Depends(get_db)):
     return dataset
 
 
+@router.get("/info/{dataset_id}")
+async def get_info(dataset_id: int, db: Session = Depends(get_db)):
+    """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.
+    """
+    try:
+        dataset = db.get(Dataset, dataset_id)
+        info = get_dataset_info(f"{dataset.file_path}/dataset")
+        if not dataset:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Dataset not found",
+            )
+
+    except exc.SQLAlchemyError as e:
+        log.exception(e)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Internal database error",
+        ) from e
+
+    return info
+
+
 @router.post("/", status_code=status.HTTP_201_CREATED)
 async def upload_dataset(
     db: Session = Depends(get_db),
diff --git a/DashAI/back/api/api_v1/endpoints/predict.py b/DashAI/back/api/api_v1/endpoints/predict.py
new file mode 100644
index 00000000..d7b9ad9f
--- /dev/null
+++ b/DashAI/back/api/api_v1/endpoints/predict.py
@@ -0,0 +1,155 @@
+import json
+import logging
+import os
+from typing import Any, List
+
+import pandas as pd
+from datasets import Dataset
+from fastapi import APIRouter, Depends, UploadFile, status
+from fastapi.exceptions import HTTPException
+from sqlalchemy import exc
+from sqlalchemy.orm import Session
+
+from DashAI.back.api.api_v1.schemas.predict_params import PredictParams
+from DashAI.back.api.deps import get_db
+from DashAI.back.core.config import component_registry, settings
+from DashAI.back.database.models import Dataset as Dt
+from DashAI.back.database.models import Experiment, Run
+from DashAI.back.dataloaders.classes.dataloader import BaseDataLoader, to_dashai_dataset
+from DashAI.back.models.base_model import BaseModel
+
+logging.basicConfig(level=logging.DEBUG)
+log = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.get("/")
+async def get_prediction():
+    """Placeholder for prediction get.
+
+    Raises
+    ------
+    HTTPException
+        Always raises exception as it was intentionally not implemented.
+    """
+    raise HTTPException(
+        status_code=status.HTTP_501_NOT_IMPLEMENTED, detail="Method not implemented"
+    )
+
+
+@router.post("/")
+async def perform_predict(
+    input_file: UploadFile,
+    params: PredictParams = Depends(),
+    db: Session = Depends(get_db),
+) -> List[Any]:
+    """
+    Endpoint to perform model prediction for a particular run, given some
+    sample values.
+
+    Parameters
+    ----------
+    input_file: UploadFile
+        File containing the sample data to be used for prediction.
+        The format of the sample data must match the format of the data set used to
+        train the run.
+    run_id: int
+        Id of the run to be used to predict.
+
+    Returns
+    -------
+    list
+        A list with the predictions given by the run.
+        The type of each prediction is given by the task associated with the run.
+    Raises
+    ------
+    HTTPException
+        If run_id does not exist in the database.
+        If experiment_id assoc. with the run does not exist in the database.
+        If dataset_id assoc. with the experiment does not exist in the database.
+    """
+    try:
+        run: Run = db.get(Run, params.run_id)
+        if not run:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND, detail="Run not found"
+            )
+
+        exp: Experiment = db.get(Experiment, run.experiment_id)
+        if not exp:
+            raise HTTPException(
+                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:
+        log.exception(e)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Internal database error",
+        ) from e
+
+    model = component_registry[run.model_name]["class"]
+    trained_model: BaseModel = model.load(run.run_path)
+
+    # Load Dataset using Dataloader
+    tmp_path = os.path.join(
+        settings.USER_DATASET_PATH, "tmp_predict", str(params.run_id)
+    )
+    try:
+        os.makedirs(tmp_path, exist_ok=True)
+    except FileExistsError as e:
+        log.exception(e)
+        raise HTTPException(
+            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=tmp_path, params={"data_key": "data"}
+    )
+    input_df = pd.DataFrame(raw_dataset["train"])
+    # TODO: Use feature_names from Experiment
+    input_df = input_df.reindex(columns=json.loads(dat.feature_names))
+    raw_dataset["train"] = Dataset.from_pandas(input_df)
+
+    # Transform into DashAIDataset
+    dataset = to_dashai_dataset(raw_dataset, raw_dataset["train"].column_names, [])
+
+    y_pred = trained_model.predict(dataset["train"])
+
+    return y_pred.tolist()
+
+
+@router.delete("/")
+async def delete_prediction():
+    """Placeholder for prediction delete.
+
+    Raises
+    ------
+    HTTPException
+        Always raises exception as it was intentionally not implemented.
+    """
+    raise HTTPException(
+        status_code=status.HTTP_501_NOT_IMPLEMENTED, detail="Method not implemented"
+    )
+
+
+@router.patch("/")
+async def update_prediction():
+    """Placeholder for prediction update.
+
+    Raises
+    ------
+    HTTPException
+        Always raises exception as it was intentionally not implemented.
+    """
+    raise HTTPException(
+        status_code=status.HTTP_501_NOT_IMPLEMENTED, detail="Method not implemented"
+    )
diff --git a/DashAI/back/api/api_v1/schemas/predict_params.py b/DashAI/back/api/api_v1/schemas/predict_params.py
new file mode 100644
index 00000000..8ffc40ce
--- /dev/null
+++ b/DashAI/back/api/api_v1/schemas/predict_params.py
@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class PredictParams(BaseModel):
+    run_id: int
diff --git a/DashAI/back/converters/__init__.py b/DashAI/back/converters/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/DashAI/back/converters/base_converter.py b/DashAI/back/converters/base_converter.py
new file mode 100644
index 00000000..45742061
--- /dev/null
+++ b/DashAI/back/converters/base_converter.py
@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from abc import ABCMeta, abstractmethod
+from typing import Final
+
+from datasets import DatasetDict
+from sklearn.base import BaseEstimator, TransformerMixin
+
+from DashAI.back.config_object import ConfigObject
+
+
+class BaseConverter(ConfigObject, BaseEstimator, TransformerMixin, metaclass=ABCMeta):
+    """Base class for all converters"""
+
+    TYPE: Final[str] = "Converter"
+
+    @abstractmethod
+    def fit(self, dataset: DatasetDict) -> "BaseConverter":
+        """Fit the converter.
+
+        Parameters
+        ----------
+        dataset : DatasetDict
+            Dataset to fit the converter
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def transform(self, dataset: DatasetDict) -> DatasetDict:
+        """Transform the dataset.
+
+        Parameters
+        ----------
+        dataset : DatasetDict
+            Dataset to be converted
+
+        Returns
+        -------
+            Dataset converted
+        """
+        raise NotImplementedError
diff --git a/DashAI/back/converters/column_dropper_by_index.py b/DashAI/back/converters/column_dropper_by_index.py
new file mode 100644
index 00000000..8d10e616
--- /dev/null
+++ b/DashAI/back/converters/column_dropper_by_index.py
@@ -0,0 +1,60 @@
+from typing import Type
+
+from beartype import beartype
+from datasets import DatasetDict
+
+from DashAI.back.converters.base_converter import BaseConverter
+from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
+
+
+class ColumnDropperByIndex(BaseConverter):
+    """Converter to drop columns from the dataset by column index"""
+
+    @beartype
+    def __init__(self, columns_index: tuple[int, int] | int):
+        """Constructor with columns to be dropped by column index
+
+        Parameters
+        ----------
+        columns : tuple[int, int] | int
+            Columns to be dropped. The tuple contains the start and end index of the
+            columns to be dropped (both included). The int contains the index of the
+            column to be dropped.
+        """
+        if isinstance(columns_index, int):
+            columns_index = [columns_index, columns_index]
+        self.columns_index = columns_index
+
+    @beartype
+    def fit(self, dataset: DatasetDict) -> Type["BaseConverter"]:
+        """Fit the converter.
+
+        Parameters
+        ----------
+        dataset : DatasetDict
+            Dataset to fit the converter
+        """
+        return self
+
+    @beartype
+    def transform(self, dataset: DatasetDict) -> DatasetDict:
+        """Convert the dataset by removing columns.
+
+        Parameters
+        ----------
+        dataset : DatasetDict
+            Dataset to be converted
+
+        Returns
+        -------
+        DatasetDict
+            Dataset converted
+        """
+        for split in dataset:
+            dataset_split: DashAIDataset = dataset[split]
+            column_names_to_drop = dataset[split].column_names[
+                self.columns_index[0] : self.columns_index[1] + 1
+            ]
+            dataset_split = dataset_split.remove_columns(column_names_to_drop)
+            dataset[split] = dataset_split
+        return dataset
diff --git a/DashAI/back/converters/column_dropper_by_name.py b/DashAI/back/converters/column_dropper_by_name.py
new file mode 100644
index 00000000..651396eb
--- /dev/null
+++ b/DashAI/back/converters/column_dropper_by_name.py
@@ -0,0 +1,56 @@
+from typing import Type
+
+from beartype import beartype
+from datasets import DatasetDict
+
+from DashAI.back.converters.base_converter import BaseConverter
+from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
+
+
+class ColumnDropperByName(BaseConverter):
+    """Converter to drop columns from the dataset"""
+
+    @beartype
+    def __init__(self, column_names: list[str] | str):
+        """Constructor with columns to be dropped by column name
+
+        Parameters
+        ----------
+        columns : list[str] | str
+            Columns to be dropped. The list contains the names of the columns to be
+            dropped. The string contains the name of the column to be dropped.
+        """
+        if isinstance(column_names, str):
+            column_names = [column_names]
+        self.column_names = column_names
+
+    @beartype
+    def fit(self, dataset: DatasetDict) -> Type["BaseConverter"]:
+        """Fit the converter.
+
+        Parameters
+        ----------
+        dataset : DatasetDict
+            Dataset to fit the converter
+        """
+        return self
+
+    @beartype
+    def transform(self, dataset: DatasetDict) -> DatasetDict:
+        """Convert the dataset by removing columns.
+
+        Parameters
+        ----------
+        dataset : DatasetDict
+            Dataset to be converted
+
+        Returns
+        -------
+        DatasetDict
+            Dataset converted
+        """
+        for split in dataset:
+            dataset_split: DashAIDataset = dataset[split]
+            dataset_split = dataset_split.remove_columns(self.column_names)
+            dataset[split] = dataset_split
+        return dataset
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 9faa9f84..572f7ef2 100644
--- a/DashAI/back/dataloaders/classes/dashai_dataset.py
+++ b/DashAI/back/dataloaders/classes/dashai_dataset.py
@@ -312,17 +312,17 @@ def save_dataset(datasetdict: DatasetDict, path: str) -> None:
             ensure_ascii=False,
         )
 
-    @beartype
-    def get_columns_type(dataset_id: int) -> Dict[str, str]:
-        """Return predefined column types
-
-        Parameters
-        ----------
-        dataset_id : int
-            The dataset to request the types
 
-        Returns
-        ---------
-        Dict[str,str]
-            Dict with the respective column and type
-        """
+@beartype
+def get_dataset_info(dataset_path: str) -> object:
+    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
diff --git a/DashAI/front/src/api/datasets.ts b/DashAI/front/src/api/datasets.ts
index d8022257..54fee4ff 100644
--- a/DashAI/front/src/api/datasets.ts
+++ b/DashAI/front/src/api/datasets.ts
@@ -12,7 +12,10 @@ export const getDatasets = async (): Promise<IDataset[]> => {
   const response = await api.get<IDataset[]>(datasetEndpoint);
   return response.data;
 };
-
+export const getDatasetInfo = async (id: number): Promise<object> => {
+  const response = await api.get<object>(`${datasetEndpoint}/info/${id}`);
+  return response.data;
+};
 export const updateDataset = async (
   id: number,
   name: string | undefined = undefined,
diff --git a/DashAI/front/src/api/experiment.ts b/DashAI/front/src/api/experiment.ts
index e152a23f..9019b403 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: JSON,
+  outputColumns: JSON,
+  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);
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/DatasetModal.jsx b/DashAI/front/src/components/datasets/DatasetModal.jsx
index 02895c3c..73edcb72 100644
--- a/DashAI/front/src/components/datasets/DatasetModal.jsx
+++ b/DashAI/front/src/components/datasets/DatasetModal.jsx
@@ -13,20 +13,17 @@ 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";
 
 const steps = [
-  { name: "selectTask", label: "Select Task" },
   { name: "selectDataloader", label: "Select a way to upload" },
   { name: "uploadDataset", label: "Configure and upload your dataset" },
 ];
 
 const defaultNewDataset = {
-  task_name: "",
   dataloader: "",
   file: null,
   url: "",
@@ -159,24 +156,16 @@ function DatasetModal({ open, setOpen, updateDatasets }) {
 
       {/* Main content - steps */}
       <DialogContent dividers>
-        {/* Step 1: select task */}
+        {/* Step 1: select dataloader */}
         {activeStep === 0 && (
-          <SelectTaskStep
-            newDataset={newDataset}
-            setNewDataset={setNewDataset}
-            setNextEnabled={setNextEnabled}
-          />
-        )}
-        {/* Step 2: select dataloader */}
-        {activeStep === 1 && (
           <SelectDataloaderStep
             newDataset={newDataset}
             setNewDataset={setNewDataset}
             setNextEnabled={setNextEnabled}
           />
         )}
-        {/* Step 3: Configure dataloader and upload file */}
-        {activeStep === 2 && (
+        {/* Step 2: Configure dataloader and upload file */}
+        {activeStep === 1 && (
           <ConfigureAndUploadDataset
             newDataset={newDataset}
             setNewDataset={setNewDataset}
@@ -194,7 +183,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/SelectDataloaderStep.jsx b/DashAI/front/src/components/datasets/SelectDataloaderStep.jsx
index ccdc3716..c1d8c170 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 !== "") {
diff --git a/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx b/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx
index c9595ad5..8880a156 100644
--- a/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx
+++ b/DashAI/front/src/components/experiments/ConfigureModelsStep.jsx
@@ -157,6 +157,13 @@ 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({
+      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..0a80eb5b
--- /dev/null
+++ b/DashAI/front/src/components/experiments/DivideDatasetColumns.jsx
@@ -0,0 +1,128 @@
+import React, { useEffect, useState } from "react";
+import PropTypes from "prop-types";
+
+import { Grid, TextField, Typography } from "@mui/material";
+
+function DivideDatasetColumns({
+  datasetInfo,
+  inputColumns,
+  setInputColumns,
+  outputColumns,
+  setOutputColumns,
+  setColumnsReady,
+  parseRangeToIndex,
+}) {
+  // 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); // TODO: input and output columns should be less than total
+      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,
+  ]);
+  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-${datasetInfo.total_columns - 1}`}
+        sx={{ mb: 2 }}
+      />
+      <TextField
+        required
+        id="dataset-output-columns"
+        label="Output"
+        fullWidth
+        autoComplete="off"
+        onChange={handleOutputColumnsChange}
+        error={parseOutputColumnsError}
+        helperText={parseOutputColumnsError ? parseOutputColumnsErrorText : ""}
+        defaultValue={datasetInfo.total_columns}
+        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,
+  parseRangeToIndex: 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..496e29f2 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,9 @@ const defaultNewExp = {
   name: "",
   dataset: null,
   task_name: "",
+  input_columns: [],
+  output_columns: [],
+  splits: { training: 0, validation: 0, testing: 0 },
   step: "SET_NAME",
   created: null,
   last_modified: null,
@@ -94,6 +99,9 @@ export default function NewExperimentModal({
         newExp.dataset.id,
         newExp.task_name,
         newExp.name,
+        JSON.stringify(newExp.input_columns),
+        JSON.stringify(newExp.output_columns),
+        JSON.stringify(newExp.splits),
       );
       const experimentId = response.id;
       await uploadRuns(experimentId);
@@ -227,6 +235,13 @@ export default function NewExperimentModal({
           />
         )}
         {activeStep === 2 && (
+          <PrepareDatasetStep
+            newExp={newExp}
+            setNewExp={setNewExp}
+            setNextEnabled={setNextEnabled}
+          />
+        )}
+        {activeStep === 3 && (
           <ConfigureModelsStep
             newExp={newExp}
             setNewExp={setNewExp}
@@ -248,7 +263,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..45ccce15
--- /dev/null
+++ b/DashAI/front/src/components/experiments/PrepareDatasetStep.jsx
@@ -0,0 +1,188 @@
+import React, { useEffect, useState } from "react";
+import PropTypes from "prop-types";
+
+import { Grid, CircularProgress, Box } from "@mui/material";
+import DivideDatasetColumns from "./DivideDatasetColumns";
+import SplitDatasetRows from "./SplitDatasetRows";
+import { getDatasetInfo as getDatasetInfoRequest } from "../../api/datasets";
+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 [loading, setLoading] = useState(true);
+  // columns index state
+  const [inputColumns, setInputColumns] = useState([]);
+  const [outputColumns, setOutputColumns] = useState([]);
+  const [columnsReady, setColumnsReady] = useState(true);
+
+  const rangeRegex = /^(\d+)(-(\d+))*(,(\d+)(-(\d+))*)*$/;
+
+  // rows index state
+  const defaultParitionsIndex = {
+    train: [],
+    validation: [],
+    test: [],
+  };
+  const defaultPartitionsPercentage = {
+    train: 60,
+    validation: 20,
+    test: 20,
+  };
+
+  const [rowsPartitionsIndex, setRowsPartitionsIndex] = useState(
+    defaultParitionsIndex,
+  );
+  const [rowsPartitionsPercentage, setRowsPartitionsPercentage] = useState(
+    defaultPartitionsPercentage,
+  );
+  const [splitsReady, setSplitsReady] = useState(false);
+
+  const getDatasetInfo = async () => {
+    setLoading(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 {
+      setLoading(false);
+    }
+  };
+
+  const parseRangeToIndex = (value, total) => {
+    const numbersArray = [];
+    if (!rangeRegex.test(value)) {
+      throw new Error("The entered text doesn't fit the example format");
+    }
+    const ranges = value.split(",");
+    ranges.forEach((range) => {
+      const [min, max] = range.split("-");
+      if (!range.includes("-") && parseInt(range) <= total) {
+        numbersArray.push(parseInt(range));
+      } else if (
+        (!range.includes("-") && parseInt(range) > total) ||
+        parseInt(max) > total
+      ) {
+        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++) {
+          numbersArray.push(i);
+        }
+      }
+    });
+    return numbersArray;
+  };
+
+  // Fetch dataset when the component is mounting
+  useEffect(() => {
+    getDatasetInfo();
+  }, []);
+  // Set the input and output columns by default
+  useEffect(() => {
+    if (!loading) {
+      setInputColumns(
+        parseRangeToIndex(
+          `1-${datasetInfo.total_columns - 1}`,
+          datasetInfo.total_columns,
+        ),
+      );
+      setOutputColumns(
+        parseRangeToIndex(
+          `${datasetInfo.total_columns}`,
+          datasetInfo.total_columns,
+        ),
+      );
+    }
+  }, [loading]);
+  useEffect(() => {
+    if (columnsReady && splitsReady) {
+      setNewExp({
+        ...newExp,
+        input_columns: inputColumns,
+        output_columns: outputColumns,
+        splits:
+          rowsPartitionsIndex !== defaultParitionsIndex
+            ? rowsPartitionsIndex
+            : rowsPartitionsPercentage,
+      }); // splits should depend on preference
+      setNextEnabled(true);
+    } else {
+      setNextEnabled(false);
+    }
+  }, [columnsReady, splitsReady]);
+
+  return (
+    <React.Fragment>
+      {!loading ? (
+        <Grid container spacing={1}>
+          <DivideDatasetColumns
+            datasetInfo={datasetInfo}
+            inputColumns={inputColumns}
+            setInputColumns={setInputColumns}
+            outputColumns={outputColumns}
+            setOutputColumns={setOutputColumns}
+            setColumnsReady={setColumnsReady}
+            parseRangeToIndex={parseRangeToIndex}
+          />
+          <SplitDatasetRows
+            datasetInfo={datasetInfo}
+            rowsPartitionsIndex={rowsPartitionsIndex}
+            setRowsPartitionsIndex={setRowsPartitionsIndex}
+            rowsPartitionsPercentage={rowsPartitionsPercentage}
+            setRowsPartitionsPercentage={setRowsPartitionsPercentage}
+            setSplitsReady={setSplitsReady}
+            parseRangeToIndex={parseRangeToIndex}
+          />
+        </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({
+      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..64ab3800 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,13 @@ 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({
+      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..2df06e14 100644
--- a/DashAI/front/src/components/experiments/SetNameAndTaskStep.jsx
+++ b/DashAI/front/src/components/experiments/SetNameAndTaskStep.jsx
@@ -150,6 +150,13 @@ 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({
+      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..9c1c00da
--- /dev/null
+++ b/DashAI/front/src/components/experiments/SplitDatasetRows.jsx
@@ -0,0 +1,291 @@
+import React, { useEffect, useState } from "react";
+import PropTypes from "prop-types";
+
+import {
+  Grid,
+  TextField,
+  Typography,
+  FormControlLabel,
+  Radio,
+  RadioGroup,
+  FormHelperText,
+  InputAdornment,
+} from "@mui/material";
+
+function SplitDatasetRows({
+  datasetInfo,
+  rowsPartitionsIndex,
+  setRowsPartitionsIndex,
+  rowsPartitionsPercentage,
+  setRowsPartitionsPercentage,
+  setSplitsReady,
+  parseRangeToIndex,
+}) {
+  const totalRows = datasetInfo.total_rows;
+  const defaultParitionsIndex = {
+    training: [],
+    validation: [],
+    testing: [],
+  };
+  const defaultPartitionsPercentage = {
+    training: 60,
+    validation: 20,
+    testing: 20,
+  };
+
+  // handle rows numbers change state
+  const [rowsPreference, setRowsPreference] = useState("random");
+  const [rowsPartitionsError, setRowsPartitionsError] = useState(false);
+  const [rowsPartitionsErrorText, setRowsPartitionsErrorText] = useState("");
+
+  const checkSplit = (train, validation, test) => {
+    return train + validation + test === 100;
+  };
+
+  const handleRowsPreferenceChange = (event) => {
+    if (event.target.value === "splitByIndex") {
+      setRowsPartitionsPercentage(defaultPartitionsPercentage);
+    } else {
+      setRowsPartitionsIndex(defaultParitionsIndex);
+    }
+    setRowsPartitionsError(false);
+    setRowsPartitionsErrorText("");
+    setRowsPreference(event.target.value);
+  };
+
+  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 (rowsPreference === "splitByIndex") {
+      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: parseInt(value) };
+          break;
+        case "validation":
+          newSplit = { ...newSplit, validation: parseInt(value) };
+          break;
+        case "test":
+          newSplit = { ...newSplit, test: parseInt(value) };
+          break;
+      }
+      setRowsPartitionsPercentage(newSplit);
+      if (!checkSplit(newSplit.train, newSplit.validation, newSplit.test)) {
+        setRowsPartitionsErrorText("Splits should add 100%");
+        setRowsPartitionsError(true);
+      } else {
+        setRowsPartitionsError(false);
+      }
+    }
+  };
+
+  useEffect(() => {
+    // check if splits doesnt have errors and arent empty
+    if (
+      rowsPreference === "splitByIndex" &&
+      !rowsPartitionsError &&
+      rowsPartitionsIndex.train.length >= 1 &&
+      rowsPartitionsIndex.validation.length >= 1 &&
+      rowsPartitionsIndex.test.length >= 1
+    ) {
+      setSplitsReady(true);
+    } else if (
+      rowsPreference === "random" &&
+      !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 percentage"
+          sx={{ my: 1 }}
+        />
+        {rowsPreference === "random" ? (
+          <React.Fragment>
+            <Grid container direction="row" spacing={4}>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="train"
+                  label="Train"
+                  autoComplete="off"
+                  type="number"
+                  size="small"
+                  InputProps={{
+                    endAdornment: (
+                      <InputAdornment position="end">%</InputAdornment>
+                    ),
+                  }}
+                  error={rowsPartitionsError}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="validation"
+                  label="Validation"
+                  autoComplete="off"
+                  type="number"
+                  size="small"
+                  InputProps={{
+                    endAdornment: (
+                      <InputAdornment position="end">%</InputAdornment>
+                    ),
+                  }}
+                  error={rowsPartitionsError}
+                  onChange={handleRowsChange}
+                />
+              </Grid>
+              <Grid item sx={{ xs: 4 }}>
+                <TextField
+                  id="test"
+                  label="Test"
+                  type="number"
+                  size="small"
+                  InputProps={{
+                    endAdornment: (
+                      <InputAdornment position="end">%</InputAdornment>
+                    ),
+                  }}
+                  autoComplete="off"
+                  error={rowsPartitionsError}
+                  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 }}
+        />
+        {rowsPreference === "splitByIndex" ? (
+          <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,
+  parseRangeToIndex: PropTypes.func.isRequired,
+};
+export default SplitDatasetRows;
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/README.rst b/README.rst
index 14759d4e..a6dc2de6 100644
--- a/README.rst
+++ b/README.rst
@@ -175,15 +175,9 @@ There are three ways to run DashAI:
 
 .. code:: bash
 
-    $ python -c "import DashAI;DashAI.run()"
+    $ python -m DashAI
 
-2. By executing DashAI as a module and enabling the plugins system:
-
-.. code:: bash
-
-    $ python -c "import DashAI;DashAI.run_plugins()"
-
-3. Or,  installing the default build:
+2. Or,  installing the default build:
 
 .. code:: bash
 
diff --git a/requirements-dev.txt b/requirements-dev.txt
index eff80622..ebbba9d1 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,7 +1,7 @@
 pytest>=7.1.2
 hypothesis==6.52.1
 pre-commit>=3.3.3
-ruff>=0.0.284
+ruff>=0.0.291
 black>=23.7.0
 isort>=5.12.0
 sphinx_rtd_theme==1.2.0
diff --git a/requirements.txt b/requirements.txt
index 09922c5d..2292110b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,7 +12,8 @@ evaluate>=0.4.0
 accelerate>=0.20.3
 Pillow>=9.5.0
 beartype>=0.15.0
-
+typer>=0.9.0
+rich>=13.5.3
 torch==1.13.0 #+cu116
 --find-links https://download.pytorch.org/whl/torch_stable.html
 transformers>=4.23.1


Test patch:
diff --git a/tests/back/api/input_iris.json b/tests/back/api/input_iris.json
new file mode 100644
index 00000000..a32236c4
--- /dev/null
+++ b/tests/back/api/input_iris.json
@@ -0,0 +1,23 @@
+{
+    "task_info": "TabularClassificationTask",
+    "data": [
+        {
+            "feature_0": 4.9,
+            "feature_1": 3.0,
+            "feature_2": 1.4,
+            "feature_3": 0.2
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 3.0,
+            "feature_2": 4.2,
+            "feature_3": 1.2
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.1,
+            "feature_2": 5.6,
+            "feature_3": 2.4
+        }
+    ]
+}
diff --git a/tests/back/api/irisDataset.json b/tests/back/api/irisDataset.json
new file mode 100644
index 00000000..7130cbd7
--- /dev/null
+++ b/tests/back/api/irisDataset.json
@@ -0,0 +1,1055 @@
+{
+    "task_info": "TabularClassificationTask",
+    "data": [
+        {
+            "feature_0": 4.9,
+            "feature_1": 3.0,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.7,
+            "feature_1": 3.2,
+            "feature_2": 1.3,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.6,
+            "feature_1": 3.1,
+            "feature_2": 1.5,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.6,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.6,
+            "feature_1": 3.4,
+            "feature_2": 1.4,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.4,
+            "feature_2": 1.5,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.4,
+            "feature_1": 2.9,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.9,
+            "feature_1": 3.1,
+            "feature_2": 1.5,
+            "feature_3": 0.1,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.8,
+            "feature_1": 3.4,
+            "feature_2": 1.6,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.8,
+            "feature_1": 3.0,
+            "feature_2": 1.4,
+            "feature_3": 0.1,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.3,
+            "feature_1": 3.0,
+            "feature_2": 1.1,
+            "feature_3": 0.1,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 4.0,
+            "feature_2": 1.2,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.4,
+            "feature_1": 3.9,
+            "feature_2": 1.3,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.5,
+            "feature_2": 1.4,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 3.8,
+            "feature_2": 1.7,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.8,
+            "feature_2": 1.5,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.7,
+            "feature_2": 1.5,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.6,
+            "feature_1": 3.6,
+            "feature_2": 1.0,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.3,
+            "feature_2": 1.7,
+            "feature_3": 0.5,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.8,
+            "feature_1": 3.4,
+            "feature_2": 1.9,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.4,
+            "feature_2": 1.6,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.2,
+            "feature_1": 3.5,
+            "feature_2": 1.5,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.2,
+            "feature_1": 3.4,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.7,
+            "feature_1": 3.2,
+            "feature_2": 1.6,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.4,
+            "feature_1": 3.4,
+            "feature_2": 1.5,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.2,
+            "feature_1": 4.1,
+            "feature_2": 1.5,
+            "feature_3": 0.1,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 4.2,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.9,
+            "feature_1": 3.1,
+            "feature_2": 1.5,
+            "feature_3": 0.1,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 3.5,
+            "feature_2": 1.3,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.9,
+            "feature_1": 3.1,
+            "feature_2": 1.5,
+            "feature_3": 0.1,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.4,
+            "feature_1": 3.0,
+            "feature_2": 1.3,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.4,
+            "feature_2": 1.5,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.5,
+            "feature_1": 2.3,
+            "feature_2": 1.3,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.4,
+            "feature_1": 3.2,
+            "feature_2": 1.3,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.5,
+            "feature_2": 1.6,
+            "feature_3": 0.6,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.8,
+            "feature_2": 1.9,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.8,
+            "feature_2": 1.6,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.6,
+            "feature_1": 3.2,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.3,
+            "feature_1": 3.7,
+            "feature_2": 1.5,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.3,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 3.2,
+            "feature_2": 4.5,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.9,
+            "feature_1": 3.1,
+            "feature_2": 4.9,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 2.3,
+            "feature_2": 4.0,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.5,
+            "feature_1": 2.8,
+            "feature_2": 4.6,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 3.3,
+            "feature_2": 4.7,
+            "feature_3": 1.6,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 4.9,
+            "feature_1": 2.4,
+            "feature_2": 3.3,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.6,
+            "feature_1": 2.9,
+            "feature_2": 4.6,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.2,
+            "feature_1": 2.7,
+            "feature_2": 3.9,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.9,
+            "feature_1": 3.0,
+            "feature_2": 4.2,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.0,
+            "feature_1": 2.2,
+            "feature_2": 4.0,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.1,
+            "feature_1": 2.9,
+            "feature_2": 4.7,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.6,
+            "feature_1": 2.9,
+            "feature_2": 3.6,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.6,
+            "feature_1": 3.0,
+            "feature_2": 4.5,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 2.7,
+            "feature_2": 4.1,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.2,
+            "feature_1": 2.2,
+            "feature_2": 4.5,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.6,
+            "feature_1": 2.5,
+            "feature_2": 3.9,
+            "feature_3": 1.1,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.1,
+            "feature_1": 2.8,
+            "feature_2": 4.0,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 2.5,
+            "feature_2": 4.9,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.1,
+            "feature_1": 2.8,
+            "feature_2": 4.7,
+            "feature_3": 1.2,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 2.9,
+            "feature_2": 4.3,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.8,
+            "feature_1": 2.8,
+            "feature_2": 4.8,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.0,
+            "feature_2": 5.0,
+            "feature_3": 1.7,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.0,
+            "feature_1": 2.9,
+            "feature_2": 4.5,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 2.6,
+            "feature_2": 3.5,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 2.4,
+            "feature_2": 3.7,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 2.7,
+            "feature_2": 3.9,
+            "feature_3": 1.2,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.0,
+            "feature_1": 2.7,
+            "feature_2": 5.1,
+            "feature_3": 1.6,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.4,
+            "feature_1": 3.0,
+            "feature_2": 4.5,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.1,
+            "feature_2": 4.7,
+            "feature_3": 1.5,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 2.3,
+            "feature_2": 4.4,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.6,
+            "feature_1": 3.0,
+            "feature_2": 4.1,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 2.5,
+            "feature_2": 4.0,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.1,
+            "feature_1": 3.0,
+            "feature_2": 4.6,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 2.6,
+            "feature_2": 4.0,
+            "feature_3": 1.2,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 2.3,
+            "feature_2": 3.3,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.6,
+            "feature_1": 2.7,
+            "feature_2": 4.2,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 2.9,
+            "feature_2": 4.2,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.2,
+            "feature_1": 2.9,
+            "feature_2": 4.3,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 2.5,
+            "feature_2": 3.0,
+            "feature_3": 1.1,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 2.8,
+            "feature_2": 4.1,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 2.7,
+            "feature_2": 5.1,
+            "feature_3": 1.9,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.1,
+            "feature_1": 3.0,
+            "feature_2": 5.9,
+            "feature_3": 2.1,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 2.9,
+            "feature_2": 5.6,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.5,
+            "feature_1": 3.0,
+            "feature_2": 5.8,
+            "feature_3": 2.2,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 4.9,
+            "feature_1": 2.5,
+            "feature_2": 4.5,
+            "feature_3": 1.7,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.3,
+            "feature_1": 2.9,
+            "feature_2": 6.3,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 2.5,
+            "feature_2": 5.8,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.2,
+            "feature_1": 3.6,
+            "feature_2": 6.1,
+            "feature_3": 2.5,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 2.7,
+            "feature_2": 5.3,
+            "feature_3": 1.9,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.8,
+            "feature_1": 3.0,
+            "feature_2": 5.5,
+            "feature_3": 2.1,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 2.5,
+            "feature_2": 5.0,
+            "feature_3": 2.0,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 2.8,
+            "feature_2": 5.1,
+            "feature_3": 2.4,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.5,
+            "feature_1": 3.0,
+            "feature_2": 5.5,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.7,
+            "feature_1": 3.8,
+            "feature_2": 6.7,
+            "feature_3": 2.2,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.7,
+            "feature_1": 2.6,
+            "feature_2": 6.9,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.0,
+            "feature_1": 2.2,
+            "feature_2": 5.0,
+            "feature_3": 1.5,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 5.6,
+            "feature_1": 2.8,
+            "feature_2": 4.9,
+            "feature_3": 2.0,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.7,
+            "feature_1": 2.8,
+            "feature_2": 6.7,
+            "feature_3": 2.0,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 2.7,
+            "feature_2": 4.9,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.3,
+            "feature_2": 5.7,
+            "feature_3": 2.1,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.2,
+            "feature_1": 2.8,
+            "feature_2": 4.8,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.1,
+            "feature_1": 3.0,
+            "feature_2": 4.9,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 2.8,
+            "feature_2": 5.6,
+            "feature_3": 2.1,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.2,
+            "feature_1": 3.0,
+            "feature_2": 5.8,
+            "feature_3": 1.6,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.9,
+            "feature_1": 3.8,
+            "feature_2": 6.4,
+            "feature_3": 2.0,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 2.8,
+            "feature_2": 5.6,
+            "feature_3": 2.2,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 2.8,
+            "feature_2": 5.1,
+            "feature_3": 1.5,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.1,
+            "feature_1": 2.6,
+            "feature_2": 5.6,
+            "feature_3": 1.4,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 3.4,
+            "feature_2": 5.6,
+            "feature_3": 2.4,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 3.1,
+            "feature_2": 5.5,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.0,
+            "feature_1": 3.0,
+            "feature_2": 4.8,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.9,
+            "feature_1": 3.1,
+            "feature_2": 5.4,
+            "feature_3": 2.1,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.9,
+            "feature_1": 3.1,
+            "feature_2": 5.1,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 5.8,
+            "feature_1": 2.7,
+            "feature_2": 5.1,
+            "feature_3": 1.9,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.8,
+            "feature_1": 3.2,
+            "feature_2": 5.9,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.3,
+            "feature_2": 5.7,
+            "feature_3": 2.5,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 2.5,
+            "feature_2": 5.0,
+            "feature_3": 1.9,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.5,
+            "feature_1": 3.0,
+            "feature_2": 5.2,
+            "feature_3": 2.0,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.2,
+            "feature_1": 3.4,
+            "feature_2": 5.4,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 5.9,
+            "feature_1": 3.0,
+            "feature_2": 5.1,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 5.1,
+            "feature_1": 3.5,
+            "feature_2": 1.4,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.4,
+            "feature_1": 3.9,
+            "feature_2": 1.7,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.4,
+            "feature_1": 3.7,
+            "feature_2": 1.5,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 4.4,
+            "feature_2": 1.5,
+            "feature_3": 0.4,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.4,
+            "feature_1": 3.4,
+            "feature_2": 1.7,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.0,
+            "feature_2": 1.6,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.8,
+            "feature_1": 3.1,
+            "feature_2": 1.6,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.2,
+            "feature_2": 1.2,
+            "feature_3": 0.2,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 3.5,
+            "feature_2": 1.3,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 4.8,
+            "feature_1": 3.0,
+            "feature_2": 1.4,
+            "feature_3": 0.3,
+            "class": "Iris-setosa"
+        },
+        {
+            "feature_0": 7.0,
+            "feature_1": 3.2,
+            "feature_2": 4.7,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 2.8,
+            "feature_2": 4.5,
+            "feature_3": 1.3,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.0,
+            "feature_1": 2.0,
+            "feature_2": 3.5,
+            "feature_3": 1.0,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.1,
+            "feature_2": 4.4,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.9,
+            "feature_1": 3.2,
+            "feature_2": 4.8,
+            "feature_3": 1.8,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.6,
+            "feature_1": 3.0,
+            "feature_2": 4.4,
+            "feature_3": 1.4,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 2.4,
+            "feature_2": 3.8,
+            "feature_3": 1.1,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.0,
+            "feature_1": 3.4,
+            "feature_2": 4.5,
+            "feature_3": 1.6,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.5,
+            "feature_1": 2.6,
+            "feature_2": 4.4,
+            "feature_3": 1.2,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 5.7,
+            "feature_1": 3.0,
+            "feature_2": 4.2,
+            "feature_3": 1.2,
+            "class": "Iris-versicolor"
+        },
+        {
+            "feature_0": 6.3,
+            "feature_1": 3.3,
+            "feature_2": 6.0,
+            "feature_3": 2.5,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.6,
+            "feature_1": 3.0,
+            "feature_2": 6.6,
+            "feature_3": 2.1,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.5,
+            "feature_1": 3.2,
+            "feature_2": 5.1,
+            "feature_3": 2.0,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.4,
+            "feature_1": 3.2,
+            "feature_2": 5.3,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.9,
+            "feature_1": 3.2,
+            "feature_2": 5.7,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.2,
+            "feature_1": 3.2,
+            "feature_2": 6.0,
+            "feature_3": 1.8,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.4,
+            "feature_1": 2.8,
+            "feature_2": 6.1,
+            "feature_3": 1.9,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 7.7,
+            "feature_1": 3.0,
+            "feature_2": 6.1,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.1,
+            "feature_2": 5.6,
+            "feature_3": 2.4,
+            "class": "Iris-virginica"
+        },
+        {
+            "feature_0": 6.7,
+            "feature_1": 3.0,
+            "feature_2": 5.2,
+            "feature_3": 2.3,
+            "class": "Iris-virginica"
+        }
+    ]
+}
diff --git a/tests/back/api/test_predict.py b/tests/back/api/test_predict.py
new file mode 100644
index 00000000..2d4f866f
--- /dev/null
+++ b/tests/back/api/test_predict.py
@@ -0,0 +1,147 @@
+import json
+import os
+
+import pytest
+from fastapi.testclient import TestClient
+from sqlalchemy.orm import sessionmaker
+
+from DashAI.back.database.models import Experiment, Run
+
+
+@pytest.fixture(scope="module", name="dataset_id")
+def fixture_dataset_id(client: TestClient):
+    script_dir = os.path.dirname(__file__)
+    test_dataset = "irisDataset.json"
+    abs_file_path = os.path.join(script_dir, test_dataset)
+    with open(abs_file_path, "rb") as json_file:
+        response = client.post(
+            "/api/v1/dataset/",
+            data={
+                "params": """{  "task_name": "TabularClassificationTask",
+                                    "dataloader": "JSONDataLoader",
+                                    "dataset_name": "test_json_2",
+                                    "outputs_columns": ["class"],
+                                    "splits_in_folders": false,
+                                    "splits": {
+                                        "train_size": 0.5,
+                                        "test_size": 0.2,
+                                        "val_size": 0.3,
+                                        "seed": 42,
+                                        "shuffle": false,
+                                        "stratify": false
+                                    },
+                                    "dataloader_params": {
+                                        "data_key": "data"
+                                    }
+                                }""",
+                "url": "",
+            },
+            files={"file": ("filename", json_file, "text/json")},
+        )
+    assert response.status_code == 201, response.text
+    dataset = response.json()
+
+    yield dataset["id"]
+
+    response = client.delete(f"/api/v1/dataset/{dataset['id']}")
+    assert response.status_code == 204, response.text
+
+
+@pytest.fixture(scope="module", name="experiment_id")
+def fixture_experiment_id(session: sessionmaker, dataset_id: int):
+    db = session()
+
+    experiment = Experiment(
+        dataset_id=dataset_id,
+        name="Experiment",
+        task_name="TabularClassificationTask",
+        input_columns=json.dumps(
+            ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"]
+        ),
+        output_columns=json.dumps(["Species"]),
+        splits=json.dumps(
+            {
+                "train_size": 0.8,
+                "test_size": 0.1,
+                "val_size": 0.1,
+                "seed": 42,
+                "shuffle": True,
+                "stratify": False,
+            }
+        ),
+    )
+    db.add(experiment)
+    db.commit()
+    db.refresh(experiment)
+
+    yield experiment.id
+
+    db.delete(experiment)
+    db.commit()
+    db.close()
+
+
+@pytest.fixture(scope="module", name="run_id")
+def fixture_run_id(session: sessionmaker, experiment_id: int):
+    db = session()
+
+    run = Run(
+        experiment_id=experiment_id,
+        model_name="RandomForestClassifier",
+        parameters={},
+        name="Run",
+    )
+    db.add(run)
+    db.commit()
+    db.refresh(run)
+
+    yield run.id
+
+    db.delete(run)
+    db.commit()
+    db.close()
+
+
+@pytest.fixture(scope="module", name="trained_run_id")
+def fixture_trained_run_id(client: TestClient, run_id: int):
+    response = client.post(
+        "/api/v1/job/runner/",
+        json={"run_id": run_id},
+    )
+    assert response.status_code == 201, response.text
+
+    response = client.post("/api/v1/job/start/?stop_when_queue_empties=True")
+    assert response.status_code == 202, response.text
+
+    return run_id
+
+
+def test_get_prediction(client: TestClient):
+    response = client.get("/api/v1/predict/")
+    assert response.status_code == 501, response.text
+
+
+def test_make_prediction(client: TestClient, trained_run_id: int):
+    script_dir = os.path.dirname(__file__)
+    test_dataset = "input_iris.json"
+    abs_file_path = os.path.join(script_dir, test_dataset)
+    with open(abs_file_path, "rb") as json_file:
+        response = client.post(
+            "/api/v1/predict/",
+            params={
+                "run_id": trained_run_id,
+            },
+            files={"input_file": ("filename", json_file, "text/json")},
+        )
+        data = response.json()
+        assert len(data) == 3
+
+
+def test_delete_prediction(client: TestClient):
+    response = client.delete("/api/v1/predict/")
+    assert response.status_code == 501, response.text
+
+
+def test_patch_prediction(client: TestClient):
+    response = client.patch("/api/v1/predict/")
+    assert response.status_code == 501, response.text
diff --git a/tests/back/converters/iris.csv b/tests/back/converters/iris.csv
new file mode 100644
index 00000000..70120eba
--- /dev/null
+++ b/tests/back/converters/iris.csv
@@ -0,0 +1,151 @@
+SepalLengthCm,SepalWidthCm,PetalLengthCm,PetalWidthCm,Species
+5.1,3.5,1.4,0.2,Iris-setosa
+4.9,3.0,1.4,0.2,Iris-setosa
+4.7,3.2,1.3,0.2,Iris-setosa
+4.6,3.1,1.5,0.2,Iris-setosa
+5.0,3.6,1.4,0.2,Iris-setosa
+5.4,3.9,1.7,0.4,Iris-setosa
+4.6,3.4,1.4,0.3,Iris-setosa
+5.0,3.4,1.5,0.2,Iris-setosa
+4.4,2.9,1.4,0.2,Iris-setosa
+4.9,3.1,1.5,0.1,Iris-setosa
+5.4,3.7,1.5,0.2,Iris-setosa
+4.8,3.4,1.6,0.2,Iris-setosa
+4.8,3.0,1.4,0.1,Iris-setosa
+4.3,3.0,1.1,0.1,Iris-setosa
+5.8,4.0,1.2,0.2,Iris-setosa
+5.7,4.4,1.5,0.4,Iris-setosa
+5.4,3.9,1.3,0.4,Iris-setosa
+5.1,3.5,1.4,0.3,Iris-setosa
+5.7,3.8,1.7,0.3,Iris-setosa
+5.1,3.8,1.5,0.3,Iris-setosa
+5.4,3.4,1.7,0.2,Iris-setosa
+5.1,3.7,1.5,0.4,Iris-setosa
+4.6,3.6,1.0,0.2,Iris-setosa
+5.1,3.3,1.7,0.5,Iris-setosa
+4.8,3.4,1.9,0.2,Iris-setosa
+5.0,3.0,1.6,0.2,Iris-setosa
+5.0,3.4,1.6,0.4,Iris-setosa
+5.2,3.5,1.5,0.2,Iris-setosa
+5.2,3.4,1.4,0.2,Iris-setosa
+4.7,3.2,1.6,0.2,Iris-setosa
+4.8,3.1,1.6,0.2,Iris-setosa
+5.4,3.4,1.5,0.4,Iris-setosa
+5.2,4.1,1.5,0.1,Iris-setosa
+5.5,4.2,1.4,0.2,Iris-setosa
+4.9,3.1,1.5,0.1,Iris-setosa
+5.0,3.2,1.2,0.2,Iris-setosa
+5.5,3.5,1.3,0.2,Iris-setosa
+4.9,3.1,1.5,0.1,Iris-setosa
+4.4,3.0,1.3,0.2,Iris-setosa
+5.1,3.4,1.5,0.2,Iris-setosa
+5.0,3.5,1.3,0.3,Iris-setosa
+4.5,2.3,1.3,0.3,Iris-setosa
+4.4,3.2,1.3,0.2,Iris-setosa
+5.0,3.5,1.6,0.6,Iris-setosa
+5.1,3.8,1.9,0.4,Iris-setosa
+4.8,3.0,1.4,0.3,Iris-setosa
+5.1,3.8,1.6,0.2,Iris-setosa
+4.6,3.2,1.4,0.2,Iris-setosa
+5.3,3.7,1.5,0.2,Iris-setosa
+5.0,3.3,1.4,0.2,Iris-setosa
+7.0,3.2,4.7,1.4,Iris-versicolor
+6.4,3.2,4.5,1.5,Iris-versicolor
+6.9,3.1,4.9,1.5,Iris-versicolor
+5.5,2.3,4.0,1.3,Iris-versicolor
+6.5,2.8,4.6,1.5,Iris-versicolor
+5.7,2.8,4.5,1.3,Iris-versicolor
+6.3,3.3,4.7,1.6,Iris-versicolor
+4.9,2.4,3.3,1.0,Iris-versicolor
+6.6,2.9,4.6,1.3,Iris-versicolor
+5.2,2.7,3.9,1.4,Iris-versicolor
+5.0,2.0,3.5,1.0,Iris-versicolor
+5.9,3.0,4.2,1.5,Iris-versicolor
+6.0,2.2,4.0,1.0,Iris-versicolor
+6.1,2.9,4.7,1.4,Iris-versicolor
+5.6,2.9,3.6,1.3,Iris-versicolor
+6.7,3.1,4.4,1.4,Iris-versicolor
+5.6,3.0,4.5,1.5,Iris-versicolor
+5.8,2.7,4.1,1.0,Iris-versicolor
+6.2,2.2,4.5,1.5,Iris-versicolor
+5.6,2.5,3.9,1.1,Iris-versicolor
+5.9,3.2,4.8,1.8,Iris-versicolor
+6.1,2.8,4.0,1.3,Iris-versicolor
+6.3,2.5,4.9,1.5,Iris-versicolor
+6.1,2.8,4.7,1.2,Iris-versicolor
+6.4,2.9,4.3,1.3,Iris-versicolor
+6.6,3.0,4.4,1.4,Iris-versicolor
+6.8,2.8,4.8,1.4,Iris-versicolor
+6.7,3.0,5.0,1.7,Iris-versicolor
+6.0,2.9,4.5,1.5,Iris-versicolor
+5.7,2.6,3.5,1.0,Iris-versicolor
+5.5,2.4,3.8,1.1,Iris-versicolor
+5.5,2.4,3.7,1.0,Iris-versicolor
+5.8,2.7,3.9,1.2,Iris-versicolor
+6.0,2.7,5.1,1.6,Iris-versicolor
+5.4,3.0,4.5,1.5,Iris-versicolor
+6.0,3.4,4.5,1.6,Iris-versicolor
+6.7,3.1,4.7,1.5,Iris-versicolor
+6.3,2.3,4.4,1.3,Iris-versicolor
+5.6,3.0,4.1,1.3,Iris-versicolor
+5.5,2.5,4.0,1.3,Iris-versicolor
+5.5,2.6,4.4,1.2,Iris-versicolor
+6.1,3.0,4.6,1.4,Iris-versicolor
+5.8,2.6,4.0,1.2,Iris-versicolor
+5.0,2.3,3.3,1.0,Iris-versicolor
+5.6,2.7,4.2,1.3,Iris-versicolor
+5.7,3.0,4.2,1.2,Iris-versicolor
+5.7,2.9,4.2,1.3,Iris-versicolor
+6.2,2.9,4.3,1.3,Iris-versicolor
+5.1,2.5,3.0,1.1,Iris-versicolor
+5.7,2.8,4.1,1.3,Iris-versicolor
+6.3,3.3,6.0,2.5,Iris-virginica
+5.8,2.7,5.1,1.9,Iris-virginica
+7.1,3.0,5.9,2.1,Iris-virginica
+6.3,2.9,5.6,1.8,Iris-virginica
+6.5,3.0,5.8,2.2,Iris-virginica
+7.6,3.0,6.6,2.1,Iris-virginica
+4.9,2.5,4.5,1.7,Iris-virginica
+7.3,2.9,6.3,1.8,Iris-virginica
+6.7,2.5,5.8,1.8,Iris-virginica
+7.2,3.6,6.1,2.5,Iris-virginica
+6.5,3.2,5.1,2.0,Iris-virginica
+6.4,2.7,5.3,1.9,Iris-virginica
+6.8,3.0,5.5,2.1,Iris-virginica
+5.7,2.5,5.0,2.0,Iris-virginica
+5.8,2.8,5.1,2.4,Iris-virginica
+6.4,3.2,5.3,2.3,Iris-virginica
+6.5,3.0,5.5,1.8,Iris-virginica
+7.7,3.8,6.7,2.2,Iris-virginica
+7.7,2.6,6.9,2.3,Iris-virginica
+6.0,2.2,5.0,1.5,Iris-virginica
+6.9,3.2,5.7,2.3,Iris-virginica
+5.6,2.8,4.9,2.0,Iris-virginica
+7.7,2.8,6.7,2.0,Iris-virginica
+6.3,2.7,4.9,1.8,Iris-virginica
+6.7,3.3,5.7,2.1,Iris-virginica
+7.2,3.2,6.0,1.8,Iris-virginica
+6.2,2.8,4.8,1.8,Iris-virginica
+6.1,3.0,4.9,1.8,Iris-virginica
+6.4,2.8,5.6,2.1,Iris-virginica
+7.2,3.0,5.8,1.6,Iris-virginica
+7.4,2.8,6.1,1.9,Iris-virginica
+7.9,3.8,6.4,2.0,Iris-virginica
+6.4,2.8,5.6,2.2,Iris-virginica
+6.3,2.8,5.1,1.5,Iris-virginica
+6.1,2.6,5.6,1.4,Iris-virginica
+7.7,3.0,6.1,2.3,Iris-virginica
+6.3,3.4,5.6,2.4,Iris-virginica
+6.4,3.1,5.5,1.8,Iris-virginica
+6.0,3.0,4.8,1.8,Iris-virginica
+6.9,3.1,5.4,2.1,Iris-virginica
+6.7,3.1,5.6,2.4,Iris-virginica
+6.9,3.1,5.1,2.3,Iris-virginica
+5.8,2.7,5.1,1.9,Iris-virginica
+6.8,3.2,5.9,2.3,Iris-virginica
+6.7,3.3,5.7,2.5,Iris-virginica
+6.7,3.0,5.2,2.3,Iris-virginica
+6.3,2.5,5.0,1.9,Iris-virginica
+6.5,3.0,5.2,2.0,Iris-virginica
+6.2,3.4,5.4,2.3,Iris-virginica
+5.9,3.0,5.1,1.8,Iris-virginica
diff --git a/tests/back/converters/iris_petal_cols_dropped.csv b/tests/back/converters/iris_petal_cols_dropped.csv
new file mode 100644
index 00000000..645ed279
--- /dev/null
+++ b/tests/back/converters/iris_petal_cols_dropped.csv
@@ -0,0 +1,151 @@
+SepalLengthCm,SepalWidthCm,Species
+5.1,3.5,Iris-setosa
+4.9,3.0,Iris-setosa
+4.7,3.2,Iris-setosa
+4.6,3.1,Iris-setosa
+5.0,3.6,Iris-setosa
+5.4,3.9,Iris-setosa
+4.6,3.4,Iris-setosa
+5.0,3.4,Iris-setosa
+4.4,2.9,Iris-setosa
+4.9,3.1,Iris-setosa
+5.4,3.7,Iris-setosa
+4.8,3.4,Iris-setosa
+4.8,3.0,Iris-setosa
+4.3,3.0,Iris-setosa
+5.8,4.0,Iris-setosa
+5.7,4.4,Iris-setosa
+5.4,3.9,Iris-setosa
+5.1,3.5,Iris-setosa
+5.7,3.8,Iris-setosa
+5.1,3.8,Iris-setosa
+5.4,3.4,Iris-setosa
+5.1,3.7,Iris-setosa
+4.6,3.6,Iris-setosa
+5.1,3.3,Iris-setosa
+4.8,3.4,Iris-setosa
+5.0,3.0,Iris-setosa
+5.0,3.4,Iris-setosa
+5.2,3.5,Iris-setosa
+5.2,3.4,Iris-setosa
+4.7,3.2,Iris-setosa
+4.8,3.1,Iris-setosa
+5.4,3.4,Iris-setosa
+5.2,4.1,Iris-setosa
+5.5,4.2,Iris-setosa
+4.9,3.1,Iris-setosa
+5.0,3.2,Iris-setosa
+5.5,3.5,Iris-setosa
+4.9,3.1,Iris-setosa
+4.4,3.0,Iris-setosa
+5.1,3.4,Iris-setosa
+5.0,3.5,Iris-setosa
+4.5,2.3,Iris-setosa
+4.4,3.2,Iris-setosa
+5.0,3.5,Iris-setosa
+5.1,3.8,Iris-setosa
+4.8,3.0,Iris-setosa
+5.1,3.8,Iris-setosa
+4.6,3.2,Iris-setosa
+5.3,3.7,Iris-setosa
+5.0,3.3,Iris-setosa
+7.0,3.2,Iris-versicolor
+6.4,3.2,Iris-versicolor
+6.9,3.1,Iris-versicolor
+5.5,2.3,Iris-versicolor
+6.5,2.8,Iris-versicolor
+5.7,2.8,Iris-versicolor
+6.3,3.3,Iris-versicolor
+4.9,2.4,Iris-versicolor
+6.6,2.9,Iris-versicolor
+5.2,2.7,Iris-versicolor
+5.0,2.0,Iris-versicolor
+5.9,3.0,Iris-versicolor
+6.0,2.2,Iris-versicolor
+6.1,2.9,Iris-versicolor
+5.6,2.9,Iris-versicolor
+6.7,3.1,Iris-versicolor
+5.6,3.0,Iris-versicolor
+5.8,2.7,Iris-versicolor
+6.2,2.2,Iris-versicolor
+5.6,2.5,Iris-versicolor
+5.9,3.2,Iris-versicolor
+6.1,2.8,Iris-versicolor
+6.3,2.5,Iris-versicolor
+6.1,2.8,Iris-versicolor
+6.4,2.9,Iris-versicolor
+6.6,3.0,Iris-versicolor
+6.8,2.8,Iris-versicolor
+6.7,3.0,Iris-versicolor
+6.0,2.9,Iris-versicolor
+5.7,2.6,Iris-versicolor
+5.5,2.4,Iris-versicolor
+5.5,2.4,Iris-versicolor
+5.8,2.7,Iris-versicolor
+6.0,2.7,Iris-versicolor
+5.4,3.0,Iris-versicolor
+6.0,3.4,Iris-versicolor
+6.7,3.1,Iris-versicolor
+6.3,2.3,Iris-versicolor
+5.6,3.0,Iris-versicolor
+5.5,2.5,Iris-versicolor
+5.5,2.6,Iris-versicolor
+6.1,3.0,Iris-versicolor
+5.8,2.6,Iris-versicolor
+5.0,2.3,Iris-versicolor
+5.6,2.7,Iris-versicolor
+5.7,3.0,Iris-versicolor
+5.7,2.9,Iris-versicolor
+6.2,2.9,Iris-versicolor
+5.1,2.5,Iris-versicolor
+5.7,2.8,Iris-versicolor
+6.3,3.3,Iris-virginica
+5.8,2.7,Iris-virginica
+7.1,3.0,Iris-virginica
+6.3,2.9,Iris-virginica
+6.5,3.0,Iris-virginica
+7.6,3.0,Iris-virginica
+4.9,2.5,Iris-virginica
+7.3,2.9,Iris-virginica
+6.7,2.5,Iris-virginica
+7.2,3.6,Iris-virginica
+6.5,3.2,Iris-virginica
+6.4,2.7,Iris-virginica
+6.8,3.0,Iris-virginica
+5.7,2.5,Iris-virginica
+5.8,2.8,Iris-virginica
+6.4,3.2,Iris-virginica
+6.5,3.0,Iris-virginica
+7.7,3.8,Iris-virginica
+7.7,2.6,Iris-virginica
+6.0,2.2,Iris-virginica
+6.9,3.2,Iris-virginica
+5.6,2.8,Iris-virginica
+7.7,2.8,Iris-virginica
+6.3,2.7,Iris-virginica
+6.7,3.3,Iris-virginica
+7.2,3.2,Iris-virginica
+6.2,2.8,Iris-virginica
+6.1,3.0,Iris-virginica
+6.4,2.8,Iris-virginica
+7.2,3.0,Iris-virginica
+7.4,2.8,Iris-virginica
+7.9,3.8,Iris-virginica
+6.4,2.8,Iris-virginica
+6.3,2.8,Iris-virginica
+6.1,2.6,Iris-virginica
+7.7,3.0,Iris-virginica
+6.3,3.4,Iris-virginica
+6.4,3.1,Iris-virginica
+6.0,3.0,Iris-virginica
+6.9,3.1,Iris-virginica
+6.7,3.1,Iris-virginica
+6.9,3.1,Iris-virginica
+5.8,2.7,Iris-virginica
+6.8,3.2,Iris-virginica
+6.7,3.3,Iris-virginica
+6.7,3.0,Iris-virginica
+6.3,2.5,Iris-virginica
+6.5,3.0,Iris-virginica
+6.2,3.4,Iris-virginica
+5.9,3.0,Iris-virginica
diff --git a/tests/back/converters/iris_petal_width_dropped.csv b/tests/back/converters/iris_petal_width_dropped.csv
new file mode 100644
index 00000000..0cbf3869
--- /dev/null
+++ b/tests/back/converters/iris_petal_width_dropped.csv
@@ -0,0 +1,151 @@
+SepalLengthCm,SepalWidthCm,PetalLengthCm,Species
+5.1,3.5,1.4,Iris-setosa
+4.9,3.0,1.4,Iris-setosa
+4.7,3.2,1.3,Iris-setosa
+4.6,3.1,1.5,Iris-setosa
+5.0,3.6,1.4,Iris-setosa
+5.4,3.9,1.7,Iris-setosa
+4.6,3.4,1.4,Iris-setosa
+5.0,3.4,1.5,Iris-setosa
+4.4,2.9,1.4,Iris-setosa
+4.9,3.1,1.5,Iris-setosa
+5.4,3.7,1.5,Iris-setosa
+4.8,3.4,1.6,Iris-setosa
+4.8,3.0,1.4,Iris-setosa
+4.3,3.0,1.1,Iris-setosa
+5.8,4.0,1.2,Iris-setosa
+5.7,4.4,1.5,Iris-setosa
+5.4,3.9,1.3,Iris-setosa
+5.1,3.5,1.4,Iris-setosa
+5.7,3.8,1.7,Iris-setosa
+5.1,3.8,1.5,Iris-setosa
+5.4,3.4,1.7,Iris-setosa
+5.1,3.7,1.5,Iris-setosa
+4.6,3.6,1.0,Iris-setosa
+5.1,3.3,1.7,Iris-setosa
+4.8,3.4,1.9,Iris-setosa
+5.0,3.0,1.6,Iris-setosa
+5.0,3.4,1.6,Iris-setosa
+5.2,3.5,1.5,Iris-setosa
+5.2,3.4,1.4,Iris-setosa
+4.7,3.2,1.6,Iris-setosa
+4.8,3.1,1.6,Iris-setosa
+5.4,3.4,1.5,Iris-setosa
+5.2,4.1,1.5,Iris-setosa
+5.5,4.2,1.4,Iris-setosa
+4.9,3.1,1.5,Iris-setosa
+5.0,3.2,1.2,Iris-setosa
+5.5,3.5,1.3,Iris-setosa
+4.9,3.1,1.5,Iris-setosa
+4.4,3.0,1.3,Iris-setosa
+5.1,3.4,1.5,Iris-setosa
+5.0,3.5,1.3,Iris-setosa
+4.5,2.3,1.3,Iris-setosa
+4.4,3.2,1.3,Iris-setosa
+5.0,3.5,1.6,Iris-setosa
+5.1,3.8,1.9,Iris-setosa
+4.8,3.0,1.4,Iris-setosa
+5.1,3.8,1.6,Iris-setosa
+4.6,3.2,1.4,Iris-setosa
+5.3,3.7,1.5,Iris-setosa
+5.0,3.3,1.4,Iris-setosa
+7.0,3.2,4.7,Iris-versicolor
+6.4,3.2,4.5,Iris-versicolor
+6.9,3.1,4.9,Iris-versicolor
+5.5,2.3,4.0,Iris-versicolor
+6.5,2.8,4.6,Iris-versicolor
+5.7,2.8,4.5,Iris-versicolor
+6.3,3.3,4.7,Iris-versicolor
+4.9,2.4,3.3,Iris-versicolor
+6.6,2.9,4.6,Iris-versicolor
+5.2,2.7,3.9,Iris-versicolor
+5.0,2.0,3.5,Iris-versicolor
+5.9,3.0,4.2,Iris-versicolor
+6.0,2.2,4.0,Iris-versicolor
+6.1,2.9,4.7,Iris-versicolor
+5.6,2.9,3.6,Iris-versicolor
+6.7,3.1,4.4,Iris-versicolor
+5.6,3.0,4.5,Iris-versicolor
+5.8,2.7,4.1,Iris-versicolor
+6.2,2.2,4.5,Iris-versicolor
+5.6,2.5,3.9,Iris-versicolor
+5.9,3.2,4.8,Iris-versicolor
+6.1,2.8,4.0,Iris-versicolor
+6.3,2.5,4.9,Iris-versicolor
+6.1,2.8,4.7,Iris-versicolor
+6.4,2.9,4.3,Iris-versicolor
+6.6,3.0,4.4,Iris-versicolor
+6.8,2.8,4.8,Iris-versicolor
+6.7,3.0,5.0,Iris-versicolor
+6.0,2.9,4.5,Iris-versicolor
+5.7,2.6,3.5,Iris-versicolor
+5.5,2.4,3.8,Iris-versicolor
+5.5,2.4,3.7,Iris-versicolor
+5.8,2.7,3.9,Iris-versicolor
+6.0,2.7,5.1,Iris-versicolor
+5.4,3.0,4.5,Iris-versicolor
+6.0,3.4,4.5,Iris-versicolor
+6.7,3.1,4.7,Iris-versicolor
+6.3,2.3,4.4,Iris-versicolor
+5.6,3.0,4.1,Iris-versicolor
+5.5,2.5,4.0,Iris-versicolor
+5.5,2.6,4.4,Iris-versicolor
+6.1,3.0,4.6,Iris-versicolor
+5.8,2.6,4.0,Iris-versicolor
+5.0,2.3,3.3,Iris-versicolor
+5.6,2.7,4.2,Iris-versicolor
+5.7,3.0,4.2,Iris-versicolor
+5.7,2.9,4.2,Iris-versicolor
+6.2,2.9,4.3,Iris-versicolor
+5.1,2.5,3.0,Iris-versicolor
+5.7,2.8,4.1,Iris-versicolor
+6.3,3.3,6.0,Iris-virginica
+5.8,2.7,5.1,Iris-virginica
+7.1,3.0,5.9,Iris-virginica
+6.3,2.9,5.6,Iris-virginica
+6.5,3.0,5.8,Iris-virginica
+7.6,3.0,6.6,Iris-virginica
+4.9,2.5,4.5,Iris-virginica
+7.3,2.9,6.3,Iris-virginica
+6.7,2.5,5.8,Iris-virginica
+7.2,3.6,6.1,Iris-virginica
+6.5,3.2,5.1,Iris-virginica
+6.4,2.7,5.3,Iris-virginica
+6.8,3.0,5.5,Iris-virginica
+5.7,2.5,5.0,Iris-virginica
+5.8,2.8,5.1,Iris-virginica
+6.4,3.2,5.3,Iris-virginica
+6.5,3.0,5.5,Iris-virginica
+7.7,3.8,6.7,Iris-virginica
+7.7,2.6,6.9,Iris-virginica
+6.0,2.2,5.0,Iris-virginica
+6.9,3.2,5.7,Iris-virginica
+5.6,2.8,4.9,Iris-virginica
+7.7,2.8,6.7,Iris-virginica
+6.3,2.7,4.9,Iris-virginica
+6.7,3.3,5.7,Iris-virginica
+7.2,3.2,6.0,Iris-virginica
+6.2,2.8,4.8,Iris-virginica
+6.1,3.0,4.9,Iris-virginica
+6.4,2.8,5.6,Iris-virginica
+7.2,3.0,5.8,Iris-virginica
+7.4,2.8,6.1,Iris-virginica
+7.9,3.8,6.4,Iris-virginica
+6.4,2.8,5.6,Iris-virginica
+6.3,2.8,5.1,Iris-virginica
+6.1,2.6,5.6,Iris-virginica
+7.7,3.0,6.1,Iris-virginica
+6.3,3.4,5.6,Iris-virginica
+6.4,3.1,5.5,Iris-virginica
+6.0,3.0,4.8,Iris-virginica
+6.9,3.1,5.4,Iris-virginica
+6.7,3.1,5.6,Iris-virginica
+6.9,3.1,5.1,Iris-virginica
+5.8,2.7,5.1,Iris-virginica
+6.8,3.2,5.9,Iris-virginica
+6.7,3.3,5.7,Iris-virginica
+6.7,3.0,5.2,Iris-virginica
+6.3,2.5,5.0,Iris-virginica
+6.5,3.0,5.2,Iris-virginica
+6.2,3.4,5.4,Iris-virginica
+5.9,3.0,5.1,Iris-virginica
diff --git a/tests/back/converters/test_drop_column_converters.py b/tests/back/converters/test_drop_column_converters.py
new file mode 100644
index 00000000..074ead4d
--- /dev/null
+++ b/tests/back/converters/test_drop_column_converters.py
@@ -0,0 +1,163 @@
+import io
+
+import pytest
+from datasets import DatasetDict
+from starlette.datastructures import UploadFile
+
+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
+
+
+@pytest.fixture(name="iris_dataset")
+def prepare_iris_dataset():
+    test_dataset_path = "tests/back/converters/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 = dataloader_test.load_data(
+        filepath_or_buffer=file,
+        temp_path="tests/back/converters",
+        params={"separator": ","},
+    )
+
+    inputs_columns = [
+        "SepalLengthCm",
+        "SepalWidthCm",
+        "PetalLengthCm",
+        "PetalWidthCm",
+    ]
+
+    datasetdict = to_dashai_dataset(
+        datasetdict,
+        inputs_columns,
+        outputs_columns=["Species"],
+    )
+
+    return datasetdict
+
+
+@pytest.fixture(name="iris_dataset_petal_width_dropped")
+def prepare_iris_petal_width_dropped_dataset():
+    test_dataset_path = "tests/back/converters/iris_petal_width_dropped.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 = dataloader_test.load_data(
+        filepath_or_buffer=file,
+        temp_path="tests/back/converters",
+        params={"separator": ","},
+    )
+
+    inputs_columns = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm"]
+
+    datasetdict = to_dashai_dataset(
+        datasetdict,
+        inputs_columns,
+        outputs_columns=["Species"],
+    )
+
+    return datasetdict
+
+
+@pytest.fixture(name="iris_dataset_petal_cols_dropped")
+def prepare_iris_dataset_petal_cols_dropped():
+    test_dataset_path = "tests/back/converters/iris_petal_cols_dropped.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 = dataloader_test.load_data(
+        filepath_or_buffer=file,
+        temp_path="tests/back/converters",
+        params={"separator": ","},
+    )
+
+    inputs_columns = ["SepalLengthCm", "SepalWidthCm"]
+
+    datasetdict = to_dashai_dataset(
+        datasetdict,
+        inputs_columns,
+        outputs_columns=["Species"],
+    )
+
+    return datasetdict
+
+
+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:
+        dataset_split: DashAIDataset = dataset_obtained[split]
+        iris_dataset_dropped_split: DashAIDataset = iris_dataset_petal_width_dropped[
+            split
+        ]
+        assert len(dataset_split) == len(iris_dataset_dropped_split)
+        for column_name in dataset_split.column_names:
+            assert dataset_split[column_name] == iris_dataset_dropped_split[column_name]
+        assert dataset_split.features == iris_dataset_dropped_split.features
+
+
+def test_remove_input_column_with_index(
+    iris_dataset: DatasetDict, iris_dataset_petal_width_dropped: DatasetDict
+):
+    dropper = ColumnDropperByIndex(columns_index=3)
+    dataset_obtained = dropper.transform(iris_dataset)
+    assert set(dataset_obtained.keys()) == set(iris_dataset_petal_width_dropped.keys())
+    for split in dataset_obtained:
+        dataset_split: DashAIDataset = dataset_obtained[split]
+        iris_dataset_dropped_split: DashAIDataset = iris_dataset_petal_width_dropped[
+            split
+        ]
+        assert len(dataset_split) == len(iris_dataset_dropped_split)
+        for column_name in dataset_split.column_names:
+            assert dataset_split[column_name] == iris_dataset_dropped_split[column_name]
+        assert dataset_split.features == iris_dataset_dropped_split.features
+
+
+def test_remove_2_input_columns_with_column_names(
+    iris_dataset: DatasetDict, iris_dataset_petal_cols_dropped: DatasetDict
+):
+    dropper = ColumnDropperByName(column_names=["PetalLengthCm", "PetalWidthCm"])
+    dataset_obtained = dropper.transform(iris_dataset)
+    assert set(dataset_obtained.keys()) == set(iris_dataset_petal_cols_dropped.keys())
+    for split in dataset_obtained:
+        dataset_split: DashAIDataset = dataset_obtained[split]
+        iris_dataset_dropped_split: DashAIDataset = iris_dataset_petal_cols_dropped[
+            split
+        ]
+        assert len(dataset_split) == len(iris_dataset_dropped_split)
+        for column_name in dataset_split.column_names:
+            assert dataset_split[column_name] == iris_dataset_dropped_split[column_name]
+        assert dataset_split.features == iris_dataset_dropped_split.features
+
+
+def test_remove_2_input_columns_with_index(
+    iris_dataset: DatasetDict, iris_dataset_petal_cols_dropped: DatasetDict
+):
+    dropper = ColumnDropperByIndex(columns_index=(2, 3))
+    dataset_obtained = dropper.transform(iris_dataset)
+    assert set(dataset_obtained.keys()) == set(iris_dataset_petal_cols_dropped.keys())
+    for split in dataset_obtained:
+        dataset_split: DashAIDataset = dataset_obtained[split]
+        iris_dataset_dropped_split: DashAIDataset = iris_dataset_petal_cols_dropped[
+            split
+        ]
+        assert len(dataset_split) == len(iris_dataset_dropped_split)
+        for column_name in dataset_split.column_names:
+            assert dataset_split[column_name] == iris_dataset_dropped_split[column_name]
+        assert dataset_split.features == iris_dataset_dropped_split.features