Repository: giotto-ai/giotto-deep (language: python)

Issue:
Generic preprocessing module
**Is your feature request related to a problem? Please describe.**
<!--
A clear and concise description of what the problem is. Eg. "I'm always frustrated when [...]"
-->

The pain is that, most often, plain datasets are not in the right input format or do not have the designed statistical caracterisrtics. Furthermore, standard techniques like data augmentation, need to be implemented
**Describe the solution you'd like**
<!--
A clear and concise description of what you want to happen. If you have ideas on how to implement
your solution with (pseudo-)code, include them here.
-->

We build an API class (`AbstractClass` ) for the preprocessing -- a generic one.

It should look similar to this one:
```
from abc import ABC, abstractmethod

class AbstractPreprocessing(ABC):
    """The abstract class to define the interface of preprocessing
    """
    @abstractmethod
    def __call__(self, *args, **kwargs):
        """This method deals with datum-wise transformations"""
        pass

    @abstractmethod
    def fit_to_data(self, *args, **kwargs):
        """This method deals with getting dataset-level information"""
        pass

```

Each of the methods shall be implemented, as it will be called automatically inside the `Dataset` classes:
 1. the output of `__getitem__` will be transformed by `item_transform`. the data inside `item_transform` that are needed to perform the transformation, will be stored in self. The methods `dataset_level_data` and `batch_level_data` will be called only once, before the first time that `__getitem__` is called.
 2. the big advantage of this "on the fly" approach is that it will save a lot of memory -- and hopefully the transformations are not to heavy to compute
 3. In case the transformations are computationally very heavy, then it would be advisable first to transform all the data (not via this class! just create a new dataset) and then use that for the next steps.

**Describe alternatives you've considered**
<!--
A clear and concise description of any alternative solutions or features you've considered.
-->

Only doing point 3 above (without 1 and 2), however I find it is always possible to only use that approach and it is much easier to implement and is less bind to the generic pipeline

**Additional context**
<!--
Add any other context or screenshots about the feature request here.
-->

<!-- Thanks for contributing! -->

Dataset refactoring polishing



Discussion / hints:
Thanks, Matteo, for the suggestion.
I think I understand what the dataset_level_data is for. But I'm not sure what batch_level_data is for.
Could you give an example of what would go in batch_level_data?

When you train a model with preprocessed data, the transformations you apply are crucial for the model; hence you want to have a way of storing them and loading them later on. This is especially important when you want to deploy your model in a production environment, where you will not have access to the preprocessing transformations. 
This can be quickly done by inheriting from the Huggingface feature extractor mixin class and adding a https://huggingface.co/docs/transformers/main_classes/feature_extractor.
Another advantage of inheriting from that class is that developers are already familiar with it, and they will be able to understand your code more easily. It would be great to be compatible with the Huggingface API since developers can efficiently other models from the library and plug them into your framework.
The preprocessing transforms in gtda.diagrams.preprocessing don't allow for normalizing the data as well as filtering out the k-most persistent points.
I also looked at the implementation of the filtration by thresholding in gtda.homology._utils
https://github.com/giotto-ai/giotto-tda/blob/8d09a39403ca11b50605bf466c1aa9f4f3876e5f/gtda/diagrams/_utils.py#L80
and it seems like their implementation does not work for extended persistence diagrams and one-hot encoded homology dimensions.
I also don't understand the implementation; it looks much more complicated than it needs to be.


Fix patch:
diff --git a/.gitignore b/.gitignore
index e291b062..600a9a41 100755
--- a/.gitignore
+++ b/.gitignore
@@ -88,3 +88,6 @@ state_dicts/*
 # Include pkl test data
 !gdeep/topology_layers/tests/data/diagrams.pkl
 !gdeep/topology_layers/tests/data/labels.pkl
+
+# pylance files
+*.pyi
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 8ffe091c..b291cd36 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -20,7 +20,7 @@
 # -- Project information -----------------------------------------------------
 
 project = 'giotto-deep'
-copyright = '2021, L2F SA'
+copyright = '2022, L2F SA'
 author = 'Matteo Caorsi'
 
 # The full version, including alpha/beta/rc tags
diff --git a/docs/source/modules/bert.rst b/docs/source/modules/bert.rst
deleted file mode 100644
index c7be75a8..00000000
--- a/docs/source/modules/bert.rst
+++ /dev/null
@@ -1,6 +0,0 @@
-BERT
-====
-
-.. automodule:: gdeep.bert
-   :members:
-   
diff --git a/docs/source/modules/index.rst b/docs/source/modules/index.rst
index 59fff795..6fb1204f 100644
--- a/docs/source/modules/index.rst
+++ b/docs/source/modules/index.rst
@@ -7,14 +7,12 @@ This pages contains a list of available features in the library.
    :maxdepth: 3
 
    analysis
-
-   bert
    
    data
    
    models
    
-   pipeline
+   trainer
    
    search
 
diff --git a/docs/source/modules/pipeline.rst b/docs/source/modules/pipeline.rst
deleted file mode 100644
index 135fa21e..00000000
--- a/docs/source/modules/pipeline.rst
+++ /dev/null
@@ -1,6 +0,0 @@
-Pipeline
-========
-
-.. automodule:: gdeep.pipeline
-   :members:
-   
diff --git a/docs/source/modules/trainer.rst b/docs/source/modules/trainer.rst
new file mode 100644
index 00000000..f4dee4ed
--- /dev/null
+++ b/docs/source/modules/trainer.rst
@@ -0,0 +1,6 @@
+Trainer
+=======
+
+.. automodule:: gdeep.trainer
+   :members:
+   
diff --git a/examples/Ensemble_learning.ipynb b/examples/Ensemble_learning.ipynb
index 1f9e2d18..4fff4ad8 100644
--- a/examples/Ensemble_learning.ipynb
+++ b/examples/Ensemble_learning.ipynb
@@ -25,8 +25,8 @@
     "from torch.utils.tensorboard import SummaryWriter\n",
     "from gdeep.models import FFNet\n",
     "from torchensemble import VotingClassifier\n",
-    "from gdeep.data import TorchDataLoader\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
+    "from gdeep.trainer import Trainer\n",
     "from gdeep.utility.optimisation import SAMOptimizer\n",
     "from torch.optim.lr_scheduler import ExponentialLR\n",
     "from torch.optim import SGD, Adam\n",
@@ -75,9 +75,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = TorchDataLoader(name=\"DoubleTori\")\n",
-    "dl_tr, dl_ts = dl.build_dataloaders(batch_size=3)\n",
-    "dl_ts = None"
+    "bd = DatasetBuilder(name=\"DoubleTori\")\n",
+    "ds_tr, ds_val, _ = bd.build()\n",
+    "#train_indices = list(range(160))\n",
+    "dl = DataLoaderBuilder((ds_tr, ds_val))\n",
+    "dl_tr, dl_val, dl_ts = dl.build(({\"batch_size\":23},{\"batch_size\":23}))"
    ]
   },
   {
@@ -169,7 +171,7 @@
     "# initlaise the loss function\n",
     "loss_fn = nn.CrossEntropyLoss()\n",
     "# initialise the pipelien class\n",
-    "pipe = Pipeline(model, (dl_tr, dl_ts), loss_fn, writer)\n",
+    "pipe = Trainer(model, (dl_tr, dl_ts), loss_fn, writer)\n",
     "\n",
     "# initialise the SAM optimiser\n",
     "optim = SAMOptimizer(SGD)  # this is a class, not an instance!\n",
@@ -206,7 +208,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3d41c4cc-8dae-429b-a14d-1bcf3476f012",
+   "id": "a9d113dd-a654-4887-9ed9-f1c186d836ca",
    "metadata": {},
    "outputs": [],
    "source": []
diff --git a/examples/Intersection_Homology.ipynb b/examples/Intersection_Homology.ipynb
deleted file mode 100644
index 77306099..00000000
--- a/examples/Intersection_Homology.ipynb
+++ /dev/null
@@ -1,296 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# import libraries\n",
-    "import numpy as np\n",
-    "import plotly.express as px\n",
-    "import networkx as nx\n",
-    "from networkx.algorithms.clique import find_cliques, enumerate_all_cliques\n",
-    "from sklearn.neighbors import kneighbors_graph\n",
-    "from scipy.sparse import csr_matrix\n",
-    "import matplotlib.pyplot as plt\n",
-    "from networkx.drawing.nx_pylab import draw_kamada_kawai\n",
-    "import warnings\n",
-    "from networkx.algorithms.components import number_connected_components\n",
-    "from itertools import combinations, chain"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# import IH classes\n",
-    "from gdeep.utility.intersection_homology import IntersectionHomology, compute_coarsest_strata\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Intersection homology \n",
-    "\n",
-    "This notebooks shows how to use compoute intersection homology for a given perversity. The perversity can be user-defined. We have two approaches to finding a suitable stratification:\n",
-    "\n",
-    "1. Use the coarsest possible stratification\n",
-    "2. Use the singularity loci stratification as if being defined in the continuous setting (default)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Perversities\n",
-    "Let $X$ be a stratified space with strata $X_d$.\n",
-    "\n",
-    "An i-simplex $\\sigma$ in $K$ is said to be $p$-proper if the following condition holds for all $k = 0,...,d$:\n",
-    "\n",
-    "$$dim(\\bar\\sigma  \\cap X_{d−k}) \\leq i − k + p(k)$$\n",
-    "\n",
-    "where $\\bar \\sigma$ denotes the closure of the open simplex $\\sigma$, and $p:\\{0,1,...,n\\}\\to \\{0,1,...,n\\}$ is a perversisty. Notice that a perversity satisfies $p(k+1)=p(k)$ or $p(k)+1$"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Coarsest Stratification\n",
-    "$B_\\sigma := St(\\sigma) = \\{\\tau \\in K | \\sigma \\leq \\tau\\}$. Using the partial order induced by inclusions of minimal open neighborhoods, we set $\\tau \\leq \\sigma$ if $B_\\sigma \\subset B_\\tau$ \n",
-    "\n",
-    "### Algorithm\n",
-    "1. Set $i=0$, $d_0 =dim X$, $X_{d_0} = X$,and initialize $S_j = \\varnothing$,forall $0\\leq j\\leq d_0$. \n",
-    "2. While $d_i \\geq 0$, do\n",
-    "\n",
-    "   (a) For each $x\\in X_{d_i}$,set $S_{d_i} = S_{d_i} \\cup x$ if $\\delta(w\\leq y)=1$,$\\forall$ adjacent pairs $w\\leq y$ in $B_x \\cap X_{d_i}$ \n",
-    "   \n",
-    "   (b) Set $d_{i+1} = dim(X_{d_i} − S_{d_i} )$\n",
-    "   \n",
-    "   (c) Define $X_{d_{i+1}} = X_{d_i} − S_{d_i}$ \n",
-    "   \n",
-    "   (d) Set $i=i+1$\n",
-    "   \n",
-    "3. Return $S$"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Computing standard homology of clique complex"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# random point cloud\n",
-    "X = np.random.random((15,3))\n",
-    "# we will build the knn graph with knn=3 and compute the homology of th clique complex\n",
-    "ih = IntersectionHomology(knn=3, perversity_check=False)\n",
-    "\n",
-    "#ih.fit(X)\n",
-    "#print(ih.simplices)\n",
-    "homology = ih.fit_transform(X)\n",
-    "print(homology)\n",
-    "ih.plot()\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Computing intersection homology with top perversty"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "# we will build the knn graph with knn=3 and compute the homology of th clique complex\n",
-    "ih = IntersectionHomology(knn=3)\n",
-    "\n",
-    "#ih.fit(X)\n",
-    "#print(ih.simplices)\n",
-    "homology = ih.fit_transform(X)\n",
-    "print(homology)\n",
-    "ih.plot()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Computing with dual perversity"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ih2 = IntersectionHomology(knn=3,perversity=IntersectionHomology.q)\n",
-    "\n",
-    "#ih.fit(X)\n",
-    "#print(ih.simplices)\n",
-    "homology = ih2.fit_transform(X)\n",
-    "print(homology)\n",
-    "ih2.plot()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# build a graph from edges: a bouquet of 2 circles\n",
-    "G = nx.Graph()\n",
-    "edges = [(0,1),(1, 2), (0, 3),(2,3),(0,4),(0,5),(4,6),(5,6)]\n",
-    "G.add_edges_from(edges)\n",
-    "ih = IntersectionHomology()\n",
-    "homology = ih.fit_transform(G)\n",
-    "print(\"Simplices: \",ih.simplices)\n",
-    "ih.plot()\n",
-    "homology"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 2 2-spheres\n",
-    "G = nx.Graph()\n",
-    "edges = [(0,1),(1, 2), (0, 3),(2,3),(0,4),(0,5),(4,6),(5,6),(1,3),(0,2),(0,6),(4,5)]\n",
-    "G.add_edges_from(edges)\n",
-    "\n",
-    "ih = IntersectionHomology()\n",
-    "ih.fit(G)\n",
-    "# we remove the tetrahedra from the proper simplices, as we are interestd in checking if\n",
-    "# we are able to find duality between H0 and H2\n",
-    "lst = ih.list_proper[:-2]\n",
-    "homology = ih.transform(G,lst)\n",
-    "print(homology)\n",
-    "\n",
-    "print(\"dual perversity: \")\n",
-    "ih2 = IntersectionHomology(perversity=IntersectionHomology.q)\n",
-    "ih2.fit(G)\n",
-    "# removing 2 tetrehedra (which happen to be the last simplices in the list_proper)\n",
-    "lst2 = ih.list_proper[:-2]\n",
-    "homology2 = ih2.transform(G,lst2)\n",
-    "ih2.plot()\n",
-    "print(homology2)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 2 3-spheres\n",
-    "G = nx.Graph()\n",
-    "edges = [(0,1),(0,2),(0,3),(0,4),(1,2),(1,3),(1,4),(2,3),(2,4),(3,4),#sphere 1\n",
-    "         (0,9),\n",
-    "         (9,5),(9,6),(9,7),(9,8),(5,6),(5,7),(5,8),(6,7),(6,8),(7,8)]# sphere 2\n",
-    "G.add_edges_from(edges)\n",
-    "\n",
-    "ih = IntersectionHomology()\n",
-    "ih.fit(G)\n",
-    "# remove the hyper-tetrahedra\n",
-    "lst = ih.list_proper[:-2]\n",
-    "homology = ih.transform(G,lst)\n",
-    "ih.plot()\n",
-    "homology\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# cross check: 3 disconnected circles\n",
-    "G = nx.Graph()\n",
-    "edges = [(0,1),(1, 2), (2, 3),(3,0),(4,5),(5,6),(6,7),(7,4),(8,9),(9,10),(10,11),(11,8)]\n",
-    "G.add_edges_from(edges)\n",
-    "\n",
-    "ih = IntersectionHomology()\n",
-    "homology = ih.fit_transform(G)\n",
-    "ih.plot()\n",
-    "homology\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# pinched torus\n",
-    "G = nx.Graph()\n",
-    "edges = [(0,1),(1, 2), (2, 3),(3,0),(4,5),(5,6),(6,7),(7,4),(8,9),(9,10),(10,11),(11,8),# 3 loops with 4 vertices each\n",
-    "         (0,12),(1,12),(2,12),(3,12),(8,12),(9,12),(10,12),(11,12),# singular point\n",
-    "         (0,4),(1,5),(2,6),(3,7),(4,8),(5,9),(6,10),(7,11), #straight connections\n",
-    "         (0,5),(1,6),(2,7),(3,4),(4,9),(5,10),(6,11),(7,8)] # oblique conections\n",
-    "G.add_edges_from(edges)\n",
-    "\n",
-    "ih = IntersectionHomology()\n",
-    "homology = ih.fit_transform(G)\n",
-    "ih.plot()\n",
-    "print(\"Homology groups representatives\")\n",
-    "homology\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.9.7"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}
diff --git a/examples/basic_tutorial_QandA.ipynb b/examples/basic_tutorial_QandA.ipynb
index 15cd2363..d3492e36 100644
--- a/examples/basic_tutorial_QandA.ipynb
+++ b/examples/basic_tutorial_QandA.ipynb
@@ -39,8 +39,8 @@
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
     "\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data import TorchDataLoader\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.data.datasets import DatasetBuilder\n",
+    "from gdeep.trainer import Trainer\n",
     "\n",
     "from gtda.diagrams import BettiCurve\n",
     "\n",
@@ -87,11 +87,9 @@
    "source": [
     "from torch.utils.data.sampler import SubsetRandomSampler\n",
     "\n",
-    "# the only part of the training set we are interested in\n",
-    "train_indices = list(range(32*10))\n",
     "\n",
-    "dl = TorchDataLoader(name=\"SQuAD2\", convert_to_map_dataset=True)\n",
-    "dl_tr_str, dl_ts_str = dl.build_dataloaders(sampler=SubsetRandomSampler(train_indices), batch_size=1)\n"
+    "bd = DatasetBuilder(name=\"SQuAD2\", convert_to_map_dataset=True)\n",
+    "ds_tr_str, ds_val_str, ds_ts_str = bd.build()\n"
    ]
   },
   {
@@ -107,7 +105,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(\"Before preprocessing: \\n\", dl_tr_str.dataset[0])\n"
+    "print(\"Before preprocessing: \\n\", ds_tr_str[0])\n"
    ]
   },
   {
@@ -130,25 +128,37 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.data import TextDatasetQA, PreprocessTextQA, PreprocessTextQATarget\n",
-    "from torch.utils.data import DataLoader\n",
+    "from gdeep.data import PreprocessingPipeline\n",
+    "from gdeep.data import TransformingDataset\n",
+    "from gdeep.data.preprocessors import Normalization, TokenizerQA\n",
+    "from gdeep.data.datasets import  DataLoaderBuilder\n",
     "\n",
-    "prec = PreprocessTextQA()\n",
-    "textds = TextDatasetQA(dl_tr_str.dataset,\n",
-    "                     prec, \n",
-    "                     PreprocessTextQATarget())\n",
     "\n",
-    "textts = TextDatasetQA(dl_ts_str.dataset,\n",
-    "                     prec, \n",
-    "                     PreprocessTextQATarget())\n",
+    "tokenizer = TokenizerQA()\n",
     "\n",
-    "print(\"After the preprocessing: \\n\", textds[0])\n",
+    "# in case you need to combine multiple preprocessing:\n",
+    "#ppp = PreprocessingPipeline(((PreprocessTextData(), IdentityTransform(), TextDataset), \n",
+    "#                             (Normalisation(), IdentityTransform(), BasicDataset)))\n",
     "\n",
+    "\n",
+    "\n",
+    "tokenizer.fit_to_dataset(ds_tr_str)\n",
+    "transformed_textds = tokenizer.attach_transform_to_dataset(ds_tr_str)\n",
+    "\n",
+    "transformed_textts = tokenizer.attach_transform_to_dataset(ds_val_str)  # this has been fitted on the train set!\n",
+    "\n",
+    "print(\"After the preprocessing: \\n\", transformed_textds[0])\n",
+    "\n",
+    "# the only part of the training/test set we are interested in\n",
     "train_indices = list(range(64*2))\n",
     "test_indices = list(range(64*1))\n",
     "\n",
-    "dl_tr = DataLoader(textds, batch_size=16, sampler=SubsetRandomSampler(train_indices))\n",
-    "dl_ts = DataLoader(textts, batch_size=16, sampler=SubsetRandomSampler(test_indices))"
+    "dl_tr2, dl_ts2, _ = DataLoaderBuilder((transformed_textds, \n",
+    "                                       transformed_textts)).build(({\"batch_size\":16, \n",
+    "                                                                    \"sampler\":SubsetRandomSampler(train_indices)},\n",
+    "                                                                   {\"batch_size\":16, \n",
+    "                                                                    \"sampler\":SubsetRandomSampler(test_indices)}))\n",
+    "\n"
    ]
   },
   {
@@ -273,7 +283,7 @@
    "outputs": [],
    "source": [
     "# prepare a pipeline class with the model, dataloaders loss_fn and tensorboard writer\n",
-    "pipe = Pipeline(model, (dl_tr, dl_ts), loss_fn, writer)\n",
+    "pipe = Trainer(model, (dl_tr2, dl_ts2), loss_fn, writer)\n",
     "\n",
     "# train the model\n",
     "pipe.train(SGD, 3, False, {\"lr\":0.01}, {\"batch_size\":16})"
@@ -294,7 +304,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "bb = next(iter(dl_ts_str))\n",
+    "bb = next(iter(ds_val_str))\n",
     "bb[:2]"
    ]
   },
@@ -312,9 +322,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "voc = prec.vocabulary\n",
-    "context = prec.tokenizer(bb[0][0])\n",
-    "question = prec.tokenizer(bb[1][0])\n",
+    "voc = tokenizer.vocabulary\n",
+    "context = tokenizer.tokenizer(bb[0])\n",
+    "question = tokenizer.tokenizer(bb[1])\n",
     "\n",
     "# get the indexes in the vocabulary of the tokens\n",
     "context_idx = torch.tensor(list(map(voc.__getitem__,context)))\n",
@@ -327,9 +337,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "aa = next(iter(dl_tr))\n",
+    "aa = next(iter(dl_tr2))\n",
     "length_to_pad = aa[0].shape[-1]\n",
-    "pad_fn = lambda item : torch.cat([item, prec.pad_item * torch.ones(length_to_pad - item.shape[0])])\n",
+    "pad_fn = lambda item : torch.cat([item, tokenizer.pad_item * torch.ones(length_to_pad - item.shape[0])]).to(torch.long)\n",
     "\n",
     "# these tansors are ready to be fitted into the model\n",
     "context_ready_for_model = pad_fn(context_idx)\n",
@@ -357,7 +367,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The output is the ligits for the start and end tokens of the answer. It is now time to extract them with `torch.argmax`"
+    "The output corresponds to the digits for the start and end tokens of the answer. It is now time to extract them with `torch.argmax`"
    ]
   },
   {
@@ -375,7 +385,7 @@
     "        print(\"The model proposes: '\", context[answer_idx[0][0]],\"...'\")\n",
     "except IndexError:\n",
     "    print(\"The model was not able to find the answer.\")\n",
-    "print(\"The actual answer was: '\" + bb[2][0][0]+\"'\")"
+    "print(\"The actual answer was: '\" + bb[2][0]+\"'\")"
    ]
   },
   {
@@ -408,7 +418,7 @@
    "outputs": [],
    "source": [
     "DEVICE = torch.device(\"cpu\")\n",
-    "x = next(iter(dl_tr))[0]\n",
+    "x = next(iter(dl_tr2))[0]\n",
     "pipe.model.eval()\n",
     "pipe.model(x.to(DEVICE))\n",
     "\n",
@@ -422,7 +432,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "x = next(iter(dl_tr))[0][0]\n",
+    "x = next(iter(dl_tr2))[0][0]\n",
     "if x.dtype is not torch.int64:\n",
     "    res = me.get_decision_boundary(x, n_epochs=1)\n",
     "    res.shape"
@@ -434,7 +444,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "x, target = next(iter(dl_tr))\n",
+    "x, target = next(iter(dl_tr2))\n",
     "if x.dtype is torch.float:\n",
     "    for gradient in me.get_gradients(x, target=target)[1]:\n",
     "        print(gradient.shape)"
diff --git a/examples/basic_tutorial_dataset_cloud.ipynb b/examples/basic_tutorial_dataset_cloud.ipynb
index a009319a..513d763e 100644
--- a/examples/basic_tutorial_dataset_cloud.ipynb
+++ b/examples/basic_tutorial_dataset_cloud.ipynb
@@ -38,7 +38,7 @@
     "\n",
     "import torch\n",
     "\n",
-    "from gdeep.data import DlBuilderFromDataCloud, DatasetCloud\n",
+    "from gdeep.data.datasets import DlBuilderFromDataCloud, DatasetCloud\n",
     "from gdeep.utility.utils import get_checksum\n",
     "\n"
    ]
@@ -167,7 +167,7 @@
     "\n",
     "# create the dataset from the downloaded dataset\n",
     "train_dataloader, val_dataloader, test_dataloader = \\\n",
-    "    dl_cloud_builder.build_dataloaders(batch_size=10)\n",
+    "    dl_cloud_builder.build(({\"batch_size\":10},))\n",
     "\n",
     "del train_dataloader, val_dataloader, test_dataloader\n"
    ]
diff --git a/examples/basic_tutorial_gridsearch_benchmarking.ipynb b/examples/basic_tutorial_gridsearch_benchmarking.ipynb
index fc409153..02bba6af 100644
--- a/examples/basic_tutorial_gridsearch_benchmarking.ipynb
+++ b/examples/basic_tutorial_gridsearch_benchmarking.ipynb
@@ -43,11 +43,11 @@
     "\n",
     "from torch.optim import SGD, Adam, RMSprop\n",
     "from torch.optim.lr_scheduler import ExponentialLR\n",
-    "from gdeep.pipeline import Pipeline\n",
-    "from gdeep.data import TorchDataLoader\n",
+    "from gdeep.trainer import Trainer\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
     "# today's protagonists\n",
     "from gdeep.search import Benchmark\n",
-    "from gdeep.search import Gridsearch, GiottoSummaryWriter\n",
+    "from gdeep.search import HyperParameterOptimization, GiottoSummaryWriter\n",
     "\n",
     "from gtda.diagrams import BettiCurve\n",
     "\n",
@@ -100,23 +100,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = TorchDataLoader(name=\"CIFAR10\")\n",
+    "from torchvision.transforms import ToTensor\n",
+    "\n",
+    "bd = DatasetBuilder(name=\"CIFAR10\")\n",
     "train_indices = list(range(32*10))\n",
     "\n",
     "print(len(train_indices))\n",
     "\n",
-    "dl_tr, dl_temp = dl.build_dataloaders(batch_size=32, \n",
-    "                                      sampler=SubsetRandomSampler(train_indices))\n",
-    "\n",
-    "print(len(dl_tr))\n",
+    "ds_tr, _, _ = bd.build( download=True)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Preprocessing\n",
     "\n",
-    "test_indices = [32*10 + x for x in list(range(32*2))]\n",
+    "from gdeep.data.preprocessors import ToTensorImage\n",
     "\n",
-    "dl_ts, dl_temp = dl.build_dataloaders(batch_size=25, sampler=SubsetRandomSampler(test_indices))\n",
+    "transformation = ToTensorImage((32,32))\n",
+    "transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation\n",
     "\n",
-    "dl_val = dl_ts\n",
+    "transformed_ds_tr = transformation.attach_transform_to_dataset(ds_tr)\n",
     "\n",
-    "print(len(dl_ts))"
+    "# use only 320 images from cifar10\n",
+    "train_indices = list(range(32*10))\n",
+    "dl_tr, *_ = DataLoaderBuilder((transformed_ds_tr, )).build(\n",
+    "    ({\"batch_size\":32, \"sampler\":SubsetRandomSampler(train_indices)},))"
    ]
   },
   {
@@ -159,7 +172,7 @@
     "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
     "# initialise pipeline class\n",
-    "pipe = Pipeline(model, [dl_tr, dl_val, dl_ts], loss_fn, writer, StratifiedKFold(2, shuffle=True))"
+    "pipe = Trainer(model, [dl_tr, None], loss_fn, writer, StratifiedKFold(2, shuffle=True))"
    ]
   },
   {
@@ -178,14 +191,14 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Gridsearch\n",
+    "# HyperParameter Optimization\n",
     "\n",
-    "We run a gridsearch over different hyperparametrs: \n",
+    "We run a search over different hyperparametrs: \n",
     " - the learning rate `lr`\n",
     " - the batch size `batch_size`\n",
     " - the network parameter `arch`\n",
     " \n",
-    "The scope of the gridsearch is to find the optimum set of hyperparameters."
+    "The scope of the search is to find the optimum set of hyperparameters."
    ]
   },
   {
@@ -196,7 +209,7 @@
    "source": [
     "\n",
     "# initialise gridsearch\n",
-    "search = Gridsearch(pipe, \"accuracy\", 2, best_not_last=True)\n",
+    "search = HyperParameterOptimization(pipe, \"accuracy\", 2, best_not_last=True)\n",
     "\n",
     "# if you want to store pickle files of the models instead of the state_dicts\n",
     "search.store_pickle = True\n",
@@ -283,24 +296,31 @@
    "outputs": [],
    "source": [
     "dataloaders_dicts = []\n",
-    "dl = TorchDataLoader(name=\"CIFAR10\")\n",
+    "bd = DatasetBuilder(name=\"CIFAR10\")\n",
     "\n",
-    "train_indices = list(range(64*5))\n",
+    "ds_tr, *_ = bd.build()\n",
+    "transformation = ToTensorImage((32,32))\n",
     "\n",
-    "dl_tr, dl_temp = dl.build_dataloaders(batch_size=32, sampler=SubsetRandomSampler(train_indices))\n",
+    "transformed_ds_tr = transformation.attach_transform_to_dataset(ds_tr)\n",
     "\n",
-    "test_indices = [64*5 + x for x in list(range(64))]\n",
     "\n",
-    "dl_ts, dl_temp = dl.build_dataloaders(batch_size=32, sampler=SubsetRandomSampler(test_indices))\n",
+    "test_indices = [64*5 + x for x in range(32*3)]\n",
+    "train_indices = [x for x in range(32*2)]\n",
+    "\n",
+    "dl = DataLoaderBuilder((transformed_ds_tr, transformed_ds_tr))\n",
+    "dl_tr, dl_val, _ = dl.build(({\"batch_size\":32,\"sampler\":SubsetRandomSampler(train_indices)},{\"batch_size\":32,\"sampler\":SubsetRandomSampler(test_indices)}))\n",
+    "\n",
     "\n",
     "temp_dict = {}\n",
     "temp_dict[\"name\"] = \"CIFAR10_1000\"\n",
-    "temp_dict[\"dataloaders\"] = (dl_tr, dl_ts)\n",
+    "temp_dict[\"dataloaders\"] = (dl_tr, dl_val, _)\n",
     "\n",
     "dataloaders_dicts.append(temp_dict)\n",
     "\n",
-    "dl = TorchDataLoader(name=\"DoubleTori\")\n",
-    "dl_tr, dl_ts = dl.build_dataloaders(batch_size=48)\n",
+    "db = DatasetBuilder(name=\"DoubleTori\")\n",
+    "ds_tr, ds_val, _ = db.build()\n",
+    "\n",
+    "dl_tr, dl_ts, _ = DataLoaderBuilder((ds_tr, ds_val)).build(({\"batch_size\":48},{\"batch_size\":32}))\n",
     "\n",
     "temp_dict = {}\n",
     "temp_dict[\"name\"] = \"double_tori\"\n",
@@ -376,7 +396,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Benchmarking + Gridsearch + CV\n",
+    "# Benchmarking + HyperParameter Optimization + CV\n",
     "\n",
     "In this last section we consider the possibility of running a gridsearch within each pair *(model, dataset)*.\n",
     "\n",
@@ -398,7 +418,7 @@
     "bench = Benchmark(models_dicts, dataloaders_dicts, loss_fn, writer, KFold(3))\n",
     "\n",
     "# initialise gridsearch with benchmark instance\n",
-    "search2 = Gridsearch(bench, \"loss\", 2)\n",
+    "search2 = HyperParameterOptimization(bench, \"loss\", 2)\n",
     "\n",
     "# yperparameters\n",
     "optimizers_params = {\"lr\": [0.001, 0.01, None, True]}  # to have the log sampler\n",
@@ -427,7 +447,7 @@
    "source": [
     "## Custom pruner and sampler\n",
     "\n",
-    "It is possible to pass to the gridsearch class a customer `optuna.Pruners` and `optuna.Samplers`"
+    "It is possible to pass to the HyperParameterOptimization class a customer `optuna.Pruners` and `optuna.Samplers`"
    ]
   },
   {
@@ -440,7 +460,7 @@
     "from optuna.samplers import TPESampler\n",
     "\n",
     "# initialise gridsearch\n",
-    "gs = Gridsearch(pipe, \"accuracy\", 8, best_not_last=False, pruner=MedianPruner(n_startup_trials=2,\n",
+    "gs = HyperParameterOptimization(pipe, \"accuracy\", 8, best_not_last=False, pruner=MedianPruner(n_startup_trials=2,\n",
     "                                                                              n_warmup_steps=0,\n",
     "                                                                              interval_steps=1,\n",
     "                                                                              n_min_trials=1), sampler=TPESampler())\n",
@@ -453,6 +473,13 @@
     "# starting the gridsearch\n",
     "gs.start((SGD, Adam), 3, False, optimizers_params, dataloaders_params, models_hyperparams, n_accumulated_grads=2)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/examples/basic_tutorial_image.ipynb b/examples/basic_tutorial_image.ipynb
index b3d2781c..2b348c94 100644
--- a/examples/basic_tutorial_image.ipynb
+++ b/examples/basic_tutorial_image.ipynb
@@ -37,7 +37,7 @@
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
     "\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data import TorchDataLoader\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
     "\n",
     "\n",
     "from gtda.diagrams import BettiCurve\n",
@@ -84,11 +84,38 @@
    "outputs": [],
    "source": [
     "from torch.utils.data.sampler import SubsetRandomSampler\n",
-    "dl = TorchDataLoader(name=\"CIFAR10\")\n",
+    "from torchvision.transforms import ToTensor  # for preprocessing\n",
+    "\n",
+    "db = DatasetBuilder(name=\"CIFAR10\")\n",
+    "ds_tr, ds_val, ds_ts = db.build()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Preprocessing\n",
+    "\n",
+    "from gdeep.data.preprocessors import ToTensorImage\n",
+    "\n",
+    "transformation = ToTensorImage((32,32))\n",
+    "transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation\n",
+    "\n",
+    "transformed_ds_tr = transformation.attach_transform_to_dataset(ds_tr)\n",
+    "transformed_ds_val = transformation.attach_transform_to_dataset(ds_val)\n",
+    "transformed_ds_ts = transformation.attach_transform_to_dataset(ds_ts)\n",
     "\n",
     "# use only 320 images from cifar10\n",
     "train_indices = list(range(32*10))\n",
-    "dl_tr, dl_ts = dl.build_dataloaders(batch_size=32, sampler=SubsetRandomSampler(train_indices))"
+    "val_indices = list(range(32*5))\n",
+    "test_indices = list(range(32*5))\n",
+    "dl_tr, dl_val, dl_ts = DataLoaderBuilder((transformed_ds_tr, transformed_ds_val, transformed_ds_ts)).build(\n",
+    "    ({\"batch_size\":32, \"sampler\":SubsetRandomSampler(train_indices)},\n",
+    "     {\"batch_size\":32, \"sampler\":SubsetRandomSampler(val_indices)},\n",
+    "     {\"batch_size\":32, \"sampler\":SubsetRandomSampler(test_indices)}\n",
+    "    ))"
    ]
   },
   {
@@ -105,7 +132,7 @@
    "outputs": [],
    "source": [
     "import torchvision.models as models\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.trainer import Trainer\n",
     "\n",
     "# wrap a sequential model in a torch nn.Module\n",
     "class model3(nn.Module):\n",
@@ -129,10 +156,11 @@
     "# print(model)\n",
     "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
-    "pipe = Pipeline(model, (dl_tr, dl_ts), loss_fn, writer)\n",
+    "pipe = Trainer(model, (dl_tr, dl_ts), loss_fn, writer)\n",
     "\n",
     "# train the model\n",
-    "pipe.train(SGD, 3, False, {\"lr\":0.01}, {\"batch_size\":32, \"sampler\":SubsetRandomSampler(train_indices)})\n",
+    "pipe.train(SGD, 3, False, {\"lr\":0.01}, {\"batch_size\":32, \n",
+    "                                        \"sampler\":SubsetRandomSampler(train_indices)})\n",
     "\n"
    ]
   },
diff --git a/examples/basic_tutorial_regression.ipynb b/examples/basic_tutorial_regression.ipynb
index af715db4..c3ea4be5 100644
--- a/examples/basic_tutorial_regression.ipynb
+++ b/examples/basic_tutorial_regression.ipynb
@@ -42,10 +42,10 @@
     "from torch.optim import SGD, Adam\n",
     "\n",
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.trainer import Trainer\n",
     "\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data import TorchDataLoader, DataLoaderFromArray\n",
+    "from gdeep.data.datasets import DatasetBuilder, FromArray, DataLoaderBuilder\n",
     "\n",
     "from gtda.diagrams import BettiCurve\n",
     "\n",
@@ -90,14 +90,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "X_train = np.random.rand(100,3)\n",
-    "y_train = 0.3*np.array(list(map(sum,X_train)))  # a hyperplane\n",
+    "X_train = np.array(np.random.rand(100,3), dtype=np.float32)\n",
+    "y_train = np.array(0.3*np.array(list(map(sum,X_train))).reshape(-1,1), dtype=np.float32)  # a hyperplane\n",
     "\n",
-    "X_val = np.random.rand(50,3)\n",
-    "y_val = 0.3*np.array(list(map(sum,X_train)))\n",
+    "X_val = np.array(np.random.rand(50,3), dtype=np.float32)\n",
+    "y_val = np.array(0.3*np.array(list(map(sum,X_train))).reshape(-1,1), dtype=np.float32)\n",
     "\n",
-    "dl = DataLoaderFromArray(X_train, y_train, X_val, y_val)\n",
-    "dl_tr, dl_val, _ = dl.build_dataloaders(batch_size=32)\n"
+    "dl_builder = DataLoaderBuilder((FromArray(X_train, y_train), FromArray( X_val, y_val)))\n",
+    "dl_tr, dl_val, _ = dl_builder.build(({\"batch_size\":32},\n",
+    "                                    {\"batch_size\":16}))\n"
    ]
   },
   {
@@ -118,6 +119,7 @@
     "    def __init__(self):\n",
     "        super(model1, self).__init__()\n",
     "        self.seqmodel = FFNet(arch=[3, 5, 1])\n",
+    "\n",
     "    def forward(self, x):\n",
     "        return self.seqmodel(x)\n",
     "\n",
@@ -133,7 +135,7 @@
     "\n",
     "loss_fn = nn.MSELoss()\n",
     "\n",
-    "pipe = Pipeline(model, (dl_tr, dl_val), loss_fn, writer)\n",
+    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer)\n",
     "\n",
     "# train the model with learning rate scheduler\n",
     "pipe.train(Adam, 3, False, lr_scheduler=ExponentialLR, scheduler_params={\"gamma\": 0.9}, \n",
diff --git a/examples/basic_tutorial_tabular.ipynb b/examples/basic_tutorial_tabular.ipynb
index 35fd76a9..cd432796 100644
--- a/examples/basic_tutorial_tabular.ipynb
+++ b/examples/basic_tutorial_tabular.ipynb
@@ -44,10 +44,10 @@
     "from torch.utils.data import SubsetRandomSampler\n",
     "\n",
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.trainer import Trainer\n",
     "\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data import TorchDataLoader\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
     "\n",
     "from gtda.diagrams import BettiCurve\n",
     "\n",
@@ -94,19 +94,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = TorchDataLoader(name=\"DoubleTori\")\n",
+    "bd = DatasetBuilder(name=\"DoubleTori\")\n",
+    "ds_tr, ds_val, _ = bd.build()\n",
     "#train_indices = list(range(160))\n",
-    "dl_tr, dl_val = dl.build_dataloaders(batch_size=23) #, sampler=SubsetRandomSampler(train_indices))\n",
-    "dl_ts = None"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "isinstance(next(iter(dl_tr))[0],torch.Tensor)"
+    "dl_builder = DataLoaderBuilder((ds_tr, ds_val))\n",
+    "dl_tr, dl_val, dl_ts = dl_builder.build(({\"batch_size\":23},{\"batch_size\":23}))\n"
    ]
   },
   {
@@ -144,7 +136,7 @@
     "# initlaise the loss function\n",
     "loss_fn = nn.CrossEntropyLoss()\n",
     "# initialise the pipelien class\n",
-    "pipe = Pipeline(model, (dl_tr, dl_val, _), loss_fn, writer, StratifiedKFold(3, shuffle=True))\n",
+    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer, StratifiedKFold(3, shuffle=True))\n",
     "\n",
     "# initialise the SAM optimiser\n",
     "optim = SAMOptimizer(SGD)  # this is a class, not an instance!\n",
diff --git a/examples/basic_tutorial_textClass.ipynb b/examples/basic_tutorial_textClass.ipynb
index 85622f68..27a3f107 100644
--- a/examples/basic_tutorial_textClass.ipynb
+++ b/examples/basic_tutorial_textClass.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Basic tutorial: text data\n",
+    "# Basic tutorial: text classification\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
     "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
@@ -37,7 +37,7 @@
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
     "\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data import TorchDataLoader\n",
+    "from gdeep.data.datasets import DatasetBuilder\n",
     "\n",
     "from gtda.diagrams import BettiCurve\n",
     "\n",
@@ -86,8 +86,8 @@
     "\n",
     "# many time we get an IterableDataset which is good for memory consumption, but cannot be sampled!\n",
     "# we can load the whole dataset in memory and sample it\n",
-    "dl = TorchDataLoader(name=\"AG_NEWS\", convert_to_map_dataset=True)\n",
-    "dl_tr, dl_ts = dl.build_dataloaders()\n",
+    "bd = DatasetBuilder(name=\"AG_NEWS\", convert_to_map_dataset=True)\n",
+    "ds_tr_str, ds_val_str, ds_ts_str = bd.build()\n",
     "\n"
    ]
   },
@@ -104,7 +104,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(\"Before the preprocessing: \\n\", dl_tr.dataset[0])"
+    "print(\"Before the preprocessing: \\n\", ds_tr_str[0])"
    ]
   },
   {
@@ -113,31 +113,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.data import PreprocessingPipeline, Normalisation, TextDataset, PreprocessTextData, PreprocessTextLabel\n",
-    "from torch.utils.data import DataLoader\n",
+    "from gdeep.data import PreprocessingPipeline\n",
+    "from gdeep.data import TransformingDataset\n",
+    "from gdeep.data.preprocessors import Normalization, TokenizerTextClassification\n",
+    "from gdeep.data.datasets import DataLoaderBuilder\n",
     "\n",
-    "ptd = PreprocessTextData()\n",
+    "ptd = TokenizerTextClassification()\n",
     "\n",
     "# in case you need to combine multiple preprocessing:\n",
-    "# ppp = PreprocessingPipeline(((PreprocessTextData(), TextDataset), (Normalisation(), lambda x, __:x)))\n",
+    "#ppp = PreprocessingPipeline(((PreprocessTextData(), IdentityTransform(), TextDataset), \n",
+    "#                             (Normalisation(), IdentityTransform(), BasicDataset)))\n",
     "\n",
     "\n",
-    "textds = TextDataset(dl_tr.dataset,\n",
-    "                     ptd, \n",
-    "                     PreprocessTextLabel())\n",
+    "ptd.fit_to_dataset(ds_tr_str)\n",
+    "transformed_textds = ptd.attach_transform_to_dataset(ds_tr_str)\n",
     "\n",
-    "textts = TextDataset(dl_ts.dataset,\n",
-    "                     ptd, \n",
-    "                     PreprocessTextLabel())\n",
+    "transformed_textts = ptd.attach_transform_to_dataset(ds_val_str)\n",
     "\n",
-    "print(\"After the preprocessing: \\n\", textds[0])\n",
+    "print(\"After the preprocessing: \\n\", transformed_textds[0])\n",
     "\n",
     "# the only part of the training/test set we are interested in\n",
     "train_indices = list(range(64*20))\n",
     "test_indices = list(range(64*10))\n",
     "\n",
-    "dl_tr2 = DataLoader(textds, batch_size=16, sampler=SubsetRandomSampler(train_indices))\n",
-    "dl_ts2 = DataLoader(textts, batch_size=16, sampler=SubsetRandomSampler(test_indices))\n",
+    "dl_tr2, dl_ts2, _ = DataLoaderBuilder((transformed_textds, \n",
+    "                                      transformed_textts)).build(({\"batch_size\":16,\n",
+    "                                                                   \"sampler\":SubsetRandomSampler(train_indices)},\n",
+    "                                                                  {\"batch_size\":16,\n",
+    "                                                                   \"sampler\":SubsetRandomSampler(test_indices)}\n",
+    "                                                                 ))\n",
     "\n",
     "\n"
    ]
@@ -155,7 +159,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.trainer import Trainer\n",
     "from torch.optim import Adam, SparseAdam, SGD\n",
     "from captum.attr import Occlusion, IntegratedGradients, visualization\n",
     "from torchvision import transforms\n",
@@ -188,8 +192,6 @@
    "outputs": [],
    "source": [
     "vocab_size = len(ptd.vocabulary)\n",
-    "print(\"len\", vocab_size)\n",
-    "vocab_size = 225981 # to be discussed, since len(prec.vocabulary) is returning 95810 whereas the index of '.' is 225971\n",
     "emsize = 64\n",
     "# print(vocab_size, emsize)\n",
     "model = TextClassificationModel(vocab_size, emsize, 4)\n",
@@ -205,7 +207,7 @@
     "\n",
     "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
-    "pipe = Pipeline(model, (dl_tr2, dl_ts2), loss_fn, writer)\n",
+    "pipe = Trainer(model, (dl_tr2, dl_ts2), loss_fn, writer)\n",
     "\n",
     "# train the model\n",
     "pipe.train(SGD, 7, False, {\"lr\":0.01}, {\"batch_size\":20})\n",
@@ -307,7 +309,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "\n",
     "vs.plot_data_model()\n",
     "vs.plot_activations(x)\n",
diff --git a/examples/basic_tutorial_translation.ipynb b/examples/basic_tutorial_translation.ipynb
index 85a1f57c..9d10eca7 100644
--- a/examples/basic_tutorial_translation.ipynb
+++ b/examples/basic_tutorial_translation.ipynb
@@ -7,16 +7,41 @@
     "# Basic tutorial: translation\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
+    "This short tutorial provides you with the basic functionalities of *giotto-deep* API.\n",
     "\n",
-    "The example described in this tutorial is the one of translation.\n",
+    "The example described in this tutorial is the one of translation following a sequence-to-sequence transformer model: all teh details can be found [here](https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf). \n",
+    "\n",
+    "## Probabilistic modelling\n",
+    "\n",
+    "The basic idea would be to predict the target sentence $y$ given source sentence $x$, more precisely you ant to predict $P(y|x)$. But this is very hard to do directly, since one would need a lot of data to have a reasonable estimate of this conditional probability.\n",
+    "Hence, one would ratehr factor the probability using the chain rule, then predict each factor:\n",
+    "\n",
+    "$$\n",
+    "P(y_1,...,y_N|x) = \\prod_i^N P(y_i|y_{<i},x)\n",
+    "$$\n",
+    "\n",
+    "In the NLP language, this corresponds to predicting the next token of the output sequence given the already predicted tokens and the input sequence.\n",
+    "\n",
+    "## Scope\n",
+    "\n",
+    "You will build your own transformer model, train it, and use it to translate German sentences to English!\n",
+    "\n",
+    "Here is an example of what a translation task is about:\n",
+    "\n",
+    " - German sentence: 'Ich mag Pizza.'\n",
+    " - English translation: 'I like pizza.'\n",
+    "\n",
+    "## The plan for this tutorial\n",
     "\n",
     "The main steps of the tutorial are the following:\n",
     " 1. creation of a dataset\n",
     " 2. creation of a model\n",
-    " 3. define metrics and losses\n",
-    " 4. train the model\n",
-    " 5. extract some features of the network"
+    " 3. definition of the metrics and losses\n",
+    " 4. trainining of the model\n",
+    " 5. using the model to translate some sentences\n",
+    " 6. (extra) extract some features of the network for interpretability analysis\n",
+    " \n",
+    "Let's start with importing the requried libraries!"
    ]
   },
   {
@@ -29,22 +54,31 @@
     "%autoreload 2\n",
     "%matplotlib inline\n",
     "\n",
-    "import numpy as np\n",
+    "import copy\n",
+    "import sys\n",
     "\n",
+    "import numpy as np\n",
     "import torch\n",
     "from torch import nn\n",
-    "\n",
-    "from gdeep.models import FFNet\n",
-    "\n",
-    "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
-    "\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data import TorchDataLoader\n",
-    "from gdeep.pipeline import Pipeline\n",
-    "\n",
     "from gtda.diagrams import BettiCurve\n",
+    "from gtda.plotting import plot_betti_surfaces\n",
+    "from torchtext.data.utils import get_tokenizer\n",
+    "from torch.utils.data.sampler import SubsetRandomSampler\n",
+    "from torch.nn import Transformer\n",
+    "from torch.optim import Adam, SparseAdam, SGD\n",
     "\n",
-    "from gtda.plotting import plot_betti_surfaces"
+    "\n",
+    "# our special guests!\n",
+    "from gdeep.models import FFNet\n",
+    "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
+    "from gdeep.data.datasets import DatasetBuilder\n",
+    "from gdeep.trainer import Trainer\n",
+    "from gdeep.data import TransformingDataset\n",
+    "from gdeep.data.preprocessors import TokenizerTranslation\n",
+    "from gdeep.data.datasets import DataLoaderBuilder\n",
+    "from gdeep.models import ModelExtractor\n",
+    "from gdeep.visualisation import Visualiser"
    ]
   },
   {
@@ -53,14 +87,16 @@
    "source": [
     "# Initialize the tensorboard writer\n",
     "\n",
-    "In order to analyse the reuslts of your models, you need to start tensorboard.\n",
-    "On the terminal, move inside the `/example` folder. There run the following command:\n",
+    "In order to analyse the results of your models, you need to start tensorboard. All data about the model, the training, the hyperparameters... will be stored there.\n",
+    "\n",
+    "## How to start tensorboard\n",
+    "On the terminal, move inside the `/examples` folder. There run the following command:\n",
     "\n",
     "```\n",
     "tensorboard --logdir=runs\n",
     "```\n",
     "\n",
-    "Then go [here](http://localhost:6006/) after the training to see all the visualisation results."
+    "Then go [here](http://localhost:6006/) after the training step to visualise all the results."
    ]
   },
   {
@@ -76,7 +112,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Create your dataset"
+    "# Create your dataset\n",
+    "\n",
+    "We propose to use the [Multi30k](https://github.com/multi30k/dataset) dataset: in the next cell you will see *how easy it is to import datasets in giotto-deep*."
    ]
   },
   {
@@ -85,20 +123,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from torch.utils.data.sampler import SubsetRandomSampler\n",
+    "# prepare the builder class\n",
+    "bd = DatasetBuilder(name=\"Multi30k\", convert_to_map_dataset=True)\n",
     "\n",
-    "# the only part of the training set we are interested in\n",
-    "train_indices = list(range(32*10))\n",
-    "\n",
-    "dl = TorchDataLoader(name=\"Multi30k\", convert_to_map_dataset=True)\n",
-    "dl_tr_str, dl_ts_str = dl.build_dataloaders(sampler=SubsetRandomSampler(train_indices))\n"
+    "# build the dataset\n",
+    "ds_tr_str, ds_val_str, ds_ts_str = bd.build()\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The dataset contains a list of pairs of sentences: the English sentecen and its German translation."
+    "The dataset contains a list of pairs of sentences: the German sentence and its English translation."
    ]
   },
   {
@@ -107,7 +143,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(\"Before preprocessing: \\n\", dl_tr_str.dataset[0])"
+    "print(\"Before preprocessing: \\n\", ds_tr_str[0])"
    ]
   },
   {
@@ -117,11 +153,11 @@
     "## Required preprocessing\n",
     "\n",
     "Neural networks cannot direcly deal with strings. We have first to preprocess the dataset in three main ways:\n",
-    " 1. Tokenise the string into its words\n",
-    " 2. Build a vocabulary out of these words\n",
-    " 3. Embed each word into a vector, so that each sentence becomes a list of vectors\n",
+    " 1. Tokenise each string into its words (and maybe adjust each word to remove plurals, interjections, capital letters...)\n",
+    " 2. Build a vocabulary out of these tokens (each modified word of point a. is called a token)\n",
+    " 3. Embed each token into a vector, so that each sentence becomes a list of vectors\n",
     "\n",
-    "The first two steps are performed by the `PreprocessTextTranslation`. The embedding will be added directly to the model."
+    "The **first two steps** are performed by the `TokenizerTranslation` class. The embedding will be added directly to the model (see below)."
    ]
   },
   {
@@ -130,24 +166,34 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.data import TextDatasetTranslation, PreprocessTextTranslation\n",
-    "from torch.utils.data import DataLoader\n",
     "\n",
-    "prec = PreprocessTextTranslation()\n",
-    "textds = TextDatasetTranslation(dl_tr_str.dataset,\n",
-    "                     prec, \n",
-    "                     None)\n",
-    "textts = TextDatasetTranslation(dl_ts_str.dataset,\n",
-    "                     prec, \n",
-    "                     None)\n",
+    "# install german tokenizer\n",
+    "!{sys.executable} -m spacy download de_core_news_sm\n",
+    "\n",
+    "# get the german tokenizer\n",
+    "de_tokenizer = get_tokenizer('spacy', language='de_core_news_sm')\n",
+    "\n",
+    "# initialise the giotto-deep tokenizer\n",
+    "tokenizer = TokenizerTranslation(tokenizer=de_tokenizer)\n",
+    "\n",
+    "# fit the tokenizer to the dataset (note that te vocabularies will be automatically built in this case\n",
+    "tokenizer.fit_to_dataset(ds_tr_str)\n",
     "\n",
-    "print(\"After the preprocessing: \\n\", textds[0])\n",
+    "# prprocess the dataset\n",
+    "transformed_textds = tokenizer.attach_transform_to_dataset(ds_tr_str)\n",
+    "transformed_textts = tokenizer.attach_transform_to_dataset(ds_val_str)\n",
     "\n",
-    "train_indices = list(range(64*20))\n",
-    "test_indices = list(range(64*10))\n",
+    "print(\"After the preprocessing: \\n\", transformed_textds[0])\n",
     "\n",
-    "dl_tr = DataLoader(textds, batch_size=16, sampler=SubsetRandomSampler(train_indices))\n",
-    "dl_ts = DataLoader(textts, batch_size=16, sampler=SubsetRandomSampler(test_indices))\n",
+    "# subsample the training and test datasets\n",
+    "train_indices = list(range(64*2))\n",
+    "test_indices = list(range(64*1))\n",
+    "\n",
+    "dl_tr, dl_val, _ = DataLoaderBuilder((transformed_textds, \n",
+    "                                   transformed_textts)).build(({\"batch_size\":16, \n",
+    "                                                                          \"sampler\":SubsetRandomSampler(train_indices)},{\"batch_size\":16, \n",
+    "                                                                          \"sampler\":SubsetRandomSampler(test_indices)}\n",
+    "                                                                          ))\n",
     "\n",
     "\n"
    ]
@@ -156,7 +202,20 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Define and train your model"
+    "As you can see, now the data is not in string format anymore: rather, the sentences have each been transformed to a `torch.Tensor` of type `long`. Each of these numbers represents the index in the vocabulary of the associated token."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define your model\n",
+    "\n",
+    "The model we play with is a simple transformer model with two embedding layers (for the German and English sentence) followed by a single transformer layer. \n",
+    "\n",
+    "The input of this model is assumed to be a tensors. The tensor is of the form `shape(transformed_textds[0][0])`. \n",
+    "\n",
+    "Of course, as in all neural networks, data are inputted to the model  in batches: hence, the first dimension will be the batch size, then a dimension to choose the source or target sentence, and the last dimension will be about the actual tokenisation."
    ]
   },
   {
@@ -165,9 +224,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from torch.nn import Transformer\n",
-    "from torch.optim import Adam, SparseAdam, SGD\n",
-    "import copy\n",
+    "\n",
     "\n",
     "# my simple transformer model\n",
     "class TranslationTransformer(nn.Module):\n",
@@ -185,62 +242,56 @@
     "        self.generator = nn.Linear(embed_dim, tgt_vocab_size)\n",
     "        \n",
     "    def forward(self, X):\n",
-    "        #print(X.shape)\n",
     "        src = X[:,0,:]\n",
     "        tgt = X[:,1,:]\n",
-    "        #print(src.shape, tgt.shape)\n",
     "        src_emb = self.embedding_src(src)\n",
     "        tgt_emb = self.embedding_tgt(tgt)\n",
-    "        #print(src_emb.shape, tgt_emb.shape)\n",
     "        self.outs = self.transformer(src_emb, tgt_emb)\n",
-    "        #print(outs.shape)\n",
     "        logits = self.generator(self.outs)\n",
-    "        #print(logits.shape)\n",
-    "        #out = torch.topk(logits, k=1, dim=2).indices.reshape(-1,44)\n",
-    "        #print(out, out.shape)\n",
     "        return logits\n",
     "    \n",
-    "    def __deepcopy__(self, memo):\n",
-    "        \"\"\"this is needed to make sure that the \n",
-    "        non-leaf nodes do not\n",
-    "        interfere with copy.deepcopy()\n",
-    "        \"\"\"\n",
-    "        cls = self.__class__\n",
-    "        result = cls.__new__(cls)\n",
-    "        memo[id(self)] = result\n",
-    "        for k, v in self.__dict__.items():\n",
-    "            setattr(result, k, copy.deepcopy(v, memo))\n",
-    "        return result\n",
-    "    \n",
     "    def encode(self, src, src_mask):\n",
-    "        \"\"\"this method is used only at the inference step\"\"\"\n",
+    "        \"\"\"this method is used only at the inference step.\n",
+    "        This method runs the data through the encoder of the\n",
+    "        transformer\"\"\"\n",
     "        return self.transformer.encoder(\n",
     "                            self.embedding_src(src), src_mask)\n",
     "\n",
     "    def decode(self, tgt, memory, tgt_mask):\n",
-    "        \"\"\"this method is used only at the inference step\"\"\"\n",
+    "        \"\"\"this method is used only at the inference step\n",
+    "        This method runs the data through the decoder of the\n",
+    "        transformer\"\"\"\n",
     "        return self.transformer.decoder(\n",
     "                          self.embedding_tgt(tgt), memory,\n",
     "                          tgt_mask)\n",
     "    \n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We initialise the model as follows:\n",
+    " - we need to set the maximum vocabulary size to fix the `Embedding` architectures\n",
+    " - we need to set the embedding dimension\n",
+    " - initialise the model class with the needed parameters."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "vocab_size = 30000 # to be discussed\n",
     "\n",
-    "src_vocab_size = vocab_size # len(prec.vocabulary)\n",
-    "tgt_vocab_size = vocab_size # len(prec.vocabulary_target)\n",
+    "src_vocab_size = len(tokenizer.vocabulary)\n",
+    "tgt_vocab_size = len(tokenizer.vocabulary_target)\n",
     "emb_dim = 64\n",
     "\n",
     "model = TranslationTransformer(src_vocab_size, tgt_vocab_size, emb_dim)\n",
-    "X = next(iter(dl_tr));\n",
+    "X = next(iter(dl_tr));  # a datum\n",
     "#assert model(X[0]).argmax(2).shape == X[1].shape\n",
-    "print(model)"
+    "print(\"This is our model: \\n\", model)"
    ]
   },
   {
@@ -249,7 +300,7 @@
    "source": [
     "## Define the loss function\n",
     "\n",
-    "This loss function is a adapted version of the Cross Entropy for the trnasformer architecture."
+    "This loss function is an adapted version of the Cross Entropy Loss for the transformer architecture we just defined."
    ]
   },
   {
@@ -264,6 +315,17 @@
     "    return cel(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "And here we are: all is set upo and we are ready to train the model.\n",
+    "\n",
+    "# Trainig the model\n",
+    "\n",
+    "In giotto deep all is done via a `Trainer`, meaning a class taht takes care of the training and validation steps, storing the intermediate results to tensorboard, using (or not) cross validation, ... basically everythin you need."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -271,17 +333,23 @@
    "outputs": [],
    "source": [
     "# prepare a pipeline class with the model, dataloaders loss_fn and tensorboard writer\n",
-    "pipe = Pipeline(model, (dl_tr, dl_ts), loss_fn, writer)\n",
+    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer)\n",
     "\n",
     "# train the model\n",
-    "pipe.train(SGD, 3, False, {\"lr\":0.01}, {\"batch_size\":16})"
+    "pipe.train(optimizer=SGD, \n",
+    "           n_epochs=2, \n",
+    "           cross_validation=False, \n",
+    "           optimizers_param={\"lr\":0.01}, \n",
+    "           dataloaders_param={\"batch_size\":16})"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Translation!"
+    "## Translation!\n",
+    "\n",
+    "So the model has been trained and now it is time to translate a sentece. Let's take thee following example:"
    ]
   },
   {
@@ -290,7 +358,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "de, en = next(iter(dl_tr_str))\n",
+    "de, en = next(iter(ds_tr_str))\n",
     "print(de, \"\\n\", en)"
    ]
   },
@@ -308,10 +376,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "voc = prec.vocabulary_target\n",
-    "sent = str.lower(de[0]).split()\n",
-    "de_sentence = list(map(voc.__getitem__,sent))\n",
-    "de_sentence"
+    "voc = tokenizer.vocabulary\n",
+    "#sent = str.lower(de).split()\n",
+    "#de_sentence = list(map(voc.__getitem__,sent))\n",
+    "de_sentence, en_sentence = tokenizer((de, en))[0]\n",
+    "de_sentence\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We need to build a couple of auxiliary functiosn to help us translate. The logic is to use the transformer encoder and decoder layers directly. Let's have a quick reminder o how things work:\n",
+    "\n",
+    "<img src=\"./images/translation_transformer.png\" alt=\"drawing\" width=\"400\" class=\"center\"/>\n",
+    "\n",
+    "The above drawing is a good representation of our architecture. During training, the English and German sentence are inputted from \"below\". Then, once the model is used in inference, there is an important change ongoing. Basically, the full input sentence (the German one in our case) is given as input to the encoder: the output (the probablities) are then greedily transformed into a signle token: that token is then used in the \"Outputs\" below: together with the Input sentence, the next word is predicted (at the top) and used again as input in the \"Outputs\" below. Hence, word after word, you get the translation of the original sentence."
    ]
   },
   {
@@ -320,38 +400,75 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "pad_item = 0\n",
+    "max_len = len(en_sentence)\n",
     "\n",
-    "def greedy_decode(model, src, src_mask, max_len, start_symbol):\n",
+    "\n",
+    "def greedy_decode(model, src, src_mask, max_len):\n",
     "    \"\"\"function to generate output sequence using greedy algorithm\"\"\"\n",
+    "\n",
     "    memory = model.encode(src, src_mask)\n",
-    "    out = model.decode(src, memory, None)\n",
-    "    prob = model.generator(out)\n",
-    "    greedy_out = torch.max(prob, dim=2).indices\n",
-    "    return greedy_out\n",
+    "    #ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long)\n",
+    "    ys = torch.ones(*src.shape).fill_(pad_item).type(torch.long)\n",
+    "    for i in range(max_len-1):\n",
+    "        #print(\"before:\", ys)\n",
+    "        out = model.decode(ys, memory, None)\n",
+    "        out = out.transpose(0, 1)\n",
+    "        prob = model.generator(out[:, -1])\n",
+    "        next_word = torch.argmax(prob, dim=1)\n",
+    "        next_word = next_word[i+1].item()  # the shift\n",
+    "        ys[0,i] = next_word\n",
+    "        #print(\"after:\", ys)\n",
+    "    return ys\n",
     "    \n",
     "\n",
     "def translate(model: torch.nn.Module, dl_ts_item):\n",
     "    \"\"\"actual function to translate input sentence into target language\"\"\"\n",
     "    model.eval()\n",
-    "    voc = prec.vocabulary_target\n",
+    "    voc = tokenizer.vocabulary_target\n",
     "    src = dl_ts_item\n",
     "    num_tokens = src.shape[1]\n",
     "    src_mask = None\n",
-    "    tgt_tokens_raw = greedy_decode(\n",
-    "        model,  src, src_mask, max_len=num_tokens + 5, start_symbol=111).flatten()\n",
-    "    tgt_tokens = [token for token in tgt_tokens_raw if token < len(list(voc.vocab))]\n",
-    "    return \" \".join(list(map(list(voc.vocab).__getitem__, tgt_tokens)))\n",
+    "    tgt_tokens_tensor = greedy_decode(\n",
+    "        model,  src, src_mask, max_len).flatten()\n",
+    "    tgt_tokens = tgt_tokens_tensor.tolist()\n",
+    "    #print(tgt_tokens)\n",
+    "    return \" \".join(voc.lookup_tokens(tgt_tokens))\n",
     "    \n",
     "# translation!\n",
-    "print(\"German sentence: \", de[0])\n",
-    "print(\"English translation: \", translate(pipe.model, torch.tensor([de_sentence])))"
+    "print(\"German sentence: \", de)\n",
+    "print(\"English translation: \", translate(model, de_sentence.reshape(1,-1)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# add data saliency map on translation"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Extract inner data from your models"
+    "If you have just run this notebooks, you would most likely have gotten some funny nonsensical answer: consider that you trained a very simple model on a very small subset of data for only two epochs.\n",
+    "\n",
+    "## Challenge\n",
+    "\n",
+    "Starting from this simple notebook, do you think you can enlarge the dataset, the model, the training epochs and get a decent translator? Good luck!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Huggingface transformers\n",
+    "\n",
+    "If you are actually interested in a pretrained transformers that works directly with a few lines of code, `giotto-deep` supports Hugginface transformers.\n",
+    "\n",
+    "The next section explains how to run one such transformer."
    ]
   },
   {
@@ -360,8 +477,37 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.models import ModelExtractor\n",
+    "# from transformers import ...\n",
     "\n",
+    "# take T5small\n",
+    "\n",
+    "# finetune on the Multi30k\n",
+    "\n",
+    "# go to TB to see the results + screenshot"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Appendix: Extract inner data from your models\n",
+    "\n",
+    "With `gdeep` is is pretty straight forwasrd to extract data about the inner working of the models.\n",
+    "\n",
+    "For example you can use the `ModelExtract` and get:\n",
+    " - the parameters of each layer\n",
+    " - the activation functions (given an input)\n",
+    " - the gradients, given a target"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "# initialise the ModelExtractor\n",
     "me = ModelExtractor(pipe.model, loss_fn)\n",
     "\n",
     "lista = me.get_layers_param()\n",
@@ -413,7 +559,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Visualise activations and other topological aspects of your model"
+    "You can also visualise the model graph on tensorboard interactively!"
    ]
   },
   {
@@ -422,13 +568,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.visualisation import Visualiser\n",
     "\n",
+    "# initialise the Visualiser\n",
     "vs = Visualiser(pipe)\n",
     "\n",
     "vs.plot_data_model()\n",
-    "#vs.plot_activations(x)\n",
-    "#vs.plot_persistence_diagrams(x)\n"
+    "\n"
    ]
   },
   {
diff --git a/examples/compactification.ipynb b/examples/compactification.ipynb
index f9055cd8..23b50665 100644
--- a/examples/compactification.ipynb
+++ b/examples/compactification.ipynb
@@ -31,8 +31,8 @@
     "import numpy as np\n",
     "from torch import nn\n",
     "from gdeep.models import FFNet\n",
-    "from gdeep.data import TorchDataLoader\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
+    "from gdeep.trainer import Trainer\n",
     "from torch import autograd  \n",
     "\n",
     "# plot\n",
@@ -69,8 +69,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = TorchDataLoader(name=\"Blobs\")\n",
-    "dl_tr, dl_ts = dl.build_dataloaders(batch_size=1)"
+    "bd = DatasetBuilder(name=\"Blobs\")\n",
+    "ds_tr, ds_val, _ = bd.build()\n",
+    "#train_indices = list(range(160))\n",
+    "dl = DataLoaderBuilder((ds_tr, ds_val))\n",
+    "dl_tr, dl_val, dl_ts = dl.build() #, sampler=SubsetRandomSampler(train_indices))\n"
    ]
   },
   {
@@ -79,7 +82,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "next(iter(dl_tr))"
+    "print(\"One batch from the dataloader:\", next(iter(dl_tr)))"
    ]
   },
   {
@@ -91,7 +94,7 @@
     "# train NN\n",
     "model = FFNet(arch=[3,3])\n",
     "print(model)\n",
-    "pipe = Pipeline(model, (dl_tr, dl_ts), nn.CrossEntropyLoss(), writer)\n",
+    "pipe = Trainer(model, (dl_tr, dl_ts), nn.CrossEntropyLoss(), writer)\n",
     "pipe.train(SGD, 5, False, {\"lr\": 0.01}, {\"batch_size\": 1})"
    ]
   },
diff --git a/examples/decision_boundary_tori.ipynb b/examples/decision_boundary_tori.ipynb
index 39a987fe..413d9dd5 100644
--- a/examples/decision_boundary_tori.ipynb
+++ b/examples/decision_boundary_tori.ipynb
@@ -31,8 +31,8 @@
     "import numpy as np\n",
     "from torch import nn\n",
     "from gdeep.models import FFNet\n",
-    "from gdeep.data import TorchDataLoader\n",
-    "from gdeep.pipeline import Pipeline\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
+    "from gdeep.trainer import Trainer\n",
     "from torch import autograd  \n",
     "\n",
     "# plot\n",
@@ -69,8 +69,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = TorchDataLoader(name=\"DoubleTori\")\n",
-    "dl_tr, dl_ts = dl.build_dataloaders(batch_size=1)"
+    "bd = DatasetBuilder(name=\"DoubleTori\")\n",
+    "ds_tr, ds_val, _ = bd.build()\n",
+    "#train_indices = list(range(160))\n",
+    "dl = DataLoaderBuilder((ds_tr, ds_val))\n",
+    "dl_tr, dl_val, dl_ts = dl.build()"
    ]
   },
   {
@@ -82,8 +85,8 @@
     "# train NN\n",
     "model = FFNet(arch=[3,10,10,2])\n",
     "print(model)\n",
-    "pipe = Pipeline(model, (dl_tr, dl_ts), nn.CrossEntropyLoss(), writer)\n",
-    "pipe.train(SGD, 10, False, {\"lr\":0.01}, {\"batch_size\":1})"
+    "pipe = Trainer(model, (dl_tr, dl_ts), nn.CrossEntropyLoss(), writer)\n",
+    "pipe.train(SGD, 5, False, {\"lr\":0.01}, {\"batch_size\":1})"
    ]
   },
   {
diff --git a/examples/orbit_5k_train.ipynb b/examples/orbit_5k_train.ipynb
index 57642982..bb3f06d6 100644
--- a/examples/orbit_5k_train.ipynb
+++ b/examples/orbit_5k_train.ipynb
@@ -35,10 +35,10 @@
     "from transformers.optimization import get_cosine_with_hard_restarts_schedule_with_warmup, get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup\n",
     "\n",
     "# Import the giotto-deep modules\n",
-    "from gdeep.data import OrbitsGenerator, DataLoaderKwargs\n",
+    "from gdeep.data.datasets import OrbitsGenerator, DataLoaderKwargs\n",
     "from gdeep.topology_layers import Persformer\n",
-    "from gdeep.pipeline import Pipeline\n",
-    "from gdeep.search import Gridsearch\n",
+    "from gdeep.trainer import Trainer\n",
+    "from gdeep.search import HyperParameterOptimization\n",
     "import json\n",
     "\n",
     "from optuna.pruners import MedianPruner, NopPruner"
@@ -190,7 +190,7 @@
     "optim = torch.optim.Adam(model.parameters(), 1e-3)\n",
     "\n",
     "# initialise pipeline class\n",
-    "pipe = Pipeline(model, [dl_train, None], loss_fn, writer)"
+    "pipe = Trainer(model, [dl_train, None], loss_fn, writer)"
    ]
   },
   {
@@ -233,7 +233,7 @@
     "\n",
     "# initialise gridsearch\n",
     "pruner = NopPruner()\n",
-    "search = Gridsearch(pipe, search_metric=\"accuracy\", n_trials=50, best_not_last=False, pruner=pruner)\n",
+    "search = HyperParameterOptimization(pipe, search_metric=\"accuracy\", n_trials=50, best_not_last=False, pruner=pruner)\n",
     "\n",
     "# dictionaries of hyperparameters\n",
     "optimizers_params = {\"lr\": [1e-3, 1e-0, None, True],\n",
diff --git a/gdeep/__init__.py b/gdeep/__init__.py
index 98dd0844..4d162f34 100644
--- a/gdeep/__init__.py
+++ b/gdeep/__init__.py
@@ -3,10 +3,9 @@
 
 __all__ = [
     'analysis',
-    'bert',
     'data',
     'models',
-    'pipeline',
+    'trainer',
     'search',
     'topology_layers',
     'utility',
diff --git a/gdeep/bert/__init__.py b/gdeep/bert/__init__.py
deleted file mode 100644
index 12cab5e9..00000000
--- a/gdeep/bert/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-
-from .simplified_bert_model import SimplifiedMultiHeadSelfAttention,\
-    SimplifiedBertBlock, SimplifiedBertEncoder, SimplifiedBertEmbeddings,\
-    SimplifiedBertPooler, SimplifiedBertClassifier
-
-__all__ = [
-    'SimplifiedMultiHeadSelfAttention',
-    'SimplifiedBertBlock',
-    'SimplifiedBertEncoder',
-    'SimplifiedBertEmbeddings',
-    'SimplifiedBertPooler',
-    'SimplifiedBertClassifier'
-    ]
diff --git a/gdeep/bert/bert_weight_copier.py b/gdeep/bert/bert_weight_copier.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/gdeep/bert/simplified_bert_model.py b/gdeep/bert/simplified_bert_model.py
deleted file mode 100644
index 95576cb9..00000000
--- a/gdeep/bert/simplified_bert_model.py
+++ /dev/null
@@ -1,318 +0,0 @@
-import torch
-import torch.nn as nn
-
-# Huggingface Bert implementation
-from transformers import BertConfig
-
-# extended einstein operations for tensor operations
-from einops import rearrange
-
-
-class SimplifiedMultiHeadSelfAttention(nn.Module):
-    """
-    Simplified multi head self-attention layer with einops wherever possible
-
-    Args:
-        config (BertConfig): Bert model specification
-    """
-    def __init__(self, config: BertConfig) -> None:
-        super().__init__()
-        self.config = config
-
-        # dimension of the encoder and pooler layer
-        self.dim = config.hidden_size
-        self.head = config.num_attention_heads  # number of attention heads
-        # dimension of each attention head
-        self.dim_head = int(self.dim / self.heads)
-        # linear layer splitting hidden state to query, key, value
-        self.to_qvk = nn.Linear(self.dim, 3 * self.dim, bias=True)
-        # scaling factor in row-wise softmax to prevent small gradients
-        self.scale_factor = self.dim_head ** -0.5
-
-    def forward(self, x: torch.tensor, mask: torch.tensor = None
-                ) -> torch.tensor:
-        """ Forward pass
-
-        Args:
-            x (torch.tensor): input tensor of shape [batch, tokens, dim]
-            mask (torch.tensor, optional): attention mask of shape
-                [batch, tokens, dim]. Defaults to None.
-
-        Returns:
-            torch.tensor: attention result concatenated head-wise
-        """
-        assert x.dim() == 3  # [batch, tokens, dim]
-        qkv = self.to_qvk(x)  # [batch, tokens, 3 * heads * dim_heads]
-
-        # decomposing to query, key, value each with dim [batch, heads, tokens,
-        # dim_heads]
-        q, k, v = tuple(rearrange(qkv, 'b t (qkv h d) -> qkv b h t d',
-                                  qkv=3,
-                                  h=self.head))
-
-        # attention value of shape [batch, heads, tokens, dim_heads]
-        out = self.__compute_mhsa(q, k, v, mask=mask)
-        # merge heads
-        out = rearrange(out, 'b h t d -> b t (h d)')
-        return out
-
-    def __compute_mhsa(
-                self,
-                q: torch.tensor,
-                k: torch.tensor,
-                v: torch.tensor,
-                mask: torch.tensor = None) -> torch.tensor:
-        """ Implemenatation of multi head self attention values
-
-        Args:
-            q (torch.tensor): query tensor of shape [batch, tokens, dim]
-            k (torch.tensor): key tensor of shape [batch, tokens, dim]
-            v (torch.tensor): value tensor of shape [batch, tokens, dim]
-            mask (torch.tensor, optional): attention mask of shape.
-                Defaults to None.
-
-        Returns:
-            torch.tensor: attention values of shape [batch, tokens, dim]
-        """
-        # scalar product of query and key vectors
-        scaled_dot_prod = torch.einsum('... i d, ... j d -> ... i j', q, k)\
-            * self.scale_factor
-
-        if mask is not None:
-            try:
-                scaled_dot_prod = scaled_dot_prod.masked_fill(
-                    mask == 0,
-                    -float('inf')
-                    )
-            except Exception:
-                print(
-                    """"
-                    mask and scaled dot product do not have the same
-                    dimensionality!
-                    """
-                    )
-
-        # Row-wise softmax
-        attention = torch.softmax(scaled_dot_prod, dim=-1)
-        return torch.einsum('... i j, ... j d -> ... i d', attention, v)
-
-
-class SimplifiedBertBlock(nn.Module):
-    """
-    Bert transformer block consisting of a multi head attention layer
-    and a feed-forward layer
-
-    Args:
-        config (BertConfig): Bert model specification
-    """
-    def __init__(self, config: BertConfig) -> None:
-        super().__init__()
-        # set activation function as nn.Module
-        activation = getattr(nn, config.hidden_act.upper())()
-
-        dim = config.hidden_size  # token's vector length
-        # dimension of the inner linear layer
-        intermediate_dim = config.intermediate_size
-        heads = config.num_attention_heads  # number of attention heads
-        # value added to the denominator of the normalization for numerical
-        # stability
-        layer_norm_eps = config.layer_norm_eps
-        dropout = config.dropout  # probabilty of dropping values
-
-        #  attention layer
-        self.mhsa = SimplifiedMultiHeadSelfAttention(dim, heads)
-        self.attention_linear = nn.Linear(dim, dim)
-        self.drop_1 = nn.Dropout(dropout)
-        self.norm_1 = nn.LayerNorm(dim, eps=layer_norm_eps)
-        self.drop_2 = nn.Dropout(dropout)
-
-        #  feed forwad layer, see https://arxiv.org/abs/1706.03762
-        self.linear = nn.Sequential(
-            nn.Linear(dim, intermediate_dim),
-            activation,
-            nn.Linear(intermediate_dim, dim),
-            nn.Dropout(dropout)
-        )
-        self.norm_2 = nn.LayerNorm(dim, eps=layer_norm_eps)
-
-    def forward(self, x: torch.tensor, mask: torch.tensor = None
-                ) -> torch.tensor:
-        attention_output = self.attention_linear(
-            self.drop_1(self.mhsa(x, mask))
-            )
-        attention_normalized = self.drop_2(self.norm_1(attention_output + x))
-        return self.norm_2(
-            self.linear(attention_normalized) + attention_normalized
-            )
-
-
-class SimplifiedBertEncoder(nn.Module):
-    """
-    Bert encoder consisting of stacked Bert blocks
-
-    Args:
-        config (BertConfig): Bert model specification
-    """
-    def __init__(self, config: BertConfig) -> None:
-        super().__init__()
-
-        # number of transformer blocks
-        blocks = config.num_hidden_layers
-
-        # stacking ´blocks´ bert transformer blocks
-        self.block_list = [
-                        SimplifiedBertBlock(config)
-                        for _ in range(blocks)
-        ]
-        # properly register ´block_list´ to make them visible to all ´Module´
-        # methods
-        self.layers = nn.ModuleList(self.block_list)
-
-    def forward(self, x: torch.tensor, mask: torch.tensor = None
-                ) -> torch.tensor:
-        for layer in self.layers:
-            x = layer(x, mask)
-        return x
-
-
-class SimplifiedBertEmbeddings(nn.Module):
-    """
-    Simplified Bert Embedding layer consisting of
-    * Word embedding layer
-    * token type embedding layer
-    * (absolute) positional encoding layer with trainable weights
-    Only works for absolute embeddings up to now.
-
-    Get the embeddings from one-hot encodings of the words.
-
-    Args:
-    config (BertConfig): Bert model specification
-    """
-    def __init__(self, config: BertConfig):
-        super().__init__()
-
-        dim = config.hidden_size  # token's vector length
-
-        # vocabulary size of Bert model, defaults to 30522
-        vocab_size = config.vocab_size
-
-        # Token id for padding, for example when batching sequences,
-        # defaults to 0
-        pad_token_id = config.pad_token_id
-
-        # number of different token types, default to 2
-        type_vocab_size = config.type_vocab_size
-
-        # Maximal sequence length that can be passed to the model,
-        # defaults to 512
-        max_position_embeddings = config.max_position_embeddings
-
-        # value added to the denominator of the normalization for numerical
-        # stability
-        layer_norm_eps = config.layer_norm_eps
-        dropout = config.dropout  # probabilty of dropping values
-
-        self.word_embeddings = nn.Embedding(vocab_size,
-                                            dim,
-                                            padding_idx=pad_token_id)
-        self.token_type_embeddings = nn.Embedding(type_vocab_size,
-                                                  dim)
-        # absolute positional embedding
-        self.position_embeddings = nn.Embedding(max_position_embeddings, dim)
-
-        # position_ids that will be encoded by the ´position_embeddings´ layer
-        # register_buffered are non-trainable weigths
-        self.register_buffer(
-            "position_ids",
-            torch.arange(config.max_position_embeddings).expand((1, -1))
-        )
-
-        self.LayerNorm = nn.LayerNorm(dim, layer_norm_eps)
-        self.dropout_layer = nn.Dropout(dropout)
-
-    def forward(self,
-                input_ids: torch.tensor,
-                token_type_ids: torch.tensor
-                ):
-        seq_length = input_ids.shape[1]
-
-        inputs_embeds = self.word_embeddings(input_ids)  # word embedding+
-
-        # token embedding
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
-        # shrink position_ids to seq_length
-        position_ids = self.position_ids[:, : seq_length]
-        absolute_embeddings = self.position_embeddings(position_ids)
-
-        embeddings = inputs_embeds + token_type_embeddings\
-            + absolute_embeddings
-
-        embeddings = self.LayerNorm(embeddings)
-        embeddings = self.dropout_layer(embeddings)
-        return embeddings
-
-
-class SimplifiedBertPooler(nn.Module):
-    """
-    Bert Pooler (only for classifications tasks)
-    Takes the output representation corresponding to the [CLS] token
-    and uses it for downstream tasks like sentiment analysis or next
-    sentence prediction. It applies a linear transformation over the
-    [CLS] token.
-
-    Args:
-        config (BertConfig): Bert model specification
-    """
-    def __init__(self, config):
-        super().__init__()
-        dim = config.hidden_dim  # token's vector length
-
-        self.dense = nn.Linear(dim, dim)
-        self.activation = nn.Tanh()
-
-    def forward(self, x):
-        # extract classification token
-        first_token_tensor = x[:, 0]
-        # feed-forward layer
-        pooled_output = self.dense(first_token_tensor)
-        pooled_output = self.activation(pooled_output)
-        return pooled_output
-
-
-class SimplifiedBertClassifier(nn.Module):
-    """
-    Bert classifier with ´num_labels´ of output classes consisting of
-    * Embedding layer
-    * Bert encoder consisting of stacked self attention layers
-    * Pooler layer
-    * Linear classifier layer computing the logit probabilities of the
-      different classes
-
-    Args:
-        config (BertConfig): Bert model specification
-    """
-    def __init__(self, config: BertConfig, num_labels: int = 2) -> None:
-        super().__init__()
-        dim = config.hidden_size  # token's vector length
-
-        self.embedding = SimplifiedBertEmbeddings(config)
-
-        self.encoder = SimplifiedBertEncoder(config)
-
-        self.pooler = SimplifiedBertPooler(config)
-
-        self.classifier = nn.Linear(dim, num_labels)
-
-    def forward(self,
-                input_ids: torch.tensor,
-                token_type_ids: torch.tensor,
-                attention_mask: torch.tensor
-                ) -> torch.tensor:
-        x = self.embedding(input_ids, token_type_ids)
-
-        x = self.encoder(x, attention_mask)
-
-        x = self.pooler(x)
-
-        return self.classifier(x)
diff --git a/gdeep/data/__init__.py b/gdeep/data/__init__.py
index 4d619dfd..630b9e21 100644
--- a/gdeep/data/__init__.py
+++ b/gdeep/data/__init__.py
@@ -1,44 +1,15 @@
-from .categorical_data import CategoricalDataCloud
-from .tori import Rotation, \
-    CreateToriDataset, GenericDataset
-from .dataset_cloud import DatasetCloud
-from ._data_cloud import _DataCloud
-from .torch_datasets import TorchDataLoader, \
-    DataLoaderFromImages, DataLoaderFromArray, DlBuilderFromDataCloud
-from .text_dataset import TextDataset, TextDatasetQA, \
-    TextDatasetTranslation
-from .parallel_orbit import generate_orbit_parallel, create_pd_orbits,\
-    OrbitsGenerator, DataLoaderKwargs
-from .preprocessing_pipes import Normalisation, \
-    PreprocessTextTranslation, PreprocessTextData, PreprocessingPipeline, \
-    PreprocessTextQA, PreprocessTextLabel, PreprocessTextQATarget
-from .preprocessing_interface import AbstractPreprocessing
 
+from .transforming_dataset import TransformingDataset
+from .preprocessing_pipeline import PreprocessingPipeline
+from .dataset_factory import DatasetFactory
+from .abstract_preprocessing import AbstractPreprocessing
+from ._utils import MissingVocabularyError
 
 
 __all__ = [
-    'Rotation',
-    'CategoricalDataCloud',
-    'CreateToriDataset',
-    'GenericDataset',
-    'PreprocessTextQATarget',
-    'TextDataset',
-    'AbstractPreprocessing',
+    'TransformingDataset',
     'PreprocessingPipeline',
-    'PreprocessTextQA',
-    'TextDatasetTranslation',
-    'PreprocessTextLabel',
-    'PreprocessTextData',
-    'TextDatasetQA',
-    'TorchDataLoader',
-    'generate_orbit_parallel',
-    'create_pd_orbits',
-    'OrbitsGenerator',
-    'DataLoaderKwargs',
-    'DataLoaderFromImages',
-    'PreprocessTextTranslation',
-    'TextDatasetTranslation',
-    'DataLoaderFromArray',
-    'DatasetCloud',
-    'DlBuilderFromDataCloud'
+    'DatasetFactory',
+    'AbstractPreprocessing',
+    'MissingVocabularyError'
     ]
diff --git a/gdeep/data/_utils.py b/gdeep/data/_utils.py
new file mode 100644
index 00000000..f8e14195
--- /dev/null
+++ b/gdeep/data/_utils.py
@@ -0,0 +1,6 @@
+
+
+class MissingVocabularyError(Exception):
+    """Class to raise the missing vocabulary
+    esception in the toknizers"""
+    pass
\ No newline at end of file
diff --git a/gdeep/data/abstract_preprocessing.py b/gdeep/data/abstract_preprocessing.py
new file mode 100644
index 00000000..f7b3a2a1
--- /dev/null
+++ b/gdeep/data/abstract_preprocessing.py
@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+import json
+import os
+from typing import Generic, TypeVar
+import warnings
+import jsonpickle
+
+from .transforming_dataset import TransformingDataset
+
+from torch.utils.data import Dataset
+
+from gdeep.data.transforming_dataset import TransformingDataset
+
+R = TypeVar('R')
+S = TypeVar('S')
+
+class AbstractPreprocessing(ABC, Generic[R, S]):
+    @abstractmethod
+    def fit_to_dataset(self, dataset: Dataset[R]) -> None:
+        pass
+    
+    @abstractmethod
+    def __call__(self, x: R) -> S:
+        pass
+    
+    def transform(self, x: R) -> S:
+        return self(x)
+
+    def attach_transform_to_dataset(self, dataset: Dataset[R]) -> Dataset[S]:
+        return TransformingDataset(dataset, self.transform)
+    
+    def save_pretrained(self, path:str) -> None:
+        with open(os.path.join(path, self.__class__.__name__ + ".json"), "w") as outfile:
+            whole_class = jsonpickle.encode(self)  # type: ignore
+            json.dump(whole_class, outfile)
+
+    def load_pretrained(self, path:str) -> None:
+        try:
+            with open(os.path.join(path,self.__class__.__name__ + ".json"), "r") as infile:
+                whole_class = json.load(infile)
+                self = jsonpickle.decode(whole_class)  # type: ignore
+        except FileNotFoundError:
+            warnings.warn("The transformation file does not exist; attempting to run"
+                          " the transformation anyway...")
+            
\ No newline at end of file
diff --git a/gdeep/data/dataset_factory.py b/gdeep/data/dataset_factory.py
new file mode 100644
index 00000000..65af4aba
--- /dev/null
+++ b/gdeep/data/dataset_factory.py
@@ -0,0 +1,48 @@
+
+from typing import Dict, Any
+
+
+
+# Define the dataset factory class for the tori dataset and torchvision datasets
+# using the factory design pattern
+# https://realpython.com/factory-method-python/
+
+
+class DatasetFactory(object):
+    """ Dataset factory class for the tori dataset and torchvision datasets
+    using the factory design pattern
+    
+    Examples::
+
+        # Create a dataset for the tori dataset
+        dataset = get_dataset("Tori", name="DoubleTori", n_points=100)
+        
+        # Create the MNIST dataset
+        dataset = get_dataset("Torchvision", name="MNIST")
+
+    """
+
+    _builders: Dict[str, Any]
+
+    def __init__(self):
+        self._builders = {}
+        
+    def register_builder(self, key: str, builder:Any):
+        """this method adds to the internal builders dictionary
+        new dataloader builders
+        """
+        self._builders[key] = builder    # type: ignore
+        
+    def build(self, key:str, **kwargs) -> Any:  # type: ignore
+        """This method returns the DataLoader builder
+        corresponding to the input key.
+
+        Args:
+            key:
+                the name of the dataset
+        """
+        builder = self._builders.get(key)    # type: ignore
+        if builder is None:
+            raise ValueError(f"No builder registered for key {key}")
+        return builder(**kwargs)  # type: ignore
+
diff --git a/gdeep/data/datasets/__init__.py b/gdeep/data/datasets/__init__.py
new file mode 100644
index 00000000..89ab4d68
--- /dev/null
+++ b/gdeep/data/datasets/__init__.py
@@ -0,0 +1,31 @@
+from .categorical_data import CategoricalDataCloud
+from .tori import Rotation, ToriDataset
+from .dataset_cloud import DatasetCloud
+from ._data_cloud import _DataCloud
+from .build_datasets import DatasetBuilder, get_dataset
+from .base_dataloaders import DataLoaderBuilder, \
+    AbstractDataLoaderBuilder
+from .dataset_for_image import ImageClassificationFromFiles
+from .dataset_form_array import FromArray
+from .dataloader_cloud import DlBuilderFromDataCloud
+from .parallel_orbit import generate_orbit_parallel, create_pd_orbits,\
+    OrbitsGenerator, DataLoaderKwargs
+
+
+__all__ = [
+    'Rotation',
+    'CategoricalDataCloud',
+    'AbstractDataLoaderBuilder',
+    'ToriDataset',
+    'DataLoaderBuilder',
+    'DatasetBuilder',
+    'generate_orbit_parallel',
+    'create_pd_orbits',
+    'OrbitsGenerator',
+    'DataLoaderKwargs',
+    'get_dataset',
+    'ImageClassificationFromFiles',
+    'FromArray',
+    'DatasetCloud',
+    'DlBuilderFromDataCloud'
+    ]
diff --git a/gdeep/data/_data_cloud.py b/gdeep/data/datasets/_data_cloud.py
similarity index 100%
rename from gdeep/data/_data_cloud.py
rename to gdeep/data/datasets/_data_cloud.py
diff --git a/gdeep/data/datasets/base_dataloaders.py b/gdeep/data/datasets/base_dataloaders.py
new file mode 100644
index 00000000..8d154d98
--- /dev/null
+++ b/gdeep/data/datasets/base_dataloaders.py
@@ -0,0 +1,78 @@
+import json
+import os
+import shutil
+import warnings
+from abc import ABC, abstractmethod
+from os.path import join
+from collections.abc import Iterable
+from typing import Any, Callable, Dict, Optional, Tuple, \
+    TypeVar, Union, List
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image, UnidentifiedImageError
+from sympy import false
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import Resize, ToTensor
+from tqdm import tqdm
+
+from .build_datasets import get_dataset
+from .dataset_cloud import DatasetCloud
+from ..transforming_dataset import TransformingDataset
+
+
+Tensor = torch.Tensor
+T = TypeVar('T')
+
+class AbstractDataLoaderBuilder(ABC):
+    """The abstractr class to interface the
+    Giotto dataloaders"""
+    @abstractmethod
+    def build(self):
+        pass
+
+
+class DataLoaderBuilder(AbstractDataLoaderBuilder):
+    """This class builds, out of a tuple of datasets, the
+    corresponding dataloaders. Note that this class would
+    use the same parameters for all the datasets
+
+    Args:
+        tuple_of_datasets :
+            the tuple eith the traing, validation and test
+            datasets. Also one or two elemennts are acceptable:
+            they will be considered as training first and
+            validation afterwards.
+    """
+    def __init__(self, tuple_of_datasets: List[Dataset[Any]]) -> None:
+        self.tuple_of_datasets = tuple_of_datasets
+        assert len(tuple_of_datasets) <= 3, "Too many Dataset inserted: maximum 3."
+
+    def build(self,
+              tuple_of_kwargs:Optional[List[Dict[str, Any]]]=None
+              ) -> List[DataLoader[Any]]:
+        """This method accepts the arguments of the torch
+        Dataloader and applies them when creating the
+        tuple
+
+        Args:
+            tuple_of_kwargs:
+                List of dictionaries, each one being the
+                kwargs for the corresponding DataLoader
+        """
+        if tuple_of_kwargs:
+            assert len(tuple_of_kwargs) == len(self.tuple_of_datasets), \
+                "Cannot match the dataloaders and the parameters. "
+            out: List = []
+            for dataset, kwargs in zip(self.tuple_of_datasets, tuple_of_kwargs):
+                out.append(DataLoader(dataset, **kwargs))
+            out += [None] * (3 - len(out))
+            return out
+        else:
+            out: List = []
+            for i, dataset in enumerate(self.tuple_of_datasets):
+                out.append(DataLoader(dataset))
+            out += [None] * (3 - len(out))
+            return out
+
diff --git a/gdeep/data/datasets/build_datasets.py b/gdeep/data/datasets/build_datasets.py
new file mode 100644
index 00000000..c25f5162
--- /dev/null
+++ b/gdeep/data/datasets/build_datasets.py
@@ -0,0 +1,173 @@
+import json
+import os
+import shutil
+import warnings
+from abc import ABC, abstractmethod
+from os.path import join
+from collections.abc import Iterable
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union, List
+
+from torchtext.data import to_map_style_dataset
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image, UnidentifiedImageError
+from sympy import false
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import Resize, ToTensor
+from tqdm import tqdm
+from torch.utils.data import Dataset
+from torchvision import datasets
+from torchtext import datasets as textdatasets
+
+from .tori import ToriDataset
+from gdeep.utility import DEFAULT_DOWNLOAD_DIR
+from .dataset_cloud import DatasetCloud
+from ..dataset_factory import DatasetFactory
+from ..transforming_dataset import TransformingDataset
+
+Tensor = torch.Tensor
+T = TypeVar('T')
+
+
+class TorchvisionDatasetBuilder(object):
+    """Builder class for the torchvision dataset
+        and all its variations"""
+
+    def __init__(self, dataset_name: str):
+        self.dataset_name = dataset_name
+
+    def __call__(self, **kwargs):  # type: ignore
+        return datasets.__getattribute__(self.dataset_name)(  # type: ignore
+            root=DEFAULT_DOWNLOAD_DIR, **kwargs  # type: ignore
+        )
+
+
+class TorchtextDatasetBuilder(object):
+    """Builder class for the torchtext dataset
+        and all its variations"""
+
+    def __init__(self, dataset_name: str):
+        self.dataset_name = dataset_name
+
+    def __call__(self, **kwargs):  # type: ignore
+        return textdatasets.__getattribute__(self.dataset_name)(  # type: ignore
+            root=DEFAULT_DOWNLOAD_DIR, **kwargs  # type: ignore
+        )
+
+
+class ToriDatasetBuilder(object):
+    """Builder class for the torus dataset
+    and all its variations"""
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    def __call__(self, **kwargs):  # type: ignore
+        return ToriDataset(name=self.name, **kwargs)
+
+
+def get_dataset(key: str, **kwargs) -> Tuple[Dataset[Any]]:
+    """ Get a dataset from the factory
+
+    Args:
+        key :
+            The name of the dataset, correspondiong
+            to the key in the list of builders
+        **kwargs:
+            The keyword arguments to pass to the dataset builder
+
+    Returns:
+        torch.utils.data.Dataset:
+            The dataset
+    """
+    factory = DatasetFactory()
+
+    factory.register_builder("DoubleTori", ToriDatasetBuilder("DoubleTori"))
+    factory.register_builder("Blobs", ToriDatasetBuilder("Blobs"))
+    factory.register_builder("EntangledTori", ToriDatasetBuilder("EntangledTori"))
+
+    for dataset_name in datasets.__all__:
+        factory.register_builder(dataset_name, TorchvisionDatasetBuilder(dataset_name))
+
+    for dataset_name in textdatasets.__all__:
+        factory.register_builder(dataset_name, TorchtextDatasetBuilder(dataset_name))
+
+    return factory.build(key, **kwargs)  # type: ignore
+
+
+class DatasetBuilder:
+    """Class to obtain Datasets from the classical
+    datasets available on pytorch. Also the torus dataset
+    and all its variations can be found here
+
+    Args:
+        name:
+            check the available datasets at
+            https://pytorch.org/vision/stable/datasets.html
+            and https://pytorch.org/text/stable/datasets.html
+        convert_to_map_dataset:
+            whether to conver to the MapDataset or to keep
+            IterableDataset
+
+    """
+    train_ds: Dataset[Any]
+    valid_ds: Optional[Dataset[Any]]
+    test_ds: Optional[Dataset[Any]]
+
+    def __init__(self, name: str="MNIST", convert_to_map_dataset:bool=False) -> None:
+        self.convert_to_map_dataset = convert_to_map_dataset
+        self.name = name
+
+    def build(self, **kwargs) -> Tuple[Dataset[Any],
+                                                Optional[Dataset[Any]],
+                                                Optional[Dataset[Any]]]:
+        """Method that returns the dataset.
+
+        Args:
+            kwargs:
+                the arguments to pass to the dataset builder.
+                For example, you may want to use the options
+                ``split=("train","dev")`` or ``split=("train","test")``
+        """
+        dataset_tuple: Tuple[Dataset[Any]] = get_dataset(self.name, **kwargs)  # type: ignore
+    
+        if len(dataset_tuple) == 1:
+            self.train_ds = dataset_tuple[0]
+            self.valid_ds = None
+            self.test_ds = None
+        elif len(dataset_tuple) == 2:
+            self.train_ds = dataset_tuple[0]
+            self.valid_ds = dataset_tuple[1]
+            self.test_ds = None
+        elif len(dataset_tuple) == 3:
+            self.train_ds = dataset_tuple[0]
+            self.valid_ds = dataset_tuple[1]
+            self.test_ds = dataset_tuple[2]
+        else:
+            self.train_ds = dataset_tuple
+            self.valid_ds = None
+            self.test_ds = None
+
+        if self.convert_to_map_dataset:
+            self.train_ds, self.valid_ds, self.test_ds = self._convert(self.train_ds,
+                                                                       self.valid_ds,
+                                                                       self.test_ds)
+        return self.train_ds, self.valid_ds, self.test_ds
+
+    def _convert(self, training_data: Dataset[Any],
+                 validation_data: Optional[Dataset[Any]]=None,
+                 test_data: Optional[Dataset[Any]]=None ) -> Tuple[Dataset[Any],
+                                                                   Optional[Dataset[Any]],
+                                                                   Optional[Dataset[Any]]]:
+        """This private method converts and IterableDataset
+        to a MapDataset"""
+        if isinstance(training_data, torch.utils.data.IterableDataset):
+            training_data = to_map_style_dataset(training_data)
+        if validation_data is not None:
+            if isinstance(validation_data, torch.utils.data.IterableDataset):
+                validation_data = to_map_style_dataset(validation_data)
+        if test_data is not None:
+            if isinstance(test_data, torch.utils.data.IterableDataset):
+                test_data = to_map_style_dataset(test_data)
+        return training_data, validation_data, test_data
diff --git a/gdeep/data/categorical_data.py b/gdeep/data/datasets/categorical_data.py
similarity index 100%
rename from gdeep/data/categorical_data.py
rename to gdeep/data/datasets/categorical_data.py
diff --git a/gdeep/data/datasets/dataloader_cloud.py b/gdeep/data/datasets/dataloader_cloud.py
new file mode 100644
index 00000000..91ace686
--- /dev/null
+++ b/gdeep/data/datasets/dataloader_cloud.py
@@ -0,0 +1,173 @@
+import json
+import os
+import shutil
+import warnings
+from abc import ABC, abstractmethod
+from os.path import join
+from typing import Any, Callable, Dict, Optional, Tuple, \
+    TypeVar, Union, List
+
+import numpy as np
+import pandas as pd
+import torch
+
+from sympy import false
+from torch.utils.data import DataLoader, Dataset
+from .dataset_form_array import FromArray
+from tqdm import tqdm
+
+from .dataset_cloud import DatasetCloud
+from .base_dataloaders import DataLoaderBuilder
+
+from .base_dataloaders import AbstractDataLoaderBuilder
+
+
+Tensor = torch.Tensor
+T = TypeVar('T')
+
+
+class DlBuilderFromDataCloud(AbstractDataLoaderBuilder):
+    """Class that loads data from Google Cloud Storage
+    
+    This class is useful to build dataloaders from a dataset stored in
+    the GDeep Dataset Cloud on Google Cloud Storage.
+
+    The constructor takes the name of a dataset as a string, and a string
+    for the download directory. The constructor will download the dataset
+    to the download directory. The dataset is downloaded in the version
+    used by Datasets Cloud, which may be different from the version
+    used by the dataset's original developers.
+    
+    Args:
+        dataset_name (str):
+            The name of the dataset.
+        download_dir (str):
+            The directory where the dataset will be downloaded.
+        use_public_access (bool):
+            Whether to use public access. If you want
+            to use the Google Cloud Storage API, you must set this to True.
+            Please make sure you have the appropriate credentials.
+        path_to_credentials (str):
+            Path to the credentials file.
+            Only used if public_access is False and credentials are not
+            provided. Defaults to None.
+        
+
+    Returns:
+        torch.utils.data.DataLoader: The dataloader for the dataset.
+
+    Raises:
+        ValueError:
+            If the dataset_name is not a valid dataset that exists
+            in Datasets Cloud.
+        ValueError:
+            If the download_directory is not a valid directory.
+    """
+    def __init__(self,
+                 dataset_name: str,
+                 download_directory: str,
+                 use_public_access: bool=True,
+                 path_to_credentials: Union[None, str]=None,
+                 ):
+        self.dataset_name = dataset_name
+        self.download_directory = download_directory
+        
+        # Download the dataset if it does not exist
+        self.download_directory
+        
+        self._download_dataset(use_public_access=use_public_access, 
+                               path_to_credentials=path_to_credentials)
+
+        self.dl_builder = None
+        
+        # Load the metadata of the dataset
+        with open(join(self.download_directory, self.dataset_name,
+                       "metadata.json")) as f:
+            self.dataset_metadata = json.load(f)
+            
+        # Load the data and labels of the dataset
+        if self.dataset_metadata['data_type'] == 'tabular':
+            if self.dataset_metadata['data_format'] == 'pytorch_tensor':
+                data = torch.load(join(self.download_directory, 
+                                       self.dataset_name, "data.pt"))
+                labels = torch.load(join(self.download_directory, 
+                                         self.dataset_name, "labels.pt"))
+
+                self.dl_builder = DataLoaderBuilder((FromArray(data, labels),))
+            elif self.dataset_metadata['data_format'] == 'numpy_array':
+                data = np.load(join(self.download_directory,
+                                    self.dataset_name, "data.npy"))
+                labels = np.load(join(self.download_directory,
+                                      self.dataset_name, "labels.npy"))
+                self.dl_builder = DataLoaderBuilder((FromArray(data, labels),))
+            else:
+                raise ValueError("Data format {}"\
+                    .format(self.dataset_metadata['data_format']) +
+                                 "is not yet supported.")
+        else:
+            raise ValueError("Dataset type {} is not yet supported."\
+                .format(self.dataset_metadata['data_type']))
+
+    def _download_dataset(self,
+                          path_to_credentials: Union[None, str] =None,
+                          use_public_access: bool=True,) -> None:
+        """Only download if the download directory does not exist already
+        and if download directory contains at least three files (metadata,
+        data, labels).
+        
+        Args:
+            path_to_credentials (str):
+                Path to the credentials file.
+            use_public_access (bool):
+                Whether to use public access. If you want
+                to use the Google Cloud Storage API, you must set this to True.
+                
+        Returns:
+            None
+        """
+        if (not os.path.isdir(join(self.download_directory, self.dataset_name))
+                            or len(os.listdir(join(self.download_directory,
+                                            self.dataset_name))) < 3):
+            # Delete the download directory if it exists but does not contain
+            # the wanted number of files
+            if (os.path.isdir(join(self.download_directory, self.dataset_name))
+                and
+                len(os.listdir(join(self.download_directory,
+                                   self.dataset_name))) < 3): # type: ignore
+                print("Deleting the download directory because it does "+
+                      "not contain the dataset")
+                shutil.rmtree(self.download_directory, ignore_errors=True)
+                
+            print("Downloading dataset {} to {}"\
+                    .format(self.dataset_name, self.download_directory))
+            dataset_cloud = DatasetCloud(self.dataset_name,
+                                    download_directory=self.download_directory,
+                                    path_to_credentials=path_to_credentials,
+                                    use_public_access = use_public_access,
+                                    )
+            dataset_cloud.download()
+            del dataset_cloud
+        else:
+            print("Dataset '%s' already downloaded" % self.dataset_name)
+
+    def get_metadata(self) -> Dict[str, Any]:
+        """ Returns the metadata of the dataset.
+        
+        Returns:
+            Dict[str, Any]:
+                The metadata of the dataset.
+        """
+        return self.dataset_metadata
+    
+    def build(self, tuple_of_kwargs: List[Dict[str, Any]])\
+        -> Tuple[DataLoader, DataLoader, DataLoader]:
+        """Builds the dataloaders for the dataset.
+        
+        Args:
+            **kwargs: Arguments for the dataloader builder.
+            
+        Returns:
+            Tuple[DataLoader, DataLoader, DataLoader]:
+                The dataloaders for the dataset (train, validation, test).
+        """
+        return self.dl_builder.build(tuple_of_kwargs) # type: ignore
diff --git a/gdeep/data/dataset_cloud.py b/gdeep/data/datasets/dataset_cloud.py
similarity index 96%
rename from gdeep/data/dataset_cloud.py
rename to gdeep/data/datasets/dataset_cloud.py
index 78cd68d4..e54f9d6d 100644
--- a/gdeep/data/dataset_cloud.py
+++ b/gdeep/data/datasets/dataset_cloud.py
@@ -1,16 +1,14 @@
-from ._data_cloud import _DataCloud
 
-import json
 import os
 from os import remove
 from os.path import join, exists
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Set
 
+import json
 import wget  # type: ignore
 
-
+from ._data_cloud import _DataCloud  # type: ignore
 from gdeep.utility.constants import DEFAULT_DOWNLOAD_DIR, DATASET_BUCKET_NAME
-from gdeep.utility.utils import get_checksum
 
 class DatasetCloud():
     """DatasetCloud class to handle the download and upload
@@ -122,15 +120,14 @@ def _download_using_api(self) -> None:
         """
         self._check_public_access()
         # List of existing datasets in the cloud.
-        existing_datasets = set([blob.name.split('/')[0]
-                                 for blob in
-                                 self._data_cloud.bucket.list_blobs()\
-            if blob.name != "giotto-deep-big.png"])
+        existing_datasets: Set[str] = set([blob.name.split('/')[0]  # type: ignore
+                                 for blob in  # type: ignore
+                                 self._data_cloud.bucket.list_blobs()  # type: ignore
+            if blob.name != "giotto-deep-big.png"])  # type: ignore
         if self.name not in existing_datasets:
             raise ValueError("Dataset {} does not exist in the cloud."\
                 .format(self.name) +
-                             "Available datasets are: {}."\
-                                 .format(existing_datasets))
+                             "Available datasets are: {existing_datasets}.")
         if not self._does_dataset_exist_locally():
             self._create_dataset_folder()
         self._data_cloud.download_folder(self.name + '/')
@@ -184,10 +181,13 @@ def _download_using_url(self) -> None:
         with open(join(self.download_directory,  # type: ignore
                        self.name, 'metadata.json')) as f:
             metadata = json.load(f)
+        # filetype: Literal['pt', 'npy']
         if metadata['data_format'] == "pytorch_tensor":
             filetype = "pt"
         elif metadata['data_format'] == "numpy_array":
             filetype = "npy"
+        else:
+            raise ValueError(f"Unknown data format: {metadata['data_format']}")
         self._data_cloud.download_file(self.name + "/data." + filetype)
         self._data_cloud.download_file(self.name + "/labels." + filetype)
         
@@ -201,7 +201,7 @@ def get_existing_datasets(self) -> List[str]:
         if self.use_public_access:
             datasets_local = "tmp_datasets.json"
             # Download the dataset list json file using the public URL.
-            wget.download(self.public_url + 'datasets.json', datasets_local)
+            wget.download(self.public_url + 'datasets.json', datasets_local)  # type: ignore
             datasets = json.load(open(datasets_local))
             
             # Remove duplicates. This has to be fixed in the future.
@@ -356,9 +356,6 @@ def _upload_metadata(self,
             None
         """
         self._check_public_access()
-        if self.metadata == None:
-            raise Exception("No metadata to upload. " #NOSONAR
-                            + "Please create metadata using create_metadata.")
         self._data_cloud.upload_file(
             path, str(self.metadata['name']) + '/'  # type: ignore
             + 'metadata.json',
diff --git a/gdeep/data/datasets/dataset_for_image.py b/gdeep/data/datasets/dataset_for_image.py
new file mode 100644
index 00000000..176c3a6e
--- /dev/null
+++ b/gdeep/data/datasets/dataset_for_image.py
@@ -0,0 +1,82 @@
+import json
+import os
+import shutil
+import warnings
+from abc import ABC, abstractmethod
+from os.path import join
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union
+
+import numpy as np
+from copy import deepcopy
+import pandas as pd
+import torch
+from PIL import Image, UnidentifiedImageError
+from sympy import false
+from torch.utils.data import DataLoader, Dataset
+from torchtext import datasets as textds
+from torchvision import datasets
+from torchvision.transforms import Resize, ToTensor
+from tqdm import tqdm
+
+
+
+Tensor = torch.Tensor
+T = TypeVar('T')
+
+
+class ImageClassificationFromFiles(Dataset[Any]):
+    """This class is useful to build a dataset
+    directly from image files
+    
+    Args:
+        img_folder (string):
+            The path to the folder where the training
+            images are located
+        labels_file (string):
+            The path and file name of the labels.
+            It shall be a ``.csv`` file with two columns.
+            The first columns contains the name of the
+            image and the second one contains the
+            label value
+        transform (AbstractPreprocessing):
+            the instance of the class of preprocessing.
+            It inherits from ``AbstractPreprocessing``
+        target_transform (AbstractPreprocessing):
+            the instance of the class of preprocessing.
+            It inherits from ``AbstractPreprocessing``
+    """
+
+    img_folder: str
+    labels_file: str
+    img_labels: pd.DataFrame
+    
+    def __init__(self, img_folder: str=".",
+                 labels_file:str="labels.csv",
+                 ) -> None:
+
+        self.img_folder = img_folder
+        self.img_labels = pd.read_csv(labels_file)
+
+
+    def __len__(self) -> int:
+        return len(self.img_labels)
+
+    def __getitem__(self, idx:int) -> Tuple[Any, Union[str, int]]:
+        image = self._get_image(idx)
+        image_out = deepcopy(image)
+        label = self.img_labels.iloc[idx, 1]
+        image.close()
+        return image_out, label
+
+    def _get_image(self, idx: int) -> Any:
+        """this method gets the i-th image in the labels.csv
+        file.
+        """
+        img_path: str = os.path.join(self.img_folder, self.img_labels.iloc[idx, 0])
+        try:
+            image = Image.open(img_path)
+        except UnidentifiedImageError:
+            warnings.warn(f"The image {img_path} canot be loaded. Skipping it.")
+            return None
+        return image
+
diff --git a/gdeep/data/datasets/dataset_form_array.py b/gdeep/data/datasets/dataset_form_array.py
new file mode 100644
index 00000000..b3ee1992
--- /dev/null
+++ b/gdeep/data/datasets/dataset_form_array.py
@@ -0,0 +1,60 @@
+
+import os
+import shutil
+from typing import Any, Callable, Dict, Optional, \
+    Optional, Tuple, TypeVar, Union
+
+import numpy as np
+import torch
+
+from sympy import false
+from torch.utils.data import DataLoader, Dataset
+
+Array = np.ndarray
+Tensor = torch.Tensor
+
+
+class FromArray(Dataset[Tuple[Tensor, Tensor]]):
+    """This class is useful to build dataloaders
+    from a array of X and y. Tensors are also supported.
+
+    Args:
+        X :
+            The data. The first dimension is the datum index
+        y :
+            The labels, need to match the first dimension
+            with the data
+
+    """
+    def __init__(self, x: Union[Tensor, Array],
+                 y: Union[Tensor, Array]
+                 ) -> None:
+        self.x = self._from_numpy(x)
+        y = self._from_numpy(y)
+        self.y = self._long_or_float(y)
+
+    @staticmethod
+    def _from_numpy(x: Union[Tensor, Array]) -> Tensor:
+        """this is an upgrade of ``torch.from_numpy()``
+        that also allows tensor input.
+        """
+        if isinstance(x, torch.Tensor):
+            return x
+        return torch.from_numpy(x)
+
+    @staticmethod
+    def _long_or_float(y: Union[Tensor, Array]) -> Tensor:
+        """This private method converts the labels to either
+        a long tensor of a float tensor"""
+        if isinstance(y, torch.Tensor):
+            return y
+        if y.dtype in [np.float16, np.float32, np.float64]:
+            return torch.tensor(y).float()
+        return torch.tensor(y).long()
+
+    def __len__(self):
+        return self.y.shape[0]
+
+    def __getitem__(self, idx:int) -> Tuple[Tensor, Tensor]:
+        x, y = (self.x[idx], self.y[idx])
+        return x, y
diff --git a/gdeep/data/graph_datasets.py b/gdeep/data/datasets/graph_datasets.py
similarity index 100%
rename from gdeep/data/graph_datasets.py
rename to gdeep/data/datasets/graph_datasets.py
diff --git a/gdeep/data/orbit.py b/gdeep/data/datasets/orbit.py
similarity index 100%
rename from gdeep/data/orbit.py
rename to gdeep/data/datasets/orbit.py
diff --git a/gdeep/data/parallel_orbit.py b/gdeep/data/datasets/parallel_orbit.py
similarity index 100%
rename from gdeep/data/parallel_orbit.py
rename to gdeep/data/datasets/parallel_orbit.py
diff --git a/gdeep/data/tori.py b/gdeep/data/datasets/tori.py
similarity index 75%
rename from gdeep/data/tori.py
rename to gdeep/data/datasets/tori.py
index f37ef37a..74c477ea 100644
--- a/gdeep/data/tori.py
+++ b/gdeep/data/datasets/tori.py
@@ -1,25 +1,30 @@
 import math
 import numpy as np
 import torch
+from typing import Tuple
 from torch.utils.data import Dataset
 from sklearn.datasets import make_blobs
 
+Tensor = torch.Tensor
+Array = np.ndarray
 
 class Rotation():
     """Class for rotations
     """
-    def __init__(self, axis_0, axis_1, angle):
+    def __init__(self, axis_0: int,
+                 axis_1: int,
+                 angle: float) -> None:
         self._axis_0 = axis_0
         self._axis_1 = axis_1
         self._angle = angle
 
-    def return_axis(self, idx):
+    def return_axis(self, idx:int):
         return eval('self._axis_'+str(idx))
 
-    def return_angle(self):
+    def return_angle(self) -> float:
         return self._angle
 
-    def rotation_matrix(self):
+    def rotation_matrix(self) -> Array:
         rotation_matrix = np.identity(3)
         rotation_matrix[self._axis_0, self._axis_0]\
             = math.cos(self._angle)
@@ -32,73 +37,40 @@ def rotation_matrix(self):
         return rotation_matrix
 
 
-class GenericDataset(Dataset):
-    """This class is the base class for the tori-datasets
-
-    Args:
-        data (Tensor):
-            tensor with first dimension
-            the number of samples
-        taregts (list):
-            list of labels
-        transform (Callable):
-            act on the single images
-        target_transform (Callable):
-            act on the single label
-    """
-
-    def __init__(self, data, targets,
-                 transform=None, target_transform=None):
-        self.targets = targets
-        self.data = data
-        self.transform = transform
-        self.target_transform = target_transform
-
-    def __len__(self):
-        return len(self.targets)
-
-    def __getitem__(self, idx):
-        image = self.data[idx]
-        label = self.targets[idx]
-        if self.transform:
-            image = self.transform(image)
-        if self.target_transform:
-            label = self.target_transform(label)
-        sample = [image.to(torch.float), label.to(torch.long)]
-        return sample
-
-
-class CreateToriDataset:
+class ToriDataset(Dataset[Tuple[Tensor, Tensor]]):
     """This class is used to generate data loaders for the
     family of tori-datasets
 
     Args:
-        name (string):
+        name:
             name of the torus dataset to generate
     """
-    def __init__(self, name):
+    def __init__(self, name: str, **kwargs) -> None:
         self.name = name
-
-    def generate_dataset(self, **kwargs):
-        """This method is the maion method of the class.
-        It generates the ToriDatasset class depending on
-        the self.name"""
         if self.name == "DoubleTori":
-            tup = self._make_two_tori_dataset(**kwargs)
-            return GenericDataset(*tup, target_transform=torch.tensor)
+            self.out_tuple = self._make_two_tori_dataset(**kwargs)
         elif self.name == "Blobs":
-            tup = self._make_blobs(**kwargs)
-            return GenericDataset(*tup, target_transform=torch.tensor)
+            self.out_tuple = self._make_blobs(**kwargs)
         else:
-            tup = self._make_entangled_tori_dataset(**self.kwargs)
-            return GenericDataset(*tup, target_transform=torch.tensor)
+            self.out_tuple = self._make_entangled_tori_dataset(**kwargs)
+
+    def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor]:
+        """This method is the main method of the class.
+        It extracts the i-the element
+        """
+
+        return self.out_tuple[0][idx], torch.tensor(self.out_tuple[1][idx]).to(torch.long)
+
+    def __len__(self) -> int:
+        return len(self.out_tuple[0])
 
     @staticmethod
-    def _make_torus_point_cloud(label: int, n_points: int,
+    def _make_torus_point_cloud(label: int, 
+                                n_points: int,
                                 noise: float,
                                 rotation: Rotation,
-                                base_point: np.array,
-                                radius: float = 1.):
+                                base_point: Array,
+                                radius: float=1.) -> Tuple[Array, Array]:
         """Generate point cloud of a single torus
 
         Args:
@@ -144,29 +116,29 @@ def _make_torus_point_cloud(label: int, n_points: int,
 
     @staticmethod
     def _make_torus_point_cloud2(label: int, n_points: int, noise: float,
-                                rotation: Rotation, base_point: np.array,
-                                radius1: float = 1., radius2: float = 1.):
+                                rotation: Rotation, base_point: Array,
+                                radius1: float = 1., radius2: float = 1.) -> Tuple[Array, Array]:
         """Generate point cloud of a single torus using
         2 radii for its definition
 
         Args:
-            label (int):
+            label :
                 label of the data points
-            n_points (int):
+            n_points :
                 number of sample points for each direction
-            noise (float):
+            noise :
                 noise
             rotation:
                 Rotation
-            base_point (np.array):
+            base_point :
                 center of the torus
-            radius1 (float):
+            radius1 :
                 radius of torus 1
-            radius2 (float):
+            radius2 :
                 radius of torus 2
 
         Returns:
-            (np.array, np.array):
+            (np.ndarray, np.ndarray):
                 data_points, labels
         """
         torus_point_clouds = np.asarray(
@@ -194,12 +166,12 @@ def _make_torus_point_cloud2(label: int, n_points: int, noise: float,
         return torus_point_clouds, torus_labels
 
     def _make_two_tori_dataset(self, entangled: bool = True,
-                               n_pts: int = 10) -> tuple:
+                               n_pts: int = 10) -> Tuple[Tensor, Array]:
         """Generates pandas Dataframe of two tori in 3D. The labels correspond to
         the different Tori.
 
         Args:
-            entangled (bool, optional):
+            entangled :
                 Either entangled or unentangled tori. Defaults to True.
 
         Returns:
@@ -239,17 +211,17 @@ def _make_two_tori_dataset(self, entangled: bool = True,
         tori_labels = np.concatenate((torus_labels[0],
                                       torus_labels[1]),
                                      axis=0)
-        return torch.from_numpy(tori_point_cloud), tori_labels
+        return torch.from_numpy(tori_point_cloud).to(torch.float32), tori_labels
 
     def _make_entangled_tori_dataset(self, m: int = 2,
-                                     n_pts: int = 10) -> tuple:
+                                     n_pts: int = 10) -> Tuple[Tensor, Array]:
         """Generates pandas Dataframe of m x m x m tori in 3D.
         The labels correspond to the different Tori.
 
         Args:
-            m (int, defalut=2):
+            m :
                 Number of entangled tori per axis
-            n_pts (int):
+            n_pts :
                 Number of points per torus
 
         Returns:
@@ -275,10 +247,10 @@ def _make_entangled_tori_dataset(self, m: int = 2,
             if t != [0, 0, 0]:
                 data = np.append(data, np.add(tori_entangled, t), axis=0)
                 lab = np.append(lab, labels)
-        return torch.from_numpy(data), lab
+        return torch.from_numpy(data).to(torch.float32), lab
 
     @staticmethod
-    def _make_blobs(m: int = 3, n_pts: int = 200) -> tuple:
+    def _make_blobs(m: int = 3, n_pts: int = 200) -> Tuple[Tensor, Array]:
         """Generates blobs in 3D.
         The labels correspond to the different blob.
 
@@ -294,4 +266,4 @@ def _make_blobs(m: int = 3, n_pts: int = 200) -> tuple:
         """
         data, lab = make_blobs(n_samples=n_pts, centers=m,
                                n_features=3, random_state=42)
-        return torch.from_numpy(data), lab
+        return torch.from_numpy(data).to(torch.float32), lab
diff --git a/gdeep/data/preprocessing_interface.py b/gdeep/data/preprocessing_interface.py
deleted file mode 100644
index f8901814..00000000
--- a/gdeep/data/preprocessing_interface.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import torch
-from abc import ABC, abstractmethod
-import warnings
-import os
-import json
-import jsonpickle
-
-Tensor = torch.Tensor
-
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
-
-
-class AbstractPreprocessing(ABC):
-    """The abstract class to define the interface of preprocessing
-    """
-    @abstractmethod
-    def transform(self, *args, **kwargs):
-        """This method deals with datum-wise transformations. This
-        method is called in the Datasets to transform the output
-        of ``__getitem__``"""
-        pass
-
-    @abstractmethod
-    def fit_to_data(self, *args, **kwargs):
-        """This method deals with getting dataset-level information.
-        """
-        pass
-
-    def save_pretrained(self, path):
-        with open(os.path.join(path, self.__class__.__name__ + ".json"), "w") as outfile:
-            whole_class = jsonpickle.encode(self)
-            json.dump(whole_class, outfile)
-
-    def load_pretrained(self, path):
-        try:
-            with open(os.path.join(path,self.__class__.__name__ + ".json"), "r") as infile:
-                whole_class = json.load(infile)
-                self = jsonpickle.decode(whole_class)
-        except FileNotFoundError:
-            warnings.warn("The transformation file does not exist; attempting to run"
-                          " the transformation anyway...")
diff --git a/gdeep/data/preprocessing_pipeline.py b/gdeep/data/preprocessing_pipeline.py
new file mode 100644
index 00000000..1d0ad48d
--- /dev/null
+++ b/gdeep/data/preprocessing_pipeline.py
@@ -0,0 +1,59 @@
+from typing import Any, Callable, Generic, Iterable, TypeVar
+
+from torch.utils.data import Dataset
+
+from .abstract_preprocessing import AbstractPreprocessing
+from .transforming_dataset import TransformingDataset, append_transform
+
+T = TypeVar('T')
+
+
+class PreprocessingPipeline(AbstractPreprocessing[T, Any], Generic[T]):
+    """ Pipeline to fit non-fitted preprocessors to a dataset in a sequential manner.
+    The fitted preprocessing transform can be attached to a dataset using
+    the ´attach_transform_to_dataset´ method.
+    The intended use case is to fit the preprocessors to the training dataset and
+    then attach the fitted transform to the training, validation and test datasets.
+
+    The transform is only applied to the data and not the labels.
+
+    Examples::
+
+        from gdeep.data.preeprocessors import PreprocessingPipeline, Normalization, \
+            PreprocessImageClassification
+        from gdeep.data.datasets import DatasetImageClassificationFromFiles
+
+        image_dataset = DatasetImageClassificationFromFiles(
+            os.path.join(file_path, "img_data"),
+            os.path.join(file_path, "img_data", "labels.csv"))
+
+        preprocessing_pipeline = PreprocessingPipeline((PreprocessImageClassification((32, 32)),
+                                                        Normalization()))
+        preprocessing_pipeline.fit_to_dataset(image_dataset)  # this will not change the image_dataset
+        preprocessed_dataset = preprocessing_pipeline.attach_transform_to_dataset(image_dataset)
+
+    """
+
+    transform_composition: Callable[[T], Any]
+
+    def __init__(self, preprocessors: Iterable[AbstractPreprocessing[Any, Any]]) -> None:
+        self.preprocessors = preprocessors
+        self.is_fitted = False
+
+    def attach_transform_to_dataset(self, dataset: Dataset[T]) -> TransformingDataset[T, Any]:
+        return TransformingDataset(dataset, self.transform_composition)
+
+    def fit_to_dataset(self, dataset: Dataset[Any]) -> None:
+        self.transform_composition = lambda x: x
+        transformed_dataset = TransformingDataset(dataset, self.transform_composition)
+        for preprocessor in self.preprocessors:
+            preprocessor.fit_to_dataset(transformed_dataset)
+            transformed_dataset = append_transform(transformed_dataset,
+                                                   preprocessor)
+        self.transform_composition = transformed_dataset.transform
+
+    def __call__(self, x: T) -> Any:
+        if not self.is_fitted:
+            raise ValueError("The preprocessing pipeline has not been fitted to a dataset."
+                             " Please call the ´fit_to_dataset´ method first.")
+        return self.transform_composition(x)
\ No newline at end of file
diff --git a/gdeep/data/preprocessing_pipes.py b/gdeep/data/preprocessing_pipes.py
deleted file mode 100644
index a14de6f7..00000000
--- a/gdeep/data/preprocessing_pipes.py
+++ /dev/null
@@ -1,413 +0,0 @@
-import torch
-from typing import Union, Tuple
-from torch.utils.data import DataLoader, Dataset
-from abc import ABC, abstractmethod
-from torchtext.data.utils import get_tokenizer
-from torch.nn.functional import pad
-from collections import Counter
-from torchtext.vocab import Vocab
-import warnings
-import os
-import json
-import jsonpickle
-from .preprocessing_interface import AbstractPreprocessing
-
-# type definition
-Tensor = torch.Tensor
-
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
-
-
-class Normalisation(AbstractPreprocessing):
-    """This class runs the standard normalisation on all the dimensions of
-    the input tensor. For example, in case of images where each item is of
-    shape ``(BS, C, H, W)``, the average will and the standard deviations
-    will be tensors of shape ``(C, H, W)``
-    """
-    is_fitted: bool
-    mean: Tensor
-    stddev: Tensor
-
-    def __init__(self):
-        self.is_fitted = False
-
-    def fit_to_data(self, data: Tensor):
-        self.mean = self._mean(data, 0, False)
-        self.stddev = self._stddev(data, 0, False)
-        self.is_fitted = True
-        self.save_pretrained(".")
-
-    def transform(self, datum: Tensor) -> Tensor:
-        if not self.is_fitted:
-            self.load_pretrained(".")
-        if not all(self.stddev>0):
-            warnings.warn("The standard deviation contains zeros! Adding 1e-7")
-            self.stddev = self.stddev + 1e-7
-        out = (datum - self.mean)/(self.stddev)
-        return out
-
-    def _mean(self, data: Union[Tensor, Dataset], dim:int, keep_dim:bool) -> Tensor:
-        """compute the mean of the whole dataset"""
-        if isinstance(data, Dataset):
-            data=next(iter(DataLoader(data, batch_size=len(data))))[0]
-        return torch.mean(data.float(), dim, keep_dim)
-
-    def _stddev(self, data: Union[Tensor, Dataset], dim:int, keep_dim:bool) -> Tensor:
-        """compute the stddev of the whole dataset"""
-        if isinstance(data, Dataset):
-            data=next(iter(DataLoader(data, batch_size=len(data))))[0]
-        return torch.std(data.float(), dim, keep_dim)
-
-
-class PreprocessingPipeline(AbstractPreprocessing):
-    """class to compose preprocessing classes. The
-    order is very important, and to make sure that the output of
-    one preprocessing is acceptable as input for the
-    following preprocessing
-
-    Args:
-        list_of_preproc_and_datatypes (list):
-            list of tuples in which the the first
-            element if the class instance of the
-            AbstractPreprocessing and the second element
-            is the Dataset class (not the instance!)
-
-    Examples::
-
-        from gdeep.data import PreprocessingPipeline, Normalisation
-        from gdeep.data import PreprocessTextData, TextDataset
-
-        PreprocessingPipeline(((PreprocessTextData(), TextDataset),
-                               (Normalisation(), lambda x, __:x)))
-
-    """
-    def __init__(self, list_of_preproc_and_datatypes:list):
-        self.list_of_cls = list_of_preproc_and_datatypes
-        
-    def fit_to_data(self, data, **kwargs):
-        for (preproc_cls, data_type_class) in self.list_of_cls:
-            preproc_cls.fit_to_data(data, **kwargs)
-            data = data_type_class(data, preproc_cls, **kwargs)
-
-    def transform(self, batch:Tensor) -> Tensor:
-        for (cls, dt) in self.list_of_cls:
-            batch = cls.transform(batch)
-        return batch
-
-    def __len__(self) -> int:
-        return len(self.list_of_cls)
-
-    def __getitem__(self, index: int):
-        return self.list_of_cls[index]
-
-    def __iter__(self):
-        return iter(self.list_of_cls)
-
-    def __repr__(self) -> str:
-        return f'Pipeline({self.list_of_cls})'
-
-    def __add__(self, other):
-        return PreprocessingPipeline(self.list_of_cls + other.list_of_cls)
-
-
-class PreprocessTextData(AbstractPreprocessing):
-    """Preprocessing class. This class is useful to convert the data format
-    ``(label, text)`` into the proper tensor format ``( word_embedding, label)``
-
-    Args:
-        tokenizer (torch Tokenizer):
-            the tokenizer of the source text
-        vocabulary (torch Vocabulary):
-            the vocubulary; it can be built of it can be
-            given.
-
-    """
-    def __init__(self, tokenizer=None,
-                 vocabulary=None):
-        if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
-        else:
-            self.tokenizer = tokenizer
-
-        self.vocabulary = vocabulary
-        self.MAX_LENGTH = 0
-        self.is_fitted = False
-
-    def fit_to_data(self, data):
-        """Method to extract global data, like to length of
-        the sentences to be able to pad.
-
-        Args:
-            data (iterable):
-                the data in the format ``(label, text)``
-        """
-
-        counter = Counter()  # for the text
-        for (label, text) in data:
-            if isinstance(text, tuple) or isinstance(text, list):
-                text = text[0]
-            counter.update(self.tokenizer(text))
-            self.MAX_LENGTH = max(self.MAX_LENGTH, len(self.tokenizer(text)))
-        # self.vocabulary = Vocab(counter, min_freq=1)
-        if self.vocabulary is None:
-            self.vocabulary = Vocab(counter)
-        self.is_fitted = True
-        self.save_pretrained(".")
-
-    def transform(self, datum: tuple) -> Tensor:
-        """This method is applied to each batch and
-        transforms it following the rule below
-
-        Args:
-            datum (tuple):
-                a single datum, being it a tuple
-                with ``(label, text)``
-        """
-        if not self.is_fitted:
-            self.load_pretrained(".")
-        text_pipeline = lambda x: [self.vocabulary[token] for token in
-                                   self.tokenizer(x)]
-
-        pad_item = self.vocabulary["."]
-
-        _text = datum
-        if isinstance(_text, tuple) or isinstance(_text, list):
-            _text = _text[0]
-        processed_text = torch.tensor(text_pipeline(_text),
-                                      dtype=torch.int64).to(DEVICE)
-        # convert to tensors (padded)
-        out = torch.cat([processed_text,
-                   pad_item * torch.ones(self.MAX_LENGTH - processed_text.shape[0]
-                                              ).to(DEVICE)])
-        return out
-
-
-class PreprocessTextLabel(AbstractPreprocessing):
-    def __init__(self, tokenizer=None, **kwargs):
-        pass
-
-    def fit_to_data(self, dataset):
-        pass
-
-    def transform(self, datum: Tensor) -> Tensor:
-        label_pipeline = lambda x: torch.tensor(x, dtype=torch.long) - 1
-
-        _label = datum
-        try:
-            label_pipeline(_label).to(DEVICE)
-        except TypeError:
-            if isinstance(_label, tuple) or isinstance(_label, list):
-                _label = _label[0]
-        out = label_pipeline(_label).to(DEVICE)
-
-        return out
-
-
-class PreprocessTextTranslation(AbstractPreprocessing):
-    """Class to preprocess text dataloaders for translation
-    tasks
-
-        Args:
-            vocabulary (torch Vocabulary):
-                the vocubulary of the source text;
-                it can be built automatically or it can be
-                given.
-            vocabulary_target (torch Vocabulary):
-                the vocubulary of the target text;
-                it can be built automatically or it can be
-                given.
-            tokenizer (torch Tokenizer):
-                the tokenizer of the source text
-            tokenizer_target (torch Tokenizer):
-                the tokenizer of the target text
-
-    Examples::
-
-        from gdeep.data import TorchDataLoader
-        from gdeep.data import TextDatasetTranslation, PreprocessTextTranslation
-
-        dl = TorchDataLoader(name="Multi30k", convert_to_map_dataset=True)
-        dl_tr, dl_ts = dl.build_dataloaders()
-
-        textds = TextDatasetTranslation(dl_tr.dataset,
-            PreprocessTextTranslation(), None)
-
-        """
-
-    def __init__(self, vocabulary=None,
-                 vocabulary_target=None,
-                 tokenizer=None,
-                 tokenizer_target=None):
-
-        if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
-        else:
-            self.tokenizer = tokenizer
-        if tokenizer_target is None:
-            self.tokenizer_target = get_tokenizer('basic_english')
-        else:
-            self.tokenizer_target = tokenizer_target
-        self.vocabulary = vocabulary
-        self.vocabulary_target = vocabulary_target
-        self.MAX_LENGTH = 0
-        self.is_fitted = False
-
-    def fit_to_data(self, data):
-        counter = Counter()  # for the text
-        counter_target = Counter()  # for the text
-        for text in data:
-            counter.update(self.tokenizer(text[0]))
-            counter_target.update(self.tokenizer_target(text[1]))
-            self.MAX_LENGTH = max(self.MAX_LENGTH, len(self.tokenizer(text[0])))
-            self.MAX_LENGTH = max(self.MAX_LENGTH, len(self.tokenizer_target(text[1])))
-        # self.vocabulary = Vocab(counter, min_freq=1)
-        if self.vocabulary is None:
-            self.vocabulary = Vocab(counter)
-        if self.vocabulary_target is None:
-            self.vocabulary_target = Vocab(counter_target)
-        self.is_fitted = True
-        self.save_pretrained(".")
-
-    def transform(self, datum):
-        """This method is applied to each batch and
-        transforms it following the rule below
-
-        Args:
-            datum (torch.tensor):
-                a single datum
-        """
-
-        if not self.is_fitted:
-            self.load_pretrained(".")
-        text_pipeline = lambda x: [self.vocabulary[token] for token in
-                                   self.tokenizer(x)]
-        text_pipeline_target = lambda x: [self.vocabulary_target[token] for token in
-                                   self.tokenizer_target(x)]
-
-        pad_item = self.vocabulary["."]
-        pad_item_target = self.vocabulary_target["."]
-
-        processed_text = torch.tensor(text_pipeline(datum[0]),
-                                      dtype=torch.int64).to(DEVICE)
-        processed_text_target = torch.tensor(text_pipeline(datum[1]),
-                                      dtype=torch.int64).to(DEVICE)
-        # convert to tensors (padded)
-        out = torch.cat([processed_text,
-                         pad_item * torch.ones(self.MAX_LENGTH - processed_text.shape[0]
-                                               ).to(DEVICE)])
-        out_target = torch.cat([processed_text_target,
-                         pad_item_target * torch.ones(self.MAX_LENGTH - processed_text_target.shape[0]
-                                               ).to(DEVICE)])
-        return out, out_target
-
-
-class PreprocessTextQA(AbstractPreprocessing):
-    """Class to preprocess text dataloaders for Q&A
-    tasks
-
-    Args:
-        vocabulary (torch Vocab):
-            the torch vocabulary
-        tokenizer (torch Tokenizer):
-            the tokenizer of the source text
-
-    Examples::
-
-        from gdeep.data import TorchDataLoader
-        from gdeep.data import  TextDatasetQA, PreprocessTextQA, PreprocessTextQATarget
-
-        dl = TorchDataLoader(name="SQuAD2", convert_to_map_dataset=True)
-        dl_tr, dl_ts = dl.build_dataloaders()
-
-        textds = TextDatasetQA(dl_tr_str.dataset,
-                               PreprocessTextQA(),
-                               PreprocessTextQATarget())
-
-    """
-
-    def __init__(self, vocabulary=None,
-                 tokenizer=None):
-        if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
-        else:
-            self.tokenizer = tokenizer
-        self.vocabulary = vocabulary
-        self.MAX_LENGTH = 0
-        self.is_fitted = False
-
-    def fit_to_data(self, data):
-
-        counter = Counter()  # for the text
-        for (context, question, answer, init_position) in data:
-            if isinstance(context, tuple) or isinstance(context, list):
-                context = context[0]
-            counter.update(self.tokenizer(context))
-            self.MAX_LENGTH = max(self.MAX_LENGTH, len(self.tokenizer(context)))
-            if isinstance(question, tuple) or isinstance(question, list):
-                question = question[0]
-            counter.update(self.tokenizer(question))
-            self.MAX_LENGTH = max(self.MAX_LENGTH, len(self.tokenizer(question)))
-            #if isinstance(answer, tuple) or isinstance(answer, list):
-            #    answer = answer[0]
-            #    if isinstance(answer, tuple) or isinstance(answer, list):
-            #        answer = answer[0]
-            #counter.update(self.tokenizer(answer))
-            #self.MAX_LENGTH_ANSWER = max(self.MAX_LENGTH_ANSWER, len(self.tokenizer(answer)))
-        if self.vocabulary is None:
-            self.vocabulary = Vocab(counter)
-        self.pad_item = self.vocabulary["."]
-        self.is_fitted = False
-        self.save_pretrained(".")
-
-    def transform(self, datum:tuple) -> Tuple[Tensor, Tensor]:
-        if not self.is_fitted:
-            self.load_pretrained(".")
-        text_pipeline = lambda x: [self.vocabulary[token] for token in
-                                   self.tokenizer(x)]
-
-        processed_context = torch.tensor(text_pipeline(datum[0]),
-                                      dtype=torch.int64).to(DEVICE)
-        out_context = torch.cat([processed_context,
-                         self.pad_item * torch.ones(self.MAX_LENGTH - processed_context.shape[0]
-                                               ).to(DEVICE)])
-        processed_question = torch.tensor(text_pipeline(datum[1]),
-                                         dtype=torch.int64).to(DEVICE)
-
-        out_question = torch.cat([processed_question,
-                         self.pad_item * torch.ones(self.MAX_LENGTH - processed_question.shape[0]
-                                               ).to(DEVICE)])
-
-        return out_context, out_question
-
-
-class PreprocessTextQATarget(AbstractPreprocessing):
-    """Class to preprocess text dataloaders for Q&A
-    tasks
-
-        Args:
-            tokenizer (torch Tokenizer):
-                the tokenizer of the source text
-
-    """
-
-    def __init__(self, tokenizer=None):
-        if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
-        else:
-            self.tokenizer = tokenizer
-        self.MAX_LENGTH = 0
-        self.is_fitted = False
-
-    def fit_to_data(self, data):
-        pass
-
-    def transform(self, datum):
-
-        pos_init_char = datum[3][0]
-        pos_init = len(self.tokenizer(datum[0][:pos_init_char]))
-        pos_end = pos_init + len(self.tokenizer(datum[2][0]))
-
-        return torch.tensor(pos_init, dtype=torch.long), torch.tensor(pos_end, dtype=torch.long)
\ No newline at end of file
diff --git a/gdeep/data/preprocessors/__init__.py b/gdeep/data/preprocessors/__init__.py
new file mode 100644
index 00000000..5366fb9d
--- /dev/null
+++ b/gdeep/data/preprocessors/__init__.py
@@ -0,0 +1,15 @@
+
+from .normalization import Normalization
+from .tokenizer_translation import TokenizerTranslation
+from .tokenizer_qa import TokenizerQA
+from .to_tensor_image import ToTensorImage
+from .tokenizer_text_classification import TokenizerTextClassification
+
+
+__all__ = [
+    'Normalization',
+    'TokenizerTranslation',
+    'TokenizerQA',
+    'ToTensorImage',
+    'TokenizerTextClassification'
+    ]
diff --git a/gdeep/data/preprocessors/normalization.py b/gdeep/data/preprocessors/normalization.py
new file mode 100644
index 00000000..4c9bc704
--- /dev/null
+++ b/gdeep/data/preprocessors/normalization.py
@@ -0,0 +1,61 @@
+import warnings
+from typing import TypeVar, Tuple, Any
+
+import torch
+from torch.utils.data import Dataset
+
+from ..abstract_preprocessing import AbstractPreprocessing
+from ..transforming_dataset import TransformingDataset
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+Tensor = torch.Tensor
+
+T = TypeVar('T')
+
+class Normalization(AbstractPreprocessing[Tuple[Tensor, T], Tuple[Tensor, T]]):
+    """This class runs the standard normalisation on all the dimensions of
+    the tensors of a dataloader. For example, in case of images where each item is of
+    shape ``(C, H, W)``, the average will and the standard deviations
+    will be tensors of shape ``(C, H, W)``
+    """
+    is_fitted: bool
+    mean: Tensor
+    stddev: Tensor
+
+    def __init__(self):
+        self.is_fitted = False
+
+    def fit_to_dataset(self, dataset: Dataset[Tuple[Tensor, T]]) -> None:
+        self.mean = _compute_mean_of_dataset(dataset)
+        self.stddev = _compute_stddev_of_dataset(dataset, self.mean)
+        self.is_fitted = True
+
+    def __call__(self, item: Tuple[Tensor, T]) -> Tuple[Tensor, T]:
+        if not self.is_fitted:
+            raise RuntimeError("The normalisation is not fitted to any dataset."
+                               " Please call fit_to_dataset() first.")
+        if not torch.all(self.stddev > 0):
+            warnings.warn("The standard deviation contains zeros! Adding 1e-7")
+            self.stddev = self.stddev + 1e-7
+        out = (item[0] - self.mean)/ self.stddev
+        return (out, item[1])
+
+def _compute_mean_of_dataset(dataset: Dataset[Tuple[Tensor, Any]]) -> Tensor:
+    """Compute the mean of the whole dataset"""
+    mean: Tensor = torch.zeros(dataset[0][0].shape, dtype=torch.float64, device=DEVICE)
+    for idx in range(len(dataset)):  # type: ignore
+        if idx == 0:
+            mean += dataset[idx][0]
+        else:
+            mean = (mean * idx + dataset[idx][0]) / (idx + 1)
+    return mean
+
+def _compute_stddev_of_dataset(dataset: Dataset[Tuple[Tensor, Any]], mean: Tensor) -> Tensor:
+    """Compute the stddev of the whole dataset"""
+    def square_diff_from_mean(x: Tuple[Tensor, Any]) -> Tuple[Tensor, Any]:
+        return (x[0] - mean) ** 2, x[1]
+    mean_normalized_dataset = TransformingDataset[Tuple[Tensor, Any], Tuple[Tensor, Any]] \
+        (dataset, square_diff_from_mean)
+    stddev: Tensor = _compute_mean_of_dataset(mean_normalized_dataset)
+    return stddev.sqrt()
diff --git a/gdeep/data/preprocessors/to_tensor_image.py b/gdeep/data/preprocessors/to_tensor_image.py
new file mode 100644
index 00000000..82653863
--- /dev/null
+++ b/gdeep/data/preprocessors/to_tensor_image.py
@@ -0,0 +1,38 @@
+from typing import Any, List, Tuple, Union
+
+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import Resize, ToTensor
+
+from ..abstract_preprocessing import AbstractPreprocessing
+
+# type definition
+Tensor = torch.Tensor
+
+class ToTensorImage(AbstractPreprocessing[Tuple[Any, Any],
+                                          Tuple[Tensor, Tensor]]):
+    """Class to preprocess image files for classification
+      tasks
+
+          Args:
+              size :
+                  Desired output size. If size is a sequence like (h, w),
+                  output size will be matched to this. If size is an int,
+                  smaller edge of the image will be matched to this number.
+                  I.e, if height > width, then image will be rescaled to
+                  ``(size * height / width, size)``.
+
+    """
+
+    is_fitted: bool
+    size: Union[int, List[int]]
+    def __init__(self, size: Union[int, List[int]]) -> None:
+        self.size = size
+        self.is_fitted = True
+
+    def fit_to_dataset(self, dataset: Dataset[Any]) -> None:
+        # nothing to do
+        pass
+
+    def __call__(self, datum: Tuple[Any, Any]) -> Tuple[Tensor, Tensor]:
+        return ToTensor()(Resize(self.size)(datum[0])), torch.tensor(datum[1],dtype=torch.long)  # type: ignore
diff --git a/gdeep/data/preprocessors/tokenizer_qa.py b/gdeep/data/preprocessors/tokenizer_qa.py
new file mode 100644
index 00000000..42e67551
--- /dev/null
+++ b/gdeep/data/preprocessors/tokenizer_qa.py
@@ -0,0 +1,110 @@
+from collections import Counter, OrderedDict
+from typing import Callable, Tuple, List, Optional, Dict
+
+import torch
+from torch.utils.data import Dataset
+from torchtext.vocab import vocab
+from torchtext.data.utils import get_tokenizer  # type: ignore
+from gdeep.utility import DEVICE
+
+from ..abstract_preprocessing import AbstractPreprocessing
+from .._utils import MissingVocabularyError
+# type definition
+Tensor = torch.Tensor
+
+
+
+class TokenizerQA(AbstractPreprocessing[Tuple[str,str,List[str],List[int]],
+                                        Tuple[Tensor, Tensor]]):
+    """Class to preprocess text dataloaders for Q&A
+    tasks. The type of dataset is assumed to be of the
+    form ``(string,string,list[string], list[string])``.
+
+    Args:
+        vocabulary:
+            the torch vocabulary
+        tokenizer :
+            the tokenizer of the source text
+
+    Examples::
+
+        from gdeep.data import TorchDataLoader
+        from gdeep.data import TransformingDataset
+        from gdeep.data.preprocessors import TokenizerQA
+
+        dl = TorchDataLoader(name="SQuAD2", convert_to_map_dataset=True)
+        dl_tr, dl_ts = dl.build_dataloaders()
+
+        textds = TransformingDataset(dl_tr_str.dataset,
+                               TokenizerQA())
+
+    """
+    is_fitted: bool
+    max_length: int
+    vocabulary: Optional[Dict[str, int]]
+    tokenizer: Optional[Callable[[str], List[str]]]
+    counter: Dict[str, int]
+    pad_item: int
+
+    def __init__(self, vocabulary:Optional[Dict[str, int]]=None,
+                 tokenizer:Optional[Callable[[str], List[str]]]=None):
+        if tokenizer is None:
+            self.tokenizer = get_tokenizer('basic_english')
+        else:
+            self.tokenizer = tokenizer
+        self.vocabulary = vocabulary
+        self.max_length = 0
+        self.is_fitted = False
+
+    def fit_to_dataset(self, dataset: Dataset[Tuple[str,str,List[str],List[int]]]) -> None:
+        """Method to fit the vocabulary to the input text"""
+        self.counter = Counter() 
+        self.ordered_dict = OrderedDict()
+        for (context, question, answer, init_position) in dataset:  # type: ignore
+            self.counter.update(self.tokenizer(context))  # type: ignore
+            self.max_length = max(self.max_length, len(self.tokenizer(context)))  # type: ignore
+            self.counter.update(self.tokenizer(question))  # type: ignore
+            self.max_length = max(self.max_length, len(self.tokenizer(question)))  # type: ignore
+
+        if self.vocabulary is None:
+            self.ordered_dict = OrderedDict(sorted(self.counter.items(),
+                                                   key=lambda x: x[1],
+                                                   reverse=True))
+            self.vocabulary = vocab(self.ordered_dict)
+            unk_token = '<unk>'  # type: ignore
+            if unk_token not in self.vocabulary: self.vocabulary.insert_token(unk_token, 0)
+            self.vocabulary.set_default_index(self.vocabulary[unk_token])
+        self.pad_item = 0
+        self.is_fitted = True
+
+    def __call__(self, datum: Tuple[str, str, List[str], List[int]]) -> Tuple[Tensor, Tensor]:
+        """This method implement the transformation once fitted."""
+        #if not self.is_fitted:
+        #    self.load_pretrained(".")
+        if self.vocabulary:
+            text_pipeline: Callable[[str], List[int]] = lambda x: [self.vocabulary[token] for token in # type: ignore
+                                       self.tokenizer(x)]  # type: ignore
+        else:
+            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+
+        processed_context = torch.tensor(text_pipeline(datum[0]),
+                                         dtype=torch.int64).to(DEVICE)
+        out_context = torch.cat([processed_context,
+                                 self.pad_item * torch.ones(self.max_length -
+                                                            processed_context.shape[0]).to(DEVICE)])
+        processed_question = torch.tensor(text_pipeline(datum[1]),
+                                         dtype=torch.int64).to(DEVICE)
+
+        out_question = torch.cat([processed_question,
+                         self.pad_item * torch.ones(self.max_length -
+                                                    processed_question.shape[0]).to(DEVICE)])
+
+        pos_init_char = datum[3][0]
+        pos_init = len(self.tokenizer(datum[0][:pos_init_char]))  # type: ignore
+        pos_end = pos_init + len(self.tokenizer(datum[2][0]))  # type: ignore
+
+        return (torch.stack((out_context, out_question)).to(torch.long),
+                torch.stack((torch.tensor(pos_init), torch.tensor(pos_end))).to(torch.long))
+
+
+
diff --git a/gdeep/data/preprocessors/tokenizer_text_classification.py b/gdeep/data/preprocessors/tokenizer_text_classification.py
new file mode 100644
index 00000000..8e412c7c
--- /dev/null
+++ b/gdeep/data/preprocessors/tokenizer_text_classification.py
@@ -0,0 +1,124 @@
+import json
+import os
+import warnings
+from abc import ABC, abstractmethod
+from collections import Counter, OrderedDict
+from collections.abc import Sequence
+from typing import Callable, Generic, NewType, Tuple, \
+    Union, Any, Optional, List, Dict
+
+import jsonpickle
+import torch
+from torch.nn.functional import pad
+from torch.utils.data import DataLoader, Dataset
+from torchtext.data.utils import get_tokenizer
+from torchtext.vocab import vocab
+from torchvision.transforms import Resize, ToTensor
+
+from gdeep.utility import DEVICE
+from ..abstract_preprocessing import AbstractPreprocessing
+from .._utils import MissingVocabularyError
+
+# type definition
+Tensor = torch.Tensor
+
+
+class TokenizerTextClassification(AbstractPreprocessing[Tuple[Any, str],
+                                                 Tuple[Tensor, Tensor]]):
+    """Preprocessing class. This class is useful to convert the data format
+    ``(label, text)`` into the proper tensor format ``( word_embedding, label)``.
+    The labels shouzld be interagers; f they ar string, they will be converted.
+
+    Args:
+        tokenizer :
+            the tokenizer of the source text
+        vocabulary :
+            the vocubulary; it can be built of it can be
+            given.
+
+    """
+
+
+    max_length:int
+    is_fitted:bool
+    vocabulary: Optional[Dict[str, int]]
+    tokenizer: Optional[Callable[[str], List[str]]]
+    counter: Dict[str, int]
+
+    def __init__(self, tokenizer: Optional[Callable[[str], List[str]]]=None,
+                 vocabulary: Optional[Dict[str, int]]=None):
+        if tokenizer is None:
+            self.tokenizer = get_tokenizer('basic_english')
+        else:
+            self.tokenizer = tokenizer
+
+        self.vocabulary = vocabulary
+        self.max_length = 0
+        self.is_fitted = False
+
+    def fit_to_dataset(self, dataset:Dataset[Tuple[Any, str]]) -> None:
+        """Method to extract global data, like to length of
+        the sentences to be able to pad.
+
+        Args:
+            data (iterable):
+                the data in the format ``(label, text)``
+        """
+
+        self.counter = Counter()  # for the text
+        self.counter_label = Counter()  # for the text
+        self.ordered_dict = OrderedDict()
+        self.list_of_possible_labels: List[Any] = []
+        for (label, text) in dataset:  # type: ignore
+            #if isinstance(text, tuple) or isinstance(text, list):
+            #    text = text[0]
+            self.counter.update(self.tokenizer(text))  # type: ignore
+            self.counter_label.update([label])
+            self.max_length = max(self.max_length, len(self.tokenizer(text)))  # type: ignore
+        # build the vocabulary
+        if not self.vocabulary:
+            self.ordered_dict = OrderedDict(sorted(self.counter.items(),
+                                                   key=lambda x: x[1],
+                                                   reverse=True))
+            self.list_of_possible_labels = list(self.counter_label.keys())
+            self.vocabulary = vocab(self.ordered_dict)
+            unk_token = '<unk>'  # type: ignore
+            if unk_token not in self.vocabulary: self.vocabulary.insert_token(unk_token, 0)
+            self.vocabulary.set_default_index(self.vocabulary[unk_token])
+        self.is_fitted = True
+        #self.save_pretrained(".")
+
+    def __call__(self, datum: Tuple[Any, str]) -> Tuple[Tensor, Tensor]:
+        """This method is applied to each datum and
+        transforms it following the rule below
+
+        Args:
+            datum (tuple):
+                a single datum, being it a tuple
+                with ``(label, text)``
+        """
+        #if not self.is_fitted:
+        #    self.load_pretrained(".")
+        if self.vocabulary:
+            text_pipeline: Callable[[str], List[int]] = lambda x: [self.vocabulary[token] 
+                                                                   for token in self.tokenizer(x)]  # type: ignore
+        else:
+            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+        pad_item: int = 0
+
+        _text = datum[1]
+        #if isinstance(_text, tuple) or isinstance(_text, list):
+        #    _text = _text[1]
+        processed_text = torch.tensor(text_pipeline(_text),
+                                      dtype=torch.long).to(DEVICE)
+        # convert to tensors (padded)
+        out_text = torch.cat([processed_text,
+                   pad_item * torch.ones(self.max_length - processed_text.shape[0]
+                                              ).to(DEVICE)]).to(torch.long)
+        # preprocess labels
+        label_pipeline = lambda x: torch.tensor(x, dtype=torch.long)
+
+        _label = datum[0]
+        out_label = label_pipeline(self.list_of_possible_labels.index(_label)).to(DEVICE)
+        return out_text, out_label
+
diff --git a/gdeep/data/preprocessors/tokenizer_translation.py b/gdeep/data/preprocessors/tokenizer_translation.py
new file mode 100644
index 00000000..77b968ae
--- /dev/null
+++ b/gdeep/data/preprocessors/tokenizer_translation.py
@@ -0,0 +1,152 @@
+
+from collections import Counter, OrderedDict
+from typing import Callable, Tuple, \
+    Optional, List, Dict
+
+import torch
+from torchtext.data.utils import get_tokenizer    # type: ignore
+from torchtext.vocab import vocab
+from gdeep.utility import DEVICE
+
+from ..abstract_preprocessing import AbstractPreprocessing
+from .._utils import MissingVocabularyError
+
+# type definition
+Tensor = torch.Tensor
+
+
+
+class TokenizerTranslation(AbstractPreprocessing[Tuple[str, str], 
+                                                 Tuple[Tensor, Tensor]
+                           ]):
+    """Class to preprocess text dataloaders for translation
+    tasks. The Dataset type is supposed to be ``(string, string)``.
+    The padding item is supposed to be of index 0.
+
+        Args:
+            vocabulary :
+                the vocubulary of the source text;
+                it can be built automatically or it can be
+                given.
+            vocabulary_target :
+                the vocubulary of the target text;
+                it can be built automatically or it can be
+                given.
+            tokenizer:
+                the tokenizer of the source text
+            tokenizer_target:
+                the tokenizer of the target text
+
+    Examples::
+
+        from gdeep.data import DatasetBuilder
+        from gdeep.data import TransformingDataset
+        from gdeep.data.preprocessors import TokenizerTranslation
+
+        db = DatasetBuilder(name="Multi30k", convert_to_map_dataset=True)
+        ds_tr, ds_val, _ = db.build()
+
+        textds = TransformingDataset(ds_tr,
+            TokenizerTranslation())
+
+        """
+    if_fitted: bool
+    vocabulary: Optional[Dict[str, int]]
+    vocabulary_target: Optional[Dict[str, int]]
+    tokenizer: Optional[Callable[[str], List[str]]]
+    tokenizer_target: Optional[Callable[[str], List[str]]]
+    ordered_dict: Dict[str, int]
+    ordered_dict_target: Dict[str, int]
+
+    def __init__(self, vocabulary:Optional[Dict[str, int]]=None,
+                 vocabulary_target:Optional[Dict[str, int]]=None,
+                 tokenizer:Optional[Callable[[str],List[str]]]=None,
+                 tokenizer_target:Optional[Callable[[str],List[str]]]=None):
+
+        if tokenizer is None:
+            self.tokenizer = get_tokenizer('basic_english')
+        else:
+            self.tokenizer = tokenizer
+        if tokenizer_target is None:
+            self.tokenizer_target = get_tokenizer('basic_english')
+        else:
+            self.tokenizer_target = tokenizer_target
+        self.vocabulary = vocabulary
+        self.vocabulary_target = vocabulary_target
+        self.max_length = 0
+        self.is_fitted = False
+
+    def fit_to_dataset(self, data):
+        self.counter = Counter()
+        self.counter_target = Counter()
+        self.ordered_dict = OrderedDict()
+        self.ordered_dict_target = OrderedDict()
+        for text in data:
+            self.counter.update(self.tokenizer(text[0]))
+            self.counter_target.update(self.tokenizer_target(text[1]))
+            self.max_length = max(self.max_length, len(self.tokenizer(text[0])))
+            self.max_length = max(self.max_length, len(self.tokenizer_target(text[1])))
+        # self.vocabulary = Vocab(counter, min_freq=1)
+        if self.vocabulary is None:
+            self.ordered_dict = OrderedDict(sorted(self.counter.items(),
+                                                   key=lambda x: x[1],
+                                                   reverse=True))
+            self.vocabulary = vocab(self.ordered_dict)
+            unk_token = '<unk>'  # type: ignore
+            if unk_token not in self.vocabulary: self.vocabulary.insert_token(unk_token, 0)
+            self.vocabulary.set_default_index(self.vocabulary[unk_token])
+        if self.vocabulary_target is None:
+            self.ordered_dict_target = OrderedDict(sorted(self.counter_target.items(),
+                                                   key=lambda x: x[1],
+                                                   reverse=True))
+            self.vocabulary_target = vocab(self.ordered_dict_target)
+            unk_token = '<unk>'  # type: ignore
+            if unk_token not in self.vocabulary_target: self.vocabulary_target.insert_token(unk_token, 0)
+            self.vocabulary_target.set_default_index(self.vocabulary_target[unk_token])
+            #self.vocabulary_target = Vocab(self.counter_target)
+        self.is_fitted = True
+        #self.save_pretrained(".")
+
+    def __call__(self, datum: Tuple[str, str]) -> Tuple[Tensor, Tensor]:
+        """This method is applied to each item of
+        the dataset and
+        transforms it following the rule described
+        in this method
+
+        Args:
+            datum (torch.tensor):
+                a single datum
+        """
+
+        #if not self.is_fitted:
+        #    self.load_pretrained(".")
+        if self.vocabulary:
+            text_pipeline: Callable[[str], List[int]] = lambda x: [self.vocabulary[token] for token in  # type: ignore
+                                       self.tokenizer(x)]  # type: ignore
+        else:
+            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+
+        if self.vocabulary_target:
+            text_pipeline_target: Callable[[str], List[int]] = lambda x: [self.vocabulary_target[token] for token in  # type: ignore
+                                       self.tokenizer_target(x)]   # type: ignore
+        else:
+            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+
+        pad_item: int = 0
+        pad_item_target: int = 0
+
+        processed_text = torch.tensor(text_pipeline(datum[0]),
+                                      dtype=torch.long).to(DEVICE)
+        processed_text_target = torch.tensor(text_pipeline_target(datum[1]),
+                                      dtype=torch.long).to(DEVICE)
+        # convert to tensors (padded)
+        out_source = torch.cat([processed_text,
+                                pad_item * torch.ones(self.max_length - processed_text.shape[0]
+                                               ).to(DEVICE)]).to(torch.long)
+        out_target = torch.cat([processed_text_target,
+                                pad_item_target * torch.ones(self.max_length - processed_text_target.shape[0]
+                                               ).to(DEVICE)]).to(torch.long)
+        X = torch.stack([out_source.to(torch.long), out_target.to(torch.long)])
+        y = out_target.to(torch.long)  #.clone().detach()
+        return X, y
+
diff --git a/gdeep/data/text_dataset.py b/gdeep/data/text_dataset.py
deleted file mode 100644
index 5c66a933..00000000
--- a/gdeep/data/text_dataset.py
+++ /dev/null
@@ -1,180 +0,0 @@
-import torch
-from torch.utils.data import Dataset, DataLoader
-from typing import Tuple, Optional, Union
-from .preprocessing_interface import AbstractPreprocessing
-
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
-
-Tensor = torch.Tensor
-
-
-class TextDataset(Dataset):
-    """This class is the base class for the text-datasets.
-    The source dataset for this class are expected to be
-    dataset of the form ``(label, string)``.
-
-    Args:
-        dataset (torch Dataset):
-            The source dataset for this class.
-            It is expected to be a
-            dataset of the form ``(label, string)``.
-        transform (AbstractPreprocessing):
-            the instance of the class of preprocessing.
-            It inherits from ``AbstractPreprocessing``
-        target_transform (AbstractPreprocessing):
-            the instance of the class of preprocessing.
-            It inherits from ``AbstractPreprocessing``
-
-    Examples::
-
-        from gdeep.data import TorchDataLoader
-        from gdeep.data import TextDataset, PreprocessTextData, PreprocessTextLabel
-
-        dl = TorchDataLoader(name="AG_NEWS", convert_to_map_dataset=True)
-        dl_tr, dl_ts = dl.build_dataloaders()
-
-        textds = TextDataset(dl_tr.dataset,
-                             PreprocessTextData(),
-                             PreprocessTextLabel())
-
-    """
-
-    def __init__(self, dataset: Dataset,
-                 transform: Optional[AbstractPreprocessing]=None,  # Optional[...]
-                 target_transform: Optional[AbstractPreprocessing]=None):
-        self.dataset = dataset
-        self.transform = transform
-        self.target_transform = target_transform
-        # initialise the preprocessing
-        if self.transform:
-            self.transform.fit_to_data(self.dataset)
-        if self.target_transform:
-            self.target_transform.fit_to_data(self.dataset)
-
-    def __len__(self):
-        return len(self.dataset)
-
-    def __getitem__(self, idx : int) -> Tuple[Tensor, Tensor]:
-        """The output of this method is a tuple of
-        Two tensors. The first one is the tensor
-        of tokenised words (i.e. a tensor of type ``torch.long``
-        that contains the indices of the tokens within
-        the vocabulary)"""
-        text = self.dataset[idx][1]
-        label = self.dataset[idx][0]
-        if self.transform:
-            text = self.transform.transform(text)
-            #if idx < 8: print(text)
-        if self.target_transform:
-            label = self.target_transform.transform(label)
-        if isinstance(label, torch.Tensor):
-            if text is None: 
-                print(text, label, idx)
-            sample = (text.to(torch.long), label.to(torch.long))
-        else:
-            sample = (text.to(torch.long), torch.tensor(label, dtype=torch.long))
-        return sample
-
-class TextDatasetTranslation(TextDataset):
-    """This class is the class for the text datasets
-    dealing with translation tasks. The source data is expected
-    to be of the form ``(string, string)`` containing the
-    senteces to translate (left string translates into right string).
-
-    Args:
-        dataset (torch Dataset):
-            The source dataset for this class.
-            It is expected to be a
-            dataset of the form ``(label, string)``.
-        transform (AbstractPreprocessing):
-            the instance of the class of preprocessing.
-            It inherits from ``AbstractPreprocessing``
-        target_transform (AbstractPreprocessing):
-            the instance of the class of preprocessing.
-            It inherits from ``AbstractPreprocessing``
-
-    Returns:
-        tensor, tensor:
-            the first tensor has three dimensions, ``(BS, 2, MAX_LENGTH)``;
-            this tensor represents the stacking of both the tokenisation
-            of the source and target sentence.
-
-    Examples::
-
-        from gdeep.data import TorchDataLoader
-        from gdeep.data import TextDatasetTranslation, PreprocessTextTranslation
-
-        dl = TorchDataLoader(name="Multi30k", convert_to_map_dataset=True)
-        dl_tr, dl_ts = dl.build_dataloaders()
-
-        textds = TextDatasetTranslation(dl_tr.dataset,
-            PreprocessTextTranslation(), None)
-
-    """
-
-    def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor]:
-        text = self.dataset[idx]
-        if self.transform:
-            text = self.transform.transform(text)
-        # unique tensor containing both tensors of the target and source
-        sample = torch.stack([text[0].to(torch.long), text[1].to(torch.long)])
-        y = text[1].to(torch.long).clone().detach()
-
-        return (sample, y)
-
-
-class TextDatasetQA(TextDataset):
-    """This class is the class for the text datasets
-    dealing with Q&A tasks. The source data is expected
-    to be of the form ``(string, string, list[string], list[int])`` containing the
-    senteces ``(context, question, [answer(s)],
-    [initial position (in characters) of the answer(s)]))``.
-
-    Args:
-        dataset (torch Dataset):
-            The source dataset for this class.
-            It is expected to be a
-            dataset of the form ``(label, string)``.
-        transform (AbstractPreprocessing):
-            the instance of the class of preprocessing.
-            It inherits from ``AbstractPreprocessing``
-        target_transform (AbstractPreprocessing):
-            the instance of the class of preprocessing.
-            It inherits from ``AbstractPreprocessing``
-
-    Returns:
-        tensor, tensor:
-            the first tensor has three dimensions, ``(BS, 2, MAX_LENGTH)``;
-            this tensor represents the stacking of both the tokenisation
-            of the context and question sentence.
-            The second tensor is just the batch of pairs of start and
-            end positions of the answers
-
-    Examples::
-
-        from gdeep.data import TorchDataLoader
-        from gdeep.data import  TextDatasetQA, PreprocessTextQA, PreprocessTextQATarget
-
-        dl = TorchDataLoader(name="SQuAD2", convert_to_map_dataset=True)
-        dl_tr, dl_ts = dl.build_dataloaders()
-
-        textds = TextDatasetQA(dl_tr_str.dataset,
-                               PreprocessTextQA(),
-                               PreprocessTextQATarget())
-
-    """
-
-    def __getitem__(self, idx:int) -> Tuple[Tensor, Tensor]:
-
-        if self.transform:
-            context, question = self.transform.transform(self.dataset[idx])
-        if self.target_transform:
-            pos_init, pos_end = self.target_transform.transform(self.dataset[idx])
-        #sample = [(context.to(torch.long), question.to(torch.long)), (pos, answer.to(torch.long))]
-        sample = [torch.stack((context, question)).to(torch.long),
-                  torch.stack((pos_init, pos_end)).to(torch.long)]
-        return sample
-
diff --git a/gdeep/data/torch_datasets.py b/gdeep/data/torch_datasets.py
deleted file mode 100644
index 0097fe76..00000000
--- a/gdeep/data/torch_datasets.py
+++ /dev/null
@@ -1,457 +0,0 @@
-from typing import Union, Tuple, Dict, Any
-from sympy import false
-import shutil
-
-import torch
-from torch.utils.data import DataLoader, Dataset
-from torchvision import datasets
-from torchtext import datasets as textds
-from torchvision.transforms import ToTensor, Resize
-import warnings
-import pandas as pd
-import os
-from os.path import join
-import json
-from PIL import Image, UnidentifiedImageError
-from tqdm import tqdm
-import numpy as np
-
-from abc import ABC, abstractmethod
-
-from . import CreateToriDataset
-from .dataset_cloud import DatasetCloud
-
-class AbstractDataLoader(ABC):
-    """The abstractr class to interface the
-    Giotto dataloaders"""
-    @abstractmethod
-    def build_dataloaders(self):
-        pass
-
-
-class DatasetNameError(Exception):
-    """Exception for the improper dataset
-    name"""
-    pass
-
-
-class MapDataset(Dataset):
-    """Class to get a MapDataset from
-    an iterable one.
-
-    Args:
-        data_list (list):
-            the list(IterableDataset)
-    """
-    def __init__(self, data_list):
-        self.data_list = data_list
-
-    def __getitem__(self, index):
-        return self.data_list[index]
-
-    def __len__(self):
-        return len(self.data_list)
-
-
-class TorchDataLoader(AbstractDataLoader):
-    """Class to obtain DataLoaders from the classical
-    datasets available on pytorch.
-
-    Args:
-        name (string):
-            check the available datasets at
-            https://pytorch.org/vision/stable/datasets.html
-        convert_to_map_dataset (bool):
-            whether to convert an IterableDataset to a
-            MapDataset
-
-    """
-    def __init__(self, name: str="MNIST", convert_to_map_dataset: bool=False) -> None:
-        self.name = name
-        self.convert_to_map_dataset = convert_to_map_dataset
-
-    def _build_datasets(self) -> None:
-        """Private method to build the dataset from
-        the name of the dataset.
-        """
-        if self.name in ["DoubleTori", "EntangledTori", "Blobs"]:
-            ctd = CreateToriDataset(self.name)
-            self.training_data = ctd.generate_dataset()
-            ctd2 = CreateToriDataset(self.name)
-            self.test_data = ctd2.generate_dataset(n_pts=20)
-        else:
-            try:
-                string_train = 'datasets.' + self.name + '(root="data",' + \
-                    'train=True,download=True,transform=ToTensor())'
-                string_test = 'datasets.' + self.name + '(root="data",' + \
-                    'train=False,download=True,transform=ToTensor())'
-
-                self.training_data = eval(string_train)
-                self.test_data = eval(string_test)
-            except AttributeError:
-                try:
-                    string_data = 'textds.' + self.name + \
-                        '(split=("train","test"))'
-                    self.training_data, self.test_data = eval(string_data)
-                except AttributeError:
-                    raise DatasetNameError(f"The dataset {self.name} is neither in the"
-                                           f" texts nor images datasets")
-                except TypeError:
-                    string_data = 'textds.' + self.name + \
-                                  '(split=("train","dev"))'
-                    self.training_data, self.test_data = eval(string_data)
-
-    def _convert(self):
-        """This private method converts and IterableDataset
-        to a MapDataset"""
-        if isinstance(self.training_data, torch.utils.data.IterableDataset):
-            self.training_data = MapDataset(list(self.training_data))
-
-        if isinstance(self.test_data, torch.utils.data.IterableDataset):
-            self.test_data = MapDataset(list(self.test_data))
-
-    def build_dataloaders(self, **kwargs) -> tuple:
-        """This method is to be called once the class
-        has been initialised with the dataset name
-        
-        Args:
-            kwargs (dict):
-                The keyword arguments for the torch.DataLoader
-                
-        Returns:
-            tuple:
-                The first element is the train DataLoader
-                while the second the test DataLoader
-        """
-        self._build_datasets()
-        if self.convert_to_map_dataset:
-            self._convert()
-        train_dataloader = DataLoader(self.training_data,
-                                      **kwargs)
-        test_dataloader = DataLoader(self.test_data,
-                                     **kwargs)
-        return train_dataloader, test_dataloader
-        
-
-class DataLoaderFromImages(AbstractDataLoader):
-    """This class is useful to build dataloaders
-    from different images you have in a folder
-    
-    Args:
-        training_folder (string):
-            The path to the folder where the training
-            images are located
-        test_folder (string):
-            The path to the folder where the
-            test images are located
-        labels_file (string):
-            The path and file name of the labels.
-            It shall be a csv file with two columns.
-            The first columns contains the name of the
-            image and the second one contains the
-            label value
-    """
-    def __init__(self, training_folder=".",
-                 test_folder=".",
-                 labels_file="labels.csv"):
-        self.training_folder = training_folder
-        self.test_folder = test_folder
-        self.labels_file = labels_file
-        
-    def build_dataloaders(self, size: tuple=(128, 128), **kwargs) -> tuple:
-        """This method builds the dataloader.
-
-        Args:
-            size (tuple):
-                a tuple (h,w) to convert all images
-                to the same size
-            kwargs (dict):
-                additional arguments to pass to the
-                DataLoaders
-
-        Returns:
-            tuple of torch.DataLoader:
-                the tuple with thhe training and
-                test dataloader directly usable in
-                the pipeline class
-        """
-        CWD = os.getcwd()
-        df = pd.read_csv(self.labels_file)
-        os.chdir(self.training_folder)  # fail if not-found
-        tr_data = []
-        list_of_file_names = os.listdir()
-        for name in tqdm(list_of_file_names):
-            try:
-                image = Image.open(name)
-                imageT = ToTensor()(Resize(size)(image))
-                label = df[df[df.columns[0]]==name][df.columns[-1]].to_numpy()[0]
-                tr_data.append((imageT, label))
-                image.close()
-            except:
-                warnings.warn("This file could not be " +
-                              "loaded due to incompatible format: " +
-                              name)
-            
-        os.chdir(CWD)
-        os.chdir(self.test_folder)  # fail if not-found
-        ts_data = []
-        list_of_file_names = os.listdir()
-        for name in list_of_file_names:
-            try:
-                image = Image.open(name)
-            except UnidentifiedImageError:
-                warnings.warn(f"The image {name} canot be loaded. Skipping it.")
-            imageT = ToTensor()(Resize(size)(image))
-            ts_data.append(imageT)
-
-        train_dataloader = DataLoader(tr_data,
-                                      **kwargs)
-        test_dataloader = DataLoader(ts_data,
-                                     **kwargs)
-        os.chdir(CWD)
-        return train_dataloader, test_dataloader
-
-
-class DataLoaderFromArray(AbstractDataLoader):
-    """This class is useful to build dataloaders
-    from a array of X and y. Tensors are also supported.
-
-    Args:
-        X_train (np.array):
-            The training data
-        y_train (np.array):
-            The training labels
-        X_val (np.array):
-            The validation data
-        y_val (np.array):
-            The validation labels
-        X_test (np.array):
-            The test data
-        labels_file (string):
-            The path and file name of the labels.
-            It shall be a csv file with two columns.
-            The first columns contains the name of the
-            image and the second one contains the
-            label value
-    """
-    def __init__(self, X_train: np.ndarray, y_train: np.ndarray, 
-                 X_val: np.ndarray=None, y_val: np.ndarray=None, 
-                 X_test: np.ndarray=None) -> None:
-        self.X_train = X_train
-        if len(y_train.shape) == 1:
-            self.y_train = y_train.reshape(-1, 1)
-        else:
-            self.y_train = y_train
-        self.y_train = y_train
-        if X_val is not None:
-            self.X_val = X_val
-            if len(y_val.shape) == 1:
-                self.y_val = y_val.reshape(-1, 1)
-            else:
-                self.y_val = y_val
-        if X_test is not None:
-            self.X_test = X_test
-
-    @staticmethod
-    def _from_numpy(X):
-        """this is torch.from_numpy() that also allows
-        for tensors"""
-        if isinstance(X, torch.Tensor):
-            return X
-        return torch.from_numpy(X)
-
-    @staticmethod
-    def _long_or_float(y):
-        if isinstance(y, torch.Tensor):
-            return y
-        if isinstance(y, np.float16) or isinstance(y, np.float32) or isinstance(y, np.float64):
-            return torch.tensor(y).float()
-        return torch.tensor(y).long()
-
-
-    def build_dataloaders(self, **kwargs) -> tuple:
-        """This method builds the dataloader.
-
-        Args:
-            kwargs (dict):
-                additional arguments to pass to the
-                DataLoaders
-
-        Returns:
-            tuple of torch.DataLoader:
-                the tuple with the training, validation and
-                test dataloader directly usable in
-                the pipeline class
-        """
-        tr_data = [(DataLoaderFromArray._from_numpy(x).float(),
-                    DataLoaderFromArray._long_or_float(y)) for x, y in zip(self.X_train, self.y_train)]
-        try:
-            val_data = [(DataLoaderFromArray._from_numpy(x).float(),
-                        DataLoaderFromArray._long_or_float(y)) for x, y in zip(self.X_val, self.y_val)]
-            #val_data = [(DataLoaderFromArray._from_numpy(x).float(),
-            #             torch.tensor(y).long() if isinstance(y, np.int64) or
-            #                                       ('__getitem__' in dir(y)
-            #            and (isinstance(y[0], np.int32) or isinstance(y[0], np.int64))) else
-            #torch.tensor(y).float()) for x, y in zip(self.X_val, self.y_val)]
-        except (TypeError, AttributeError):
-            val_data = None
-        try:
-            ts_data = [DataLoaderFromArray._from_numpy(x).float() for x in self.X_test]
-        except (TypeError, AttributeError):
-            ts_data = None
-
-        train_dataloader = DataLoader(tr_data,
-                                      **kwargs)
-        val_dataloader = DataLoader(val_data,
-                                     **kwargs)
-        test_dataloader = DataLoader(ts_data,
-                                     **kwargs)
-        return train_dataloader, val_dataloader, test_dataloader
-    
-    
-    
-class DlBuilderFromDataCloud(AbstractDataLoader):
-    """Class that loads data from Google Cloud Storage
-    
-    This class is useful to build dataloaders from a dataset stored in
-    the GDeep Dataset Cloud on Google Cloud Storage.
-
-    The constructor takes the name of a dataset as a string, and a string
-    for the download directory. The constructor will download the dataset
-    to the download directory. The dataset is downloaded in the version
-    used by Datasets Cloud, which may be different from the version
-    used by the dataset's original developers.
-    
-    Args:
-        dataset_name (str):
-            The name of the dataset.
-        download_dir (str):
-            The directory where the dataset will be downloaded.
-        use_public_access (bool):
-            Whether to use public access. If you want
-            to use the Google Cloud Storage API, you must set this to True.
-            Please make sure you have the appropriate credentials.
-        path_to_credentials (str):
-            Path to the credentials file.
-            Only used if public_access is False and credentials are not
-            provided. Defaults to None.
-        
-
-    Returns:
-        torch.utils.data.DataLoader: The dataloader for the dataset.
-
-    Raises:
-        ValueError:
-            If the dataset_name is not a valid dataset that exists
-            in Datasets Cloud.
-        ValueError:
-            If the download_directory is not a valid directory.
-    """
-    def __init__(self,
-                 dataset_name: str,
-                 download_directory: str,
-                 use_public_access: bool=True,
-                 path_to_credentials: Union[None, str] = None,
-                 ):
-        self.dataset_name = dataset_name
-        self.download_directory = download_directory
-        
-        # Download the dataset if it does not exist
-        self.download_directory
-        
-        self._download_dataset(use_public_access=use_public_access, 
-                               path_to_credentials = path_to_credentials)
-
-        self.dl_builder = None
-        
-        # Load the metadata of the dataset
-        with open(join(self.download_directory, self.dataset_name,
-                       "metadata.json")) as f:
-            self.dataset_metadata = json.load(f)
-            
-        # Load the data and labels of the dataset
-        if self.dataset_metadata['data_type'] == 'tabular':
-            if self.dataset_metadata['data_format'] == 'pytorch_tensor':
-                data = torch.load(join(self.download_directory, 
-                                       self.dataset_name, "data.pt"))
-                labels = torch.load(join(self.download_directory, 
-                                         self.dataset_name, "labels.pt"))
-
-                self.dl_builder = DataLoaderFromArray(data, labels)
-            elif self.dataset_metadata['data_format'] == 'numpy_array':
-                data = np.load(join(self.download_directory,
-                                    self.dataset_name, "data.npy"))
-                labels = np.load(join(self.download_directory,
-                                      self.dataset_name, "labels.npy"))
-                self.dl_builder = DataLoaderFromArray(data, labels)
-            else:
-                raise ValueError("Data format {}"\
-                    .format(self.dataset_metadata['data_format']) +
-                                 "is not yet supported.")
-        else:
-            raise ValueError("Dataset type {} is not yet supported."\
-                .format(self.dataset_metadata['data_type']))
-
-    def _download_dataset(self,
-                          path_to_credentials: Union[None, str] =None,
-                          use_public_access: bool=True,) -> None:
-        """Only download if the download directory does not exist already
-        and if download directory contains at least three files (metadata,
-        data, labels).
-        
-        Args:
-            path_to_credentials (str):
-                Path to the credentials file.
-            use_public_access (bool):
-                Whether to use public access. If you want
-                to use the Google Cloud Storage API, you must set this to True.
-                
-        Returns:
-            None
-        """
-        if (not os.path.isdir(join(self.download_directory, self.dataset_name))
-                            or len(os.listdir(join(self.download_directory,
-                                            self.dataset_name))) < 3):
-            # Delete the download directory if it exists but does not contain
-            # the wanted number of files
-            if (os.path.isdir(join(self.download_directory, self.dataset_name))
-                and
-                len(os.listdir(join(self.download_directory,
-                                   self.dataset_name))) < 3): # type: ignore
-                print("Deleting the download directory because it does "+
-                      "not contain the dataset")
-                shutil.rmtree(self.download_directory, ignore_errors=True)
-                
-            print("Downloading dataset {} to {}"\
-                    .format(self.dataset_name, self.download_directory))
-            dataset_cloud = DatasetCloud(self.dataset_name,
-                                    download_directory=self.download_directory,
-                                    path_to_credentials=path_to_credentials,
-                                    use_public_access = use_public_access,
-                                    )
-            dataset_cloud.download()
-            del dataset_cloud
-        else:
-            print("Dataset '%s' already downloaded" % self.dataset_name)
-
-    def get_metadata(self) -> Dict[str, Any]:
-        """ Returns the metadata of the dataset.
-        
-        Returns:
-            Dict[str, Any]:
-                The metadata of the dataset.
-        """
-        return self.dataset_metadata
-    
-    def build_dataloaders(self, **kwargs)\
-        -> Tuple[DataLoader, DataLoader, DataLoader]:
-        """Builds the dataloaders for the dataset.
-        
-        Args:
-            **kwargs: Arguments for the dataloader builder.
-            
-        Returns:
-            Tuple[DataLoader, DataLoader, DataLoader]:
-                The dataloaders for the dataset (train, validation, test).
-        """
-        return self.dl_builder.build_dataloaders(**kwargs) # type: ignore
\ No newline at end of file
diff --git a/gdeep/data/transforming_dataset.py b/gdeep/data/transforming_dataset.py
new file mode 100644
index 00000000..1c4b676e
--- /dev/null
+++ b/gdeep/data/transforming_dataset.py
@@ -0,0 +1,74 @@
+from typing import Any, Callable, Generic, TypeVar
+
+from torch.utils.data import Dataset
+
+R = TypeVar('R')
+S = TypeVar('S')
+T = TypeVar('T')
+
+
+class TransformingDataset(Dataset[S], Generic[R, S]):
+    """This class is the base class for the all
+    Datasets that need to be transformed via preprocessors.
+    This base class expects to get data from Dataset.
+
+    Args:
+        dataset :
+            The source dataset for this class.
+        transform :
+            This is either a function defined by the
+            users or a fitted preprocessor. The
+            preprocessors inherits from ``AbstractPreprocessing``
+    """
+
+    dataset: Dataset[R]
+    transform: Callable[[R], S]
+
+    def __init__(self,
+                 dataset: Dataset[R],
+                 transform: Callable[[R], S]) -> None:
+        self.dataset = dataset
+        self.transform = transform
+
+    def __len__(self) -> int:
+        if hasattr(self.dataset, '__len__'):
+            return len(self.dataset)  #type: ignore
+        else:
+            raise NotImplementedError("The dataset does not implement the __len__ method.")
+    
+    def __getitem__(self, idx: int) -> S:
+        """The output of this method is one element
+        in the dataset. The type of this element will
+        change accordingly to the dataset itself. For
+        example, in case of a text classification
+        dataset, the output would probably be a tuple
+        like ``(label, string)``
+        """
+        return self.transform(self.dataset[idx])
+
+    
+    # forward all other methods of the TransformingDataset to the Dataset
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.dataset, name)
+
+
+def append_transform(dataset: TransformingDataset[R, S], transform: Callable[[S], T]) \
+    -> TransformingDataset[R, T]:
+    """This function allows to concatenate different preprocessors
+    that the usr may want to stack.
+
+    Args:
+        dataset :
+            A dataset on which it is possible to apply transforms
+        transform :
+            A callable that can be applied to each item in the
+            ``dataset``
+
+    Returns:
+        TransformingDataset:
+            A new TransformingDataset that has been transformed
+            by the transformation
+    """
+    def new_transform(x: R) -> T:
+        return transform(dataset.transform(x))
+    return TransformingDataset(dataset.dataset, new_transform)
\ No newline at end of file
diff --git a/gdeep/search/__init__.py b/gdeep/search/__init__.py
index 2c44f972..05de8526 100644
--- a/gdeep/search/__init__.py
+++ b/gdeep/search/__init__.py
@@ -1,12 +1,12 @@
 
 from .benchmark import Benchmark, _benchmarking_param
-from .gridsearch import Gridsearch, GiottoSummaryWriter
+from .hpo import HyperParameterOptimization, GiottoSummaryWriter
 from .persformer_hyperparamter_search import PersformerHyperparameterSearch
 from ._utils import clean_up_files
 
 __all__ = [
     'Benchmark',
-    'Gridsearch',
+    'HyperParameterOptimization',
     'GiottoSummaryWriter',
     '_benchmarking_param',
     'PersformerHyperparameterSearch',
diff --git a/gdeep/search/benchmark.py b/gdeep/search/benchmark.py
index f1bb6e24..f3fa6be4 100644
--- a/gdeep/search/benchmark.py
+++ b/gdeep/search/benchmark.py
@@ -1,13 +1,20 @@
-from gdeep.pipeline import Pipeline
-from gdeep.utility import _are_compatible
+import warnings
+from typing import Tuple, Any, Dict, Callable, \
+    Type, Optional, List
+
 import torch
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from sklearn.model_selection._split import BaseCrossValidator
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+
+from gdeep.trainer import Trainer
+from gdeep.utility import _are_compatible
 from sklearn.model_selection import KFold
-import warnings
+from gdeep.utility import DEVICE
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
+Tensor = torch.Tensor
 
 class Benchmark:
     """This is the generic class that allows
@@ -34,7 +41,12 @@ class Benchmark:
 
     """
 
-    def __init__(self, models_dicts, dataloaders_dicts, loss_fn, writer, KFold_class=None):
+    def __init__(self,
+                 models_dicts:Dict[str, torch.nn.Module],
+                 dataloaders_dicts:Dict[str, List[DataLoader[Tuple[Tensor, Tensor]]]],
+                 loss_fn: Callable[[Tuple[Tensor, Tensor]], Tensor],
+                 writer: SummaryWriter,
+                 KFold_class:Optional[BaseCrossValidator]=None) -> None:
         self.models_dicts = models_dicts
         self.dataloaders_dicts = dataloaders_dicts
         self.loss_fn = loss_fn
@@ -53,18 +65,19 @@ def __init__(self, models_dicts, dataloaders_dicts, loss_fn, writer, KFold_class
         if not isinstance(self.dataloaders_dicts, list):
             raise TypeError("The provided datasets must be a Python list of dictionaries")
 
-    def start(self, optimizer,
-              n_epochs=10,
-              cross_validation=False,
-              optimizer_param=None,
-              dataloaders_param=None,
-              lr_scheduler=None,
-              scheduler_params=None,
-              profiling=False,
-              parallel_tpu=False,
-              keep_training=False,
-              store_grad_layer_hist=False,
-              n_accumulated_grads=0):
+    def start(self,
+              optimizer: Type[Optimizer],
+              n_epochs:int=10,
+              cross_validation:bool=False,
+              optimizer_param:Optional[Dict[str, Any]]=None,
+              dataloaders_param:Optional[Dict[str, Any]]=None,
+              lr_scheduler: Optional[Type[_LRScheduler]]=None,
+              scheduler_params:Optional[Dict[str, Any]]=None,
+              profiling:bool=False,
+              parallel_tpu:bool=False,
+              keep_training:bool=False,
+              store_grad_layer_hist:bool=False,
+              n_accumulated_grads:int=0) -> None:
         """Method to be called when starting the benchmarking
         
         Args:
@@ -121,20 +134,22 @@ def start(self, optimizer,
                             writer_tag="")
 
 
-    def _inner_function(self, model,
-                        dataloaders,
-                        optimizer, n_epochs,
-                        cross_validation,
-                        optimizer_param,
-                        dataloaders_param,
-                        lr_scheduler,
-                        scheduler_params,
-                        profiling,
-                        parallel_tpu,
-                        keep_training,
-                        store_grad_layer_hist,
-                        n_accumulated_grads,
-                        writer_tag=""):
+    def _inner_function(self,
+                        model:Dict[str,torch.nn.Module],
+                        dataloaders:Dict[str, List[DataLoader[Tuple[Tensor, Tensor]]]],
+                        optimizer: Type[Optimizer],
+                        n_epochs:int,
+                        cross_validation:bool,
+                        optimizer_param:Dict[str, Any],
+                        dataloaders_param:Dict[str, Any],
+                        lr_scheduler:Optional[Type[_LRScheduler]],
+                        scheduler_params:Dict[str, Any],
+                        profiling:bool,
+                        parallel_tpu:bool,
+                        keep_training:bool,
+                        store_grad_layer_hist:bool,
+                        n_accumulated_grads:int,
+                        writer_tag:str=""):
         """private method to run the inner
         function of the benchmark loops
         
@@ -147,7 +162,7 @@ def _inner_function(self, model,
                 and the actual list of dataloaders, e.g.
                 ``[dl_tr, dl_val, dl_ts]``
             optimizer (torch.optim):
-                a torch optimizers
+                a torch optimizer class
             n_epochs (int):
                 number of training epochs
             cross_validation (bool):
@@ -159,7 +174,7 @@ def _inner_function(self, model,
                 dictionary of the dataloaders
                 parameters, e.g. `{"batch_size": 32}`
             lr_scheduler (torch.optim):
-                a learning rate scheduler
+                a learning rate scheduler class
             scheduler_params (dict):
                 learning rate scheduler parameters
             profiling (bool, default=False):
@@ -180,10 +195,11 @@ def _inner_function(self, model,
             writer_tag (str):
                 the tensorboard writer tag
         """
-        pipe = Pipeline(model["model"], dataloaders["dataloaders"],
+        pipe = Trainer(model["model"], dataloaders["dataloaders"],
                         self.loss_fn, self.writer, self.KFold_class)
         writer_tag += "Dataset:" + dataloaders["name"] +"|Model:" + model["name"]
-        pipe.train(optimizer, n_epochs,
+        pipe.train(optimizer,
+                   n_epochs,
                    cross_validation,
                    optimizer_param,
                    dataloaders_param,
@@ -197,7 +213,10 @@ def _inner_function(self, model,
                    n_accumulated_grads,
                    writer_tag)
 
-def _benchmarking_param(fun, arguments, *args, **kwargs):
+def _benchmarking_param(fun: Callable[[Any],Any],
+                        arguments:Tuple[Dict[str, torch.nn.Module],
+                                             Dict[str, DataLoader[Tuple[Tensor, Tensor]]]],
+                        *args, **kwargs):  # type: ignore
     """Function to be used as pseudo-decorator for
     benchmarking loops
     
diff --git a/gdeep/search/gridsearch.py b/gdeep/search/hpo.py
similarity index 94%
rename from gdeep/search/gridsearch.py
rename to gdeep/search/hpo.py
index b9d7b53f..e669e04a 100644
--- a/gdeep/search/gridsearch.py
+++ b/gdeep/search/hpo.py
@@ -1,33 +1,34 @@
+import os
+import time
+import warnings
+from itertools import chain, combinations
+from functools import partial
+from typing import Tuple, Any, Dict, Callable, \
+    Type, Optional, List
+
 import torch
 import optuna
-import os
+from optuna.trial import TrialState
 import pandas as pd
 import numpy as np
-import time
-from gdeep.utility import _are_compatible, _inner_refactor_scalars, KnownWarningSilencer
-from optuna.trial import TrialState
-from gdeep.pipeline import Pipeline
-from gdeep.search import Benchmark, _benchmarking_param
 from sklearn.model_selection import KFold
 from torch.utils.data.sampler import SubsetRandomSampler
-from gdeep.visualisation import plotly2tensor
-from torch.optim import *
-import plotly.express as px
-from functools import partial
-import warnings
-from itertools import chain, combinations
-from optuna.pruners import MedianPruner
-from ..utility import save_model_and_optimizer
 from torch.utils.tensorboard import SummaryWriter
 from torch.utils.tensorboard.summary import hparams
 import random
 import string
+from torch.optim import *
+import plotly.express as px
+from optuna.pruners import MedianPruner
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 
-
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
+from gdeep.utility import _are_compatible, _inner_refactor_scalars, KnownWarningSilencer
+from gdeep.trainer import Trainer
+from gdeep.search import Benchmark, _benchmarking_param
+from gdeep.visualisation import plotly2tensor
+from ..utility import save_model_and_optimizer
+from gdeep.utility import DEVICE
 
 
 class GiottoSummaryWriter(SummaryWriter):
@@ -94,18 +95,19 @@ def add_hparams(self, hparam_dict, metric_dict,
                     w_hp.add_scalar("accuracy", v, t)
 
 
-class Gridsearch(Pipeline):
+class HyperParameterOptimization(Trainer):
     """This is the generic class that allows
     the user to perform gridsearch over several
     parameters such as learning rate, optimizer.
 
     Args:
-        obj (either Pipeline or Benchmark object):
-            either a pipeline of a bechmark class
+        obj (either Trainer or Benchmark object):
+            either a Trainer or a Bechmark class
+            instance
         search_metric (string):
             either ``'loss'`` or ``'accuracy'``
         n_trials (int):
-            number of total gridsearch trials
+            number of total search trials
         best_not_last (bool, default False):
             boolean flag that is ``True`` would use
             the best metric over epochs averaged over the folds in CV
@@ -124,10 +126,14 @@ class Gridsearch(Pipeline):
 
     """
 
-    def __init__(self, obj, search_metric="loss", n_trials=10,
-                 best_not_last=False,
-                 pruner=None, sampler=None,
-                 db_url=None, study_name=None):
+    def __init__(self, obj: Trainer,
+                 search_metric:str="loss",
+                 n_trials:int=10,
+                 best_not_last:bool=False,
+                 pruner=None,
+                 sampler=None,
+                 db_url:str=None,
+                 study_name:str=None):
         self.best_not_last_gs = best_not_last
         self.is_pipe = None
         self.study = None
@@ -137,7 +143,7 @@ def __init__(self, obj, search_metric="loss", n_trials=10,
         self.df_res = None
         self.db_url = db_url
         self.study_name = study_name
-        if (isinstance(obj, Pipeline)):
+        if (isinstance(obj, Trainer)):
             self.pipe = obj
             super().__init__(self.pipe.model,
                              self.pipe.dataloaders,
@@ -183,7 +189,7 @@ def _initialise_new_model(self, models_hyperparam):
                 torch nn.Module
         """
         
-        list_of_params_keys = Gridsearch._powerset(list(models_hyperparam.keys()))
+        list_of_params_keys = HyperParameterOptimization._powerset(list(models_hyperparam.keys()))
         list_of_params_keys.reverse()
         for params_keys in list_of_params_keys:
             sub_models_hyperparam = {k:models_hyperparam[k] for k in models_hyperparam.keys() if k in params_keys}
@@ -271,17 +277,17 @@ def _objective(self, trial,
         optimizer = eval(trial.suggest_categorical("optimizer", optimizers_names))
 
         # generate all the hyperparameters
-        self.optimizers_param = Gridsearch._suggest_params(trial, optimizers_params)
-        self.dataloaders_param = Gridsearch._suggest_params(trial, dataloaders_params)
-        self.models_hyperparam = Gridsearch._suggest_params(trial, models_hyperparams)
-        self.schedulers_param = Gridsearch._suggest_params(trial, schedulers_params)
+        self.optimizers_param = HyperParameterOptimization._suggest_params(trial, optimizers_params)
+        self.dataloaders_param = HyperParameterOptimization._suggest_params(trial, dataloaders_params)
+        self.models_hyperparam = HyperParameterOptimization._suggest_params(trial, models_hyperparams)
+        self.schedulers_param = HyperParameterOptimization._suggest_params(trial, schedulers_params)
         # tag for storing the results
         writer_tag += "/" + str(trial.datetime_start) # str(self.optimizers_param) + \
             #str(self.dataloaders_param) + str(self.models_hyperparam) + \
             #str(self.schedulers_param)
         # create a new model instance
         self.model = self._initialise_new_model(self.models_hyperparam)
-        self.pipe = Pipeline(self.model, self.dataloaders, self.loss_fn,
+        self.pipe = Trainer(self.model, self.dataloaders, self.loss_fn,
                              self.writer, self.KFold_class)
         # set best_not_last
         self.pipe.best_not_last = self.best_not_last_gs
@@ -360,18 +366,18 @@ def _extract_model_and_dataset_name(self, writer_tag):
         return model_name, dataset_name
 
     def start(self,
-              optimizers,
-              n_epochs=1,
-              cross_validation=False,
-              optimizers_params=None,
-              dataloaders_params=None,
-              models_hyperparams=None,
-              lr_scheduler=None,
-              schedulers_params=None,
-              profiling=False,
-              parallel_tpu=False,
-              keep_training=False,
-              store_grad_layer_hist=False,
+              optimizers: List[Optimizer],
+              n_epochs:int=1,
+              cross_validation:bool=False,
+              optimizers_params: Optional[Dict[str, Any]]=None,
+              dataloaders_params: Optional[Dict[str, Any]]=None,
+              models_hyperparams: Optional[Dict[str, Any]]=None,
+              lr_scheduler: Type[_LRScheduler]=None,
+              schedulers_params:Optional[Dict[str, Any]]=None,
+              profiling:bool=False,
+              parallel_tpu:bool=False,
+              keep_training:bool=False,
+              store_grad_layer_hist:bool=False,
               n_accumulated_grads:int=0,
               writer_tag=""):
         """method to be called when starting the gridsearch
@@ -637,7 +643,8 @@ def _store_trial_to_tb(self, trial, scalars_dict_value,
 
     def _results(self, model_name="model", dataset_name="dataset"):
         """This method returns the dataframe with all the results of
-        the gridsearch. It also saves the figures in the writer.
+        the hyperparameters optimisaton.
+        It also saves the figures in the writer.
 
         Args:
             run_name (str):
@@ -691,7 +698,7 @@ def _results(self, model_name="model", dataset_name="dataset"):
                 fig2.show()
                 img2 = plotly2tensor(fig2)
 
-                self.writer.add_images("Gridsearch correlation: " +
+                self.writer.add_images("HPO correlation: " +
                                        model_name + " " + dataset_name,
                                        img2, dataformats="HWC")
                 self.writer.flush()
@@ -788,7 +795,7 @@ def _suggest_params(trial, params):
             if (isinstance(v, list) or isinstance(v, tuple)) and len(v)==1:
                 params[k] = 2*v
         param_temp = {}
-        param_temp = {k:Gridsearch._new_suggest_float(trial, k,*v) for k,v in
+        param_temp = {k:HyperParameterOptimization._new_suggest_float(trial, k,*v) for k,v in
                       params.items() if (isinstance(v, list) or isinstance(v, tuple))
                       and (isinstance(v[0], float) or isinstance(v[1], float))}
         param_temp2 = {k:trial.suggest_int(k,*v) for k,v in
diff --git a/gdeep/search/persformer_hyperparamter_search.py b/gdeep/search/persformer_hyperparamter_search.py
index 18f50d66..a686c8df 100644
--- a/gdeep/search/persformer_hyperparamter_search.py
+++ b/gdeep/search/persformer_hyperparamter_search.py
@@ -11,9 +11,9 @@
 from optuna.pruners import NopPruner
 
 from gdeep.topology_layers import Persformer
-from gdeep.search import GiottoSummaryWriter, Gridsearch
-from gdeep.data import DlBuilderFromDataCloud
-from gdeep.pipeline import Pipeline
+from gdeep.search import GiottoSummaryWriter, HyperParameterOptimization
+from gdeep.data.datasets import DlBuilderFromDataCloud
+from gdeep.trainer import Trainer
 
 class PersformerHyperparameterSearch:
     """This class is used to perform hyperparameter search for Persfomer using 
@@ -64,7 +64,7 @@ def _get_data_loader(self) -> torch.utils.data.DataLoader:
 
         # create the dataset from the downloaded dataset
         train_dataloader, _, _ = dl_cloud_builder\
-            .build_dataloaders(batch_size=10)
+            .build(({"batch_size" : 10},))
         return train_dataloader
     
     
@@ -100,7 +100,7 @@ def search(self) -> None:
             schedulers_params = hyperparameters_dicts.schedulers_params
 
         # Initialize pipeline        
-        pipe = Pipeline(model, [train_dataloader, None], loss_fn, writer,
+        pipe = Trainer(model, [train_dataloader, None], loss_fn, writer,
                             eval(
                                 hyperparameters_dicts.fold_mode + "(" +
                                 str(hyperparameters_dicts.n_splits) + ", "+
@@ -111,11 +111,11 @@ def search(self) -> None:
 
 
         pruner = NopPruner()
-        search = Gridsearch(pipe,
-                            search_metric=hyperparameters_dicts.search_metric,
-                            n_trials=hyperparameters_dicts.n_trials,
-                            best_not_last=True,
-                            pruner=pruner)
+        search = HyperParameterOptimization(pipe,
+                                            search_metric=hyperparameters_dicts.search_metric,
+                                            n_trials=hyperparameters_dicts.n_trials,
+                                            best_not_last=True,
+                                            pruner=pruner)
 
         # starting the hyperparameter search
         search.start([eval(opt) for opt in hyperparameters_dicts.optimizer],
diff --git a/gdeep/pipeline/__init__.py b/gdeep/trainer/__init__.py
similarity index 55%
rename from gdeep/pipeline/__init__.py
rename to gdeep/trainer/__init__.py
index 6b84bde2..1b75a015 100644
--- a/gdeep/pipeline/__init__.py
+++ b/gdeep/trainer/__init__.py
@@ -1,8 +1,8 @@
 
-from .pipeline import Pipeline
+from .trainer import Trainer
 # from .gridsearch import GridSearch
 
 
 __all__ = [
-    'Pipeline'
+    'Trainer'
     ]
diff --git a/gdeep/pipeline/pipeline.py b/gdeep/trainer/trainer.py
similarity index 80%
rename from gdeep/pipeline/pipeline.py
rename to gdeep/trainer/trainer.py
index becf08ab..22148209 100644
--- a/gdeep/pipeline/pipeline.py
+++ b/gdeep/trainer/trainer.py
@@ -1,25 +1,32 @@
-import torch.nn.functional as F
-import torch
-import numpy as np
 import os
 import copy
 import time
 from functools import wraps
-from tqdm import tqdm
 import warnings
+from typing import Tuple, Union, Optional, Callable, \
+    Any, Dict, List, Type
+
+import torch.nn.functional as F
+from torch.optim import Optimizer
+import torch
+from optuna.trial._base import BaseTrial
+import numpy as np
+from tqdm import tqdm
+from sklearn.model_selection._split import BaseCrossValidator
 from sklearn.model_selection import KFold, train_test_split
 from torch.utils.data.sampler import SubsetRandomSampler
-from gdeep.models import ModelExtractor
-from gdeep.utility import _inner_refactor_scalars
-from ..utility.optimisation import MissingClosureError
+from torch.utils.data import DataLoader
 import optuna
 from datetime import datetime
+from torch.utils.tensorboard import SummaryWriter
+from torch.optim.lr_scheduler import _LRScheduler
+
+from ..utility.optimisation import MissingClosureError
+from gdeep.models import ModelExtractor
+from gdeep.utility import _inner_refactor_scalars
+from gdeep.utility import DEVICE
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-    print("Using GPU!")
-else:
-    DEVICE = torch.device("cpu")
+Tensor = torch.Tensor
 
 try:
     import torch_xla
@@ -31,14 +38,14 @@
     print("No TPUs...")
 
 
-def _add_data_to_tb(func):
-    """decorator to store data to tensorboard"""
+def _add_data_to_tb(func: Callable[[Any], Any]) -> Callable[[Any],Any]:
+    """decorator to store PR data to tensorboard"""
     @wraps(func)
     def wrapper(*args, **kwargs):
         pred, val_loss, correct = func(*args, **kwargs)
         try:
             # add data to tensorboard
-            Pipeline._add_pr_curve_tb(torch.vstack(pred), kwargs["class_label"], 
+            Trainer._add_pr_curve_tb(torch.vstack(pred), kwargs["class_label"],
                                       kwargs["class_probs"], kwargs["writer"],
                                       kwargs["writer_tag"] + "/validation")
             try:
@@ -51,21 +58,21 @@ def wrapper(*args, **kwargs):
     return wrapper
     
     
-class Pipeline:
+class Trainer:
     """This is the generic class that allows
     the user to benchhmark models over architectures
     datasets, regularisations, metrics... in one line
     of code.
 
     Args:
-        model (nn.Module):
+        model :
             standard torch model
         dataloaders (list of utils.DataLoader):
             list of standard torch DataLoaders, e.g.
             `[dl_tr, dl_val, dl_ts]`
-        loss_fn (Callables):
+        loss_fn :
             loss function to average over batches
-        writer (tensorboard SummaryWriter):
+        writer :
             tensorboard writer
         KFold_class (sklearn.model_selection, default ``KFold(5, shuffle=True)``):
             the class instance to implement the KFold, can be
@@ -78,8 +85,8 @@ class Pipeline:
         from torch.optim import SGD
         from sklearn.model_selection import StratifiedKFold
         from gdeep.models import FFNet
-        from gdeep.pipeline import Pipeline
-        from gdeep.data import TorchDataLoader
+        from gdeep.trainer import Trainer
+        from gdeep.data import BuildDatasets, BuildDataLoaders
         from gdeep.search import GiottoSummaryWriter
         # model
         class model1(nn.Module):
@@ -91,15 +98,17 @@ def forward(self, x):
 
         model = model1()
         # dataloaders
-        dl = TorchDataLoader(name="DoubleTori")
-        dl_tr, dl_val = dl.build_dataloaders(batch_size=23)
-        dl_ts = None
+        bd = BuildDatasets(name="DoubleTori")
+        ds_tr, ds_val, _ = bd.build_datasets()
+        dl = BuildDataLoaders((ds_tr, ds_val))
+        dl_tr, dl_val, dl_ts = dl.build_dataloaders(batch_size=23)
+
         # loss function
         loss_fn = nn.CrossEntropyLoss()
         # tb writer
         writer = GiottoSummaryWriter()
         # pipeline
-        pipe = Pipeline(model, [dl_tr, dl_val, dl_ts],
+        pipe = Trainer(model, [dl_tr, dl_val, dl_ts],
                         loss_fn, writer,
                         StratifiedKFold(5, shuffle=True))
         # then one needs to train the model using the pipeline!
@@ -109,7 +118,11 @@ def forward(self, x):
 
     # def __init__(self, model, dataloaders, loss_fn, writer,
     # hyperparams_search = False, search_metric = "accuracy", n_trials = 10):
-    def __init__(self, model, dataloaders, loss_fn, writer, KFold_class=None):
+    def __init__(self, model: torch.nn.Module,
+                 dataloaders: List[DataLoader[Tuple[Tensor, Tensor]]],
+                 loss_fn: Callable[[Tensor,Tensor],Tensor],
+                 writer: SummaryWriter,
+                 KFold_class:Optional[BaseCrossValidator]=None) -> None:
         self.model = model
         self.initial_model = copy.deepcopy(self.model)
         assert len(dataloaders) > 0 and len(dataloaders) < 4, "Length of dataloaders must be 1, 2, or 3"
@@ -124,17 +137,17 @@ def __init__(self, model, dataloaders, loss_fn, writer, KFold_class=None):
         if self.writer is None:
             warnings.warn("No writer detected")
         self.DEVICE = DEVICE
-        self.check_has_trained = False
+        self.check_has_trained:bool = False
         self.registered_hook = None
         # used in gradient clipping
-        self.clip = 5
+        self.clip:int = 5
         # used by gridsearch:
-        self.run_name = None
-        self.val_loss_list_hparam = []
-        self.val_acc_list_hparam = []
-        self.best_not_last = False
+        self.run_name:Optional[str] = None
+        self.val_loss_list_hparam:List[float] = []
+        self.val_acc_list_hparam:List[float] = []
+        self.best_not_last:bool = False
         # profiler
-        self.prof = None
+        self.prof:Any = None
 
         if not KFold_class:
             self.KFold_class = KFold(5, shuffle=True)
@@ -142,20 +155,26 @@ def __init__(self, model, dataloaders, loss_fn, writer, KFold_class=None):
             self.KFold_class = KFold_class
 
         
-    def _set_initial_model(self):
+    def _set_initial_model(self) -> None:
         """This private method is used to set
         the initial_model"""
         self.initial_model = copy.deepcopy(self.model)
 
-    def _reset_model(self):
+    def _reset_model(self) -> None:
         """Private method to reset the initial model weights.
         This function is essential for the cross-validation
         procedure.
         """
         self.model = copy.deepcopy(self.initial_model)
 
-    def _optimisation_step(self, dl_tr, steps, loss, 
-                           t_loss, correct, batch, closure):
+    def _optimisation_step(self,
+                           dl_tr:DataLoader[Tuple[Tensor, Tensor]],
+                           steps: int,
+                           loss: Tensor,
+                           t_loss: float,
+                           correct: float,
+                           batch: int,
+                           closure: Callable[[], Tensor]) -> float:
         """Backpropagation"""
         if self.n_accumulated_grads < 2:  # usual case for stochastic gradient descent
             self.optimizer.zero_grad()
@@ -189,12 +208,13 @@ def _optimisation_step(self, dl_tr, steps, loss,
         return t_loss
         
     def _inner_train_loop(self, 
-                            dl_tr, 
-                            writer_tag, 
-                            size, steps,
-                            loss,
-                            t_loss,
-                            correct):
+                          dl_tr: DataLoader[Tuple[Tensor, Tensor]],
+                          writer_tag: str,
+                          size: int,
+                          steps: int,
+                          loss: Tensor,
+                          t_loss: float,
+                          correct: float) -> Tuple[float, float]:
         """Private method to run the loop
         over the batches for the optimisation"""
 
@@ -211,9 +231,9 @@ def closure():
             # Compute prediction and loss
             pred = self.model(X)
             try:
-                correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+                correct += (pred.argmax(1) == y).to(torch.float).sum().item()
             except RuntimeError:
-                correct += (pred.argmax(2) == y).type(torch.float).sum().item()
+                correct += (pred.argmax(2) == y).to(torch.float).sum().item()
             loss = self.loss_fn(pred, y)
             # Save to tensorboard
             try:
@@ -249,7 +269,9 @@ def closure():
         return correct, t_loss
 
 
-    def _train_loop(self, dl_tr, writer_tag=""):
+    def _train_loop(self,
+                    dl_tr: DataLoader[Tuple[Tensor, Tensor]],
+                    writer_tag:str="") -> Tuple[float, float]:
         """private method to run a single training
         loop
         """
@@ -264,7 +286,7 @@ def _train_loop(self, dl_tr, writer_tag=""):
         except AttributeError:
             size = len(dl_tr.dataset)
         steps = len(dl_tr)
-        loss = 0
+        loss = torch.tensor(0)
         correct = 0
         t_loss = 0
         tik = time.time()
@@ -272,7 +294,8 @@ def _train_loop(self, dl_tr, writer_tag=""):
                                                   " accumulated gradients shall be diminished!"
         correct, t_loss = self._inner_train_loop(dl_tr, 
                                                  writer_tag,
-                                                 size, steps,
+                                                 size,
+                                                 steps,
                                                  loss,
                                                  t_loss,
                                                  correct)
@@ -288,7 +311,9 @@ def _train_loop(self, dl_tr, writer_tag=""):
             pass
         return t_loss, correct*100
     
-    def _val_loop(self, dl_val, writer_tag=""):
+    def _val_loop(self,
+                  dl_val: DataLoader[Tuple[Tensor, Tensor]],
+                  writer_tag:str="") -> Tuple[float, float]:
         """private method to run a single validation
         loop
         """
@@ -301,14 +326,14 @@ def _val_loop(self, dl_val, writer_tag=""):
             size = len(dl_val.sampler.indices)
         except AttributeError:
             size = len(dl_val.dataset)
-        val_loss, correct = 0, 0
-        class_label = []
-        class_probs = []
+        val_loss, correct = 0., 0.
+        class_label:List[Tensor] = []
+        class_probs:List[List[Tensor]] = []
         self.model.eval()
 
         pred, val_loss, correct = self._inner_loop(dl=dl_val,
                                                    class_probs=class_probs,
-                                                   class_label= class_label,
+                                                   class_label=class_label,
                                                    loss=val_loss,
                                                    correct=correct,
                                                    writer=self.writer,
@@ -349,7 +374,11 @@ def _val_loop(self, dl_val, writer_tag=""):
         return val_loss, 100*correct
 
     @staticmethod
-    def _add_pr_curve_tb(pred, class_label, class_probs, writer, writer_tag=""):
+    def _add_pr_curve_tb(pred: Tensor,
+                         class_label: List[Tensor],
+                         class_probs: List[List[Tensor]],
+                         writer: SummaryWriter,
+                         writer_tag: str="") -> None:
         """private function to add the PR curve
         to tensorboard"""
         probs = torch.cat([torch.stack(batch) for batch in
@@ -371,8 +400,14 @@ def _add_pr_curve_tb(pred, class_label, class_probs, writer, writer_tag=""):
                 warnings.warn("Cannot store data in the PR curve")
 
     @_add_data_to_tb
-    def _inner_loop(self, dl, class_probs, class_label,
-                    loss, correct, writer, writer_tag):
+    def _inner_loop(self,*,
+                    dl: DataLoader[Tuple[Tensor, Tensor]],
+                    class_probs: List[List[Tensor]],
+                    class_label: List[Tensor],
+                    loss: float,
+                    correct: float,
+                    writer: SummaryWriter,
+                    writer_tag: str) -> Tuple[List[Tensor], float, float]:
         """private function used inside the test
         and validation loops"""
         pred_list = []
@@ -387,14 +422,18 @@ def _inner_loop(self, dl, class_probs, class_label,
                 class_probs.append(class_probs_batch)
                 loss += self.loss_fn(pred, y).item()
                 try:
-                    correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+                    correct += (pred.argmax(1) == y).to(torch.float).sum().item()
                 except RuntimeError:
-                    correct += (pred.argmax(2) == y).type(torch.float).sum().item()
+                    correct += (pred.argmax(2) == y).to(torch.float).sum().item()
                 class_label.append(y)
 
         return pred_list, loss, correct
 
-    def _init_profiler(self, profiling, cross_validation, n_epochs, k_folds):
+    def _init_profiler(self,
+                       profiling:bool,
+                       cross_validation:bool,
+                       n_epochs:int,
+                       k_folds:int) -> None:
         """initialise the profler for profiling"""
         # profiling
         if not cross_validation:
@@ -402,7 +441,7 @@ def _init_profiler(self, profiling, cross_validation, n_epochs, k_folds):
         else:
             active = k_folds*(n_epochs-2)
 
-        active = 10
+        active:int = 10
         if profiling:
             try:
                 self.prof = torch.profiler.profile(
@@ -421,9 +460,13 @@ def _init_profiler(self, profiling, cross_validation, n_epochs, k_folds):
             except AssertionError:
                 pass
     
-    def _init_optimizer_and_scheduler(self, keep_training, cross_validation,
-                                      optimizer, optimizers_param,
-                                      lr_scheduler, scheduler_params):
+    def _init_optimizer_and_scheduler(self,
+                                      keep_training:bool,
+                                      cross_validation:bool,
+                                      optimizer:Type[Optimizer],
+                                      optimizers_param:Dict[str, Any],
+                                      lr_scheduler:Optional[Type[_LRScheduler]]=None,
+                                      scheduler_params:Optional[Dict[str, Any]]=None) -> None:
         """Reset or maintain the LR scheduler and the 
         optimizer depending on the training"""
         if not (self.check_has_trained and keep_training):
@@ -443,18 +486,21 @@ def _init_optimizer_and_scheduler(self, keep_training, cross_validation,
             if lr_scheduler is not None:  # reset scheduler
                 self.scheduler = lr_scheduler(self.optimizer, **scheduler_params)
     
-    def train(self, optimizer, n_epochs=10, cross_validation=False,
-              optimizers_param=None,
-              dataloaders_param=None,
-              lr_scheduler=None,
-              scheduler_params=None,
-              optuna_params=None,
-              profiling=False,
-              parallel_tpu=False,
-              keep_training=False,
-              store_grad_layer_hist=False,
-              n_accumulated_grads:int=0,
-              writer_tag=""):
+    def train(self,
+              optimizer: Type[Optimizer],
+              n_epochs: int=10,
+              cross_validation: bool=False,
+              optimizers_param: Optional[Dict[str, Any]]=None,
+              dataloaders_param: Optional[Dict[str, Any]]=None,
+              lr_scheduler: Optional[Type[_LRScheduler]]=None,
+              scheduler_params: Optional[Dict[str, Any]]=None,
+              optuna_params: Optional[Tuple[BaseTrial, str]]=None,
+              profiling: bool=False,
+              parallel_tpu: bool=False,
+              keep_training: bool=False,
+              store_grad_layer_hist: bool=False,
+              n_accumulated_grads: int=0,
+              writer_tag: str="") -> Tuple[float, float]:
         """Function to run all the training cycles.
 
         Args:
@@ -524,10 +570,10 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
         # dataloaders_param initialisation
         if dataloaders_param is None:
             if self.dataloaders[1] is not None:
-                dataloaders_param_val = Pipeline.copy_dataloader_params(self.dataloaders[1])
+                dataloaders_param_val = Trainer.copy_dataloader_params(self.dataloaders[1])
             else:
-                dataloaders_param_val = Pipeline.copy_dataloader_params(dl_tr)
-            dataloaders_param_tr = Pipeline.copy_dataloader_params(dl_tr)
+                dataloaders_param_val = Trainer.copy_dataloader_params(dl_tr)
+            dataloaders_param_tr = Trainer.copy_dataloader_params(dl_tr)
         else:
             dataloaders_param_val = dataloaders_param.copy()
             dataloaders_param_tr = dataloaders_param.copy()
@@ -570,7 +616,7 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
         # dataloders without crossvalidation
         if len(self.dataloaders) == 3:
             try:
-                val_idx = self.dataloaders[1].sampler.indices
+                val_idx = self.dataloaders[1].sampler.indices  # type: ignore
             except AttributeError:
                 val_idx = list(range(len(self.dataloaders[1].dataset)))
             #print(val_idx)
@@ -579,7 +625,7 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
                                                  **dataloaders_param_val,
                                                  sampler=SubsetRandomSampler(val_idx))
             try:
-                tr_idx = self.dataloaders[0].sampler.indices
+                tr_idx = self.dataloaders[0].sampler.indices  # type: ignore
             except AttributeError:
                 tr_idx = list(range(len(self.dataloaders[0].dataset)))
             #print(tr_idx)
@@ -589,7 +635,7 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
                                                 sampler=SubsetRandomSampler(tr_idx))
         else:
             try:
-                data_idx = self.dataloaders[0].sampler.indices
+                data_idx = self.dataloaders[0].sampler.indices  # type: ignore
             except AttributeError:
                 data_idx = list(range(len(self.dataloaders[0].dataset)))
             #print(data_idx)
@@ -606,7 +652,7 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
         if cross_validation:
             mean_val_loss = []
             mean_val_acc = []
-            valloss, valacc = 0, 0
+            valloss, valacc = 0., 0.
             try:
                 data_idx = self.dataloaders[0].sampler.indices
                 labels_for_split = [self.dataloaders[0].dataset[i][-1] for i in data_idx]
@@ -675,13 +721,13 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
                 valloss, valacc = self._training_loops(n_epochs, dl_tr,
                                                    dl_val, lr_scheduler, self.scheduler,
                                                    check_optuna, search_metric,
-                                                   trial, 0, writer_tag)
+                                                   trial, False, writer_tag)
             else:
                 valloss, valacc = self.parallel_tpu_training_loops(n_epochs, dl_tr,
                                                    dl_val, optimizers_param, lr_scheduler,
                                                    self.scheduler,
                                                    check_optuna, search_metric,
-                                                   trial, 0)
+                                                   trial, False)
 
         try:
             self.writer.flush()
@@ -693,10 +739,17 @@ def train(self, optimizer, n_epochs=10, cross_validation=False,
         # put the mean of the cross_val
         return valloss, valacc
         
-    def _training_loops(self, n_epochs, dl_tr,
-                        dl_val, lr_scheduler, scheduler,
-                        check_optuna, search_metric,
-                        trial, cross_validation=False, writer_tag=""):
+    def _training_loops(self,
+                        n_epochs:int,
+                        dl_tr:DataLoader[Tuple[Tensor, Tensor]],
+                        dl_val:DataLoader[Tuple[Tensor, Tensor]],
+                        lr_scheduler: Optional[Type[_LRScheduler]]=None,
+                        scheduler: Optional[_LRScheduler]=None,
+                        check_optuna:bool=False,
+                        search_metric:Optional[str]=None,
+                        trial:Optional[BaseTrial]=None,
+                        cross_validation:bool=False,
+                        writer_tag:str="") -> Tuple[float, float]:
         """private method to run the trainign loops
         
         Args:
@@ -735,7 +788,7 @@ def _training_loops(self, n_epochs, dl_tr,
                 the validation loss and validation accuracy
         """
 
-        min_valloss, max_valacc = np.inf, 0
+        min_valloss, max_valacc = np.inf, 0.
         for t in range(n_epochs):
             print(f"Epoch {t+1}\n-------------------------------")
             self.val_epoch = t
@@ -773,24 +826,30 @@ def _training_loops(self, n_epochs, dl_tr,
             self.best_val_loss = min(self.best_val_loss, valloss)
             #print(self.optimizer.param_groups[0]["lr"])
             if not lr_scheduler is None:
-                scheduler.step()
+                scheduler.step()  # type: ignore
             # pruning trials
             if check_optuna and not cross_validation:
                 if search_metric == "loss":
-                    trial.report(valloss, t) # + fold*n_epochs)
+                    trial.report(valloss, t)  # type: ignore
                 else:
-                    trial.report(valacc, t) # + fold*n_epochs)
+                    trial.report(valacc, t)  # type: ignore
                 # Handle pruning based on the intermediate value.
-                if trial.should_prune():
+                if trial.should_prune():  # type: ignore
                     raise optuna.exceptions.TrialPruned()
         return valloss, valacc
 
 
-    def parallel_tpu_training_loops(self, n_epochs, dl_tr_old,
-                                    dl_val_old, optimizers_param,
-                                    lr_scheduler, scheduler,
-                                    check_optuna, search_metric,
-                                    trial, cross_validation=False):
+    def parallel_tpu_training_loops(self,
+                                    n_epochs:int,
+                                    dl_tr_old:DataLoader[Tuple[Tensor, Tensor]],
+                                    dl_val_old:DataLoader[Tuple[Tensor, Tensor]],
+                                    optimizers_param: Dict[str, Any],
+                                    lr_scheduler: Optional[Type[_LRScheduler]]=None,
+                                    scheduler: Optional[_LRScheduler]=None,
+                                    check_optuna: bool=False,
+                                    search_metric: Optional[str]=None,
+                                    trial:Optional[BaseTrial]=None,
+                                    cross_validation: bool=False) -> Tuple[float, float]:
         """Experimental function to run all the training cycles
         on colab TPUs in parallel.
         Note: ``cross_validation`` parameter as well as
@@ -805,9 +864,9 @@ def parallel_tpu_training_loops(self, n_epochs, dl_tr_old,
                 validation dataloader
                 parameters, e.g. `{'batch_size': 32}`
             lr_scheduler (torch.optim):
-                a learning rate scheduler
+                a learning rate scheduler class
             scheduler (torch.optim):
-                the actual scheduler
+                the actual scheduler instance
             check_optuna (bool):
                 boolean to store the optuna results of
                 the trial
@@ -883,17 +942,17 @@ def map_fun_custom(index, flags):
                 # train batch loop
                 size = len(dl_tr.dataset)
                 steps = len(dl_tr)
-                loss = 0    # arbitrary starting value to avoid nan loss
-                correct = 0
+                loss = 0.    # arbitrary starting value to avoid nan loss
+                correct = 0.
                 tik = time.time()
                 # for batch, (X, y) in enumerate(self.dataloaders[0]):
                 for batch, (X, y) in enumerate(para_train_loader):
                     # Compute prediction and loss
                     pred = model2(X)
                     try:
-                        correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+                        correct += (pred.argmax(1) == y).to(torch.float).sum().item()
                     except RuntimeError:
-                        correct += (pred.argmax(2) == y).type(torch.float).sum().item()
+                        correct += (pred.argmax(2) == y).to(torch.float).sum().item()
                     loss = self.loss_fn(pred, y)
                     # Save to tensorboard
                     #self.writer.add_scalar("Parallel" + "/Loss/train",
@@ -915,11 +974,11 @@ def map_fun_custom(index, flags):
                 model2.eval()
                 
                 size = len(dl_val.dataset)
-                loss, correct = 0, 0
-                class_label = []
-                class_probs = []
+                loss, correct = 0., 0.
+                class_label:List[Tensor] = []
+                class_probs:List[List[Tensor]] = []
 
-                pred = 0
+                pred = 0.
                 para_valid_loader = pl.ParallelLoader(dl_val,
                                                       [device]).per_device_loader(device)
                 with torch.no_grad():
@@ -931,9 +990,9 @@ def map_fun_custom(index, flags):
                         class_probs.append(class_probs_batch)
                         loss += self.loss_fn(pred, y).item()
                         try:
-                            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+                            correct += (pred.argmax(1) == y).to(torch.float).sum().item()
                         except RuntimeError:
-                            correct += (pred.argmax(2) == y).type(torch.float).sum().item()
+                            correct += (pred.argmax(2) == y).to(torch.float).sum().item()
                         class_label.append(y)
                     # add data to tensorboard
                     #self._add_pr_curve_tb(pred, class_label, class_probs, "validation")
@@ -960,13 +1019,16 @@ def map_fun_custom(index, flags):
             self.val_loss += loss
             self.val_acc += correct*100
 
-        flags = {}
+        flags:Dict[Any, Any] = {}
         self.val_acc /= len(dl_val_old)*dl_val_old.batch_size
         self.val_loss /= len(dl_val_old)
         xmp.spawn(map_fun_custom, args=(flags,), nprocs=8, start_method='fork')
         return self.val_loss, self.val_acc
 
-    def evaluate_classification(self, num_class=None, dl=None):
+    def evaluate_classification(self,
+                                num_class:int=None,
+                                dl:DataLoader[Tuple[Tensor, Tensor]]=None
+                                ) -> Tuple[float, float, np.ndarray]:
         """Method to evaluate the performance of the model.
         
         Args:
@@ -982,10 +1044,10 @@ def evaluate_classification(self, num_class=None, dl=None):
         """
         if dl is None:
             dl = self.dataloaders[0]
-        class_probs = []
-        class_label = []
-        loss = 0
-        correct = 0
+        class_probs:List[List[Tensor]] = []
+        class_label:List[Tensor] = []
+        loss = 0.
+        correct = 0.
         confusion_matrix = np.zeros((num_class, num_class))
         self.model.eval()
         with torch.no_grad():
@@ -998,7 +1060,7 @@ def evaluate_classification(self, num_class=None, dl=None):
                 class_probs.append(class_probs_batch)
                 loss += self.loss_fn(pred, y).item()
                 correct += (pred.argmax(1) ==
-                            y).type(torch.float).sum().item()
+                            y).to(torch.float).sum().item()
                 class_label.append(y)
                 for t, p in zip(y.view(-1), pred.argmax(1).view(-1)):
                     confusion_matrix[t.long(), p.long()] += 1
@@ -1007,12 +1069,16 @@ def evaluate_classification(self, num_class=None, dl=None):
         return 100*correct, loss, confusion_matrix
 
 
-    def register_pipe_hook(self, callable):
+    def register_pipe_hook(self,
+                           callable:Callable[[Optimizer, int,
+                                              ModelExtractor,
+                                              SummaryWriter],
+                                             Any]) -> None:
         """This method registers a function that
         will be called after each trainign step.
 
         The arguments of the callable function are, in this order:
-         - current optimizer (torch.optim)
+         - current optimizer instance (torch.optim)
          - current epoch number (int)
          - the ModelExtractor instance at that epoch (ModelExtractor)
          - the tensorboard writer (writer)
@@ -1023,14 +1089,18 @@ def register_pipe_hook(self, callable):
         self.registered_hook = callable
 
 
-    def _run_pipe_hook(self, epoch, optim, me, writer):
+    def _run_pipe_hook(self,
+                       epoch:int,
+                       optim:Optimizer,
+                       me:ModelExtractor,
+                       writer: SummaryWriter) -> None:
         """private method that runs the hooked
         function at every epoch, after the single training loop"""
         if self.registered_hook is not None:
             self.registered_hook(epoch, optim, me, writer)
             
     @staticmethod
-    def copy_dataloader_params(original_dataloader):
+    def copy_dataloader_params(original_dataloader: DataLoader[Tuple[Tensor, Tensor]]) -> Dict[str, Any]:
         """returns the dict of init parameters"""
         return {"batch_size": original_dataloader.batch_size,
                 #"batch_sampler": original_dataloader.batch_sampler, 
diff --git a/gdeep/utility/__init__.py b/gdeep/utility/__init__.py
index fbaf1159..4228549b 100644
--- a/gdeep/utility/__init__.py
+++ b/gdeep/utility/__init__.py
@@ -1,6 +1,12 @@
 from .utils import _are_compatible, save_model_and_optimizer, \
     ensemble_wrapper, _inner_refactor_scalars, is_notebook, \
-    KnownWarningSilencer
+        autoreload_if_notebook, KnownWarningSilencer
+from .constants import ROOT_DIR, DEFAULT_DATA_DIR, DATASET_BUCKET_NAME, \
+    DEFAULT_DOWNLOAD_DIR, DATASET_BUCKET_NAME, DEFAULT_GRAPH_DIR, \
+    DEVICE
+
+from ._typing_utils import torch_transform, get_parameter_types, \
+    get_return_type
 
 __all__ = [
     '_are_compatible',
@@ -10,5 +16,16 @@
     'intersection_homology',
     '_inner_refactor_scalars',  # This should be here
     'is_notebook',
+    'autoreload_if_notebook',
+    'ROOT_DIR',
+    'DEVICE',
+    'DEFAULT_DATA_DIR',
+    'DATASET_BUCKET_NAME',
+    'DEFAULT_DOWNLOAD_DIR',
+    'DATASET_BUCKET_NAME',
+    'DEFAULT_GRAPH_DIR',
+    'get_return_type',
+    'torch_transform',
+    'get_parameter_types',
     'KnownWarningSilencer'
-    ]
+    ]
\ No newline at end of file
diff --git a/gdeep/utility/_typing_utils.py b/gdeep/utility/_typing_utils.py
new file mode 100644
index 00000000..2e52f953
--- /dev/null
+++ b/gdeep/utility/_typing_utils.py
@@ -0,0 +1,52 @@
+from inspect import signature
+from typing import Any, Callable, List, Tuple, Union
+
+import torch
+from numpy import ndarray
+
+Tensor = torch.Tensor
+
+def torch_transform(transform: Union[Callable[[Tensor], Tensor],
+                                     Callable[[ndarray], ndarray]]) \
+                                         -> Callable[[Tensor], Tensor]:
+    """ Transforms a numpy array transform to a torch transform. 
+    
+    Args:
+        transform: Either a callable that takes a numpy array and returns a
+            numpy array or a callable that takes a tensor and returns a tensor.
+        
+    Returns:
+        The torch transform.
+    """
+    if get_parameter_types(transform)[0] is Tensor and\
+        get_return_type(transform) is Tensor:
+        return transform  # type: ignore
+    elif get_parameter_types(transform)[0] is ndarray and\
+        get_return_type(transform) is ndarray:
+        return lambda x: torch.tensor(transform(x.numpy()))
+    else:
+        raise ValueError("Transform must be a function that takes a tensor or"
+                            "an array and returns a tensor or an array."
+                            "Please provide type annotations.")
+
+        
+def get_parameter_types(func: Callable) -> List[Any]:
+    """ Returns a list of the types of the parameters of a function. 
+    
+    Args:
+        func: The function to get the types of the parameters of.
+    Returns:
+        A list of the types of the parameters of the function.
+    """
+    return [t.annotation for t in signature(func).parameters.values()]
+    
+def get_return_type(func: Callable) -> Any:
+    """ Returns the type of the return value of a function.
+     
+    Args:
+        func: The function to get the type of the return value of.
+        
+    Returns:
+        The type of the return value of the function.
+    """
+    return signature(func).return_annotation
\ No newline at end of file
diff --git a/gdeep/utility/constants.py b/gdeep/utility/constants.py
index ba18efc0..3fcb3acd 100644
--- a/gdeep/utility/constants.py
+++ b/gdeep/utility/constants.py
@@ -1,10 +1,14 @@
 import os
+import torch
 
 # Define the root directory of the project which is parent of the parent of 
 # the current directory
 ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                         os.pardir, os.pardir))
 
+# Define the default data directory
+DEFAULT_DATA_DIR = os.path.join(ROOT_DIR, 'examples', 'data')
+
 # Define the default dataset download directory
 DEFAULT_DOWNLOAD_DIR = os.path.join(ROOT_DIR, "examples", "data",
                                     "DatasetCloud")
@@ -16,4 +20,7 @@
 
 # Define the default dataset download directory where the graph
 # datasets from the PyG (PyTorch Geometric) library are stored
-DEFAULT_GRAPH_DIR = os.path.join(ROOT_DIR, "example", "data", "graph_datasets")
+DEFAULT_GRAPH_DIR = os.path.join(ROOT_DIR, "examples", "data", "GraphDatasets")
+
+# Define the default deevce: "cuda" is available, "cpu" otherwise
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
\ No newline at end of file
diff --git a/gdeep/utility/intersection_homology/__init__.py b/gdeep/utility/intersection_homology/__init__.py
deleted file mode 100644
index df3b5c4f..00000000
--- a/gdeep/utility/intersection_homology/__init__.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from .intersection_homology import compute_coarsest_strata, \
-    IntersectionHomology
-
-
-__all__=[
-    'IntersectionHomology',
-    'compute_coarsest_strata'
-]
diff --git a/gdeep/utility/intersection_homology/intersection_homology.py b/gdeep/utility/intersection_homology/intersection_homology.py
deleted file mode 100644
index 55f10405..00000000
--- a/gdeep/utility/intersection_homology/intersection_homology.py
+++ /dev/null
@@ -1,386 +0,0 @@
-import numpy as np
-import plotly.express as px
-import networkx as nx
-from networkx.algorithms.clique import find_cliques, enumerate_all_cliques
-from sklearn.neighbors import kneighbors_graph
-from scipy.sparse import csr_matrix
-
-from networkx.drawing.nx_pylab import draw_kamada_kawai
-import warnings
-from networkx.algorithms.components import number_connected_components
-from itertools import combinations, chain
-from sklearn.base import BaseEstimator, TransformerMixin
-
-
-class IntersectionHomology(BaseEstimator, TransformerMixin):
-    """This class computes the Intersection Homology of a clque complex.
-    After initialising the class, a stratification is computed (during `fit`).
-    Once the stratification is established, the Intersection homology
-    can be computed (during `transform`).
-
-    Args:
-        - knn (int) : number of nearest nneigbours
-
-    """
-    def __init__(self, knn: int = 5, perversity=None,
-                 perversity_check=True, coarse_strata=False):
-        self.knn = knn
-        self.gg = None
-        if perversity is None:
-            self.perversity = self.p
-        else:
-            self.perversity = perversity
-        self.check_is_fitted = False
-        self.perversity_check = perversity_check
-        self.list_proper = None
-        self.coarse_strata = coarse_strata
-
-    def fit(self, X):
-        ''' The fit functions accepts as input a pointcloud as a
-        2darray or a networkx graph.
-
-        Args:
-            X (2darray OR networkx Graph object):
-            this is the inout simplicial complex or the point cloud
-            from which to compute the knn graph
-
-        '''
-        self.check_is_fitted = True
-        if type(X) == np.ndarray or type(X) == list:
-            self.gg = self.from_pts_to_nn_graph(X, self.knn)
-        else:
-            self.gg = X
-        self.simplices, self.maximal_cliques = \
-            self.from_graph_to_complex(self.gg)
-        if self.perversity_check:
-            self.list_proper = self.perv_list(self.perversity)
-        return self
-
-    def transform(self, X, phi_list_proper=None):
-        if self.check_is_fitted:
-            if phi_list_proper is None:
-                if self.perversity_check:
-                    return self.phi_generators(self.simplices,
-                                               self.list_proper)
-                return self.generators(self.simplices)
-            else:
-                return self.phi_generators(self.simplices,
-                                           phi_list_proper)
-        warnings.warn("You need to fit before tranforming!")
-
-    def plot(self):
-        if self.check_is_fitted:
-            draw_kamada_kawai(self.gg, with_labels=True)
-
-    @staticmethod
-    def _reduce(D1, verbose: bool = False):
-        '''Input: incidence matrix (2d array)
-        Output: same dimensioal matrix but reduced
-        '''
-        V_fin = np.eye(D1.shape[1])
-        V_temp = np.eye(D1.shape[1])
-        for j in range(D1.shape[1]):
-            j1 = 0
-            while j1 < j:
-                # print(j,j1)
-                if IntersectionHomology._low_(D1, j1) == \
-                    IntersectionHomology._low_(D1, j) and \
-                        IntersectionHomology._low_(D1, j) != 0:
-
-                    D1.T[j] = (D1.T[j] + D1.T[j1]) % 2
-                    V_temp[j1, j] = 1
-                    V_fin = V_fin.dot(V_temp).copy() % 2
-                    V_temp[j1, j] = 0
-                    if verbose:
-                        print("summing line ", j1, " to ", j)
-                        px.imshow(D1).show()
-                    j1 = -1
-                j1 += 1
-        return D1, V_fin
-
-    def _compute_homology(self, D1, proper_list=None):
-        '''returns the index of the unpaired simplices. i.e. those
-        that generate the homology group.'''
-        threshold = D1.shape[0]
-        if proper_list is not None:
-            threshold = len(proper_list)
-        # simplicial pairs
-        D1, _ = self._reduce(D1)
-        index_positive_simplices = []
-        index_pairings = []
-        for j in range(D1.shape[1]):
-            i = self._low_(D1, j)
-            if i == 0:
-                index_positive_simplices.append(j)
-            else:
-                if i < threshold:
-                    index_pairings.append([i, j])
-
-        # check if positive is in pairings
-        homology_index = []
-        for index in index_positive_simplices:
-            if not np.array([True if index in pair else False for
-                             pair in index_pairings]).any():
-                homology_index.append(index)
-        if proper_list is not None:
-            return np.array(proper_list)[homology_index]
-        return homology_index
-
-    @staticmethod
-    def _low_(M, j):
-        column = M.T[j]
-        try:
-            return np.max(np.nonzero(column))
-        except:
-            return 0
-
-    def phi(self, D, lista: list):
-        '''inputs the incidence matrix and outputs the m x s
-        incidence matrix,
-        with the rows ordered (first proper, then improper)
-        '''
-        improper = list(set(range(D.shape[1]))-set(lista))
-        proper = lista
-        # print(improper, proper)
-        return ((D.T[proper]).T)[proper + improper], proper
-
-    def from_pts_to_nn_graph(self, X, knn):
-        '''get the knn graph of the point cloud'''
-        A = kneighbors_graph(X, knn, include_self=False)
-        Mat = A.toarray()
-        G = nx.convert_matrix.from_numpy_matrix(Mat)
-        return G
-
-    def from_graph_to_complex(self, G):
-        '''Get the list of simplices (cliques) fro
-        the input graph.
-        Input: networkx graph G
-        Output: (list of simplices , list of maximal_simplices)'''
-        simplices = list(enumerate_all_cliques(G))
-        max_cliques = list(map(list,
-                           list(map(set, list(find_cliques(G))))))
-        map(list.sort, max_cliques)
-        # simplices_dupl = list(map(list,np.array(G.nodes).reshape(-1,1))) + \
-        #     list(map(list,list(G.edges))) + cliques
-        # simplices = []
-        # [list.sort(s) for s in simplices_dupl]
-        # [simplices.append(n) for n in simplices_dupl if n not in simplices]
-        # simplices.sort(key=len)
-        return simplices, max_cliques
-
-    @staticmethod
-    def is_sublist(sub_list, test_list):
-        if (all(x in test_list for x in sub_list)):
-            return True
-        return False
-
-    def _boundary_matrix(self, simplices):
-        '''build the incidence matrix from the list of simplices.
-        If there are `n` simplices, the matrx is `n * n`. '''
-        row = []
-        col = []
-        data = []
-        length = len(simplices)
-        for i in range(length):
-            for j in range(length):
-                # print(i,j)
-                # print(simplices[i],simplices[j])
-                if len(simplices[i]) == len(simplices[j]) - 1:
-                    if self.is_sublist(simplices[i], simplices[j]):
-                        row.append(i)
-                        col.append(j)
-                        data.append(1)
-
-        boundary_matrix = csr_matrix((data, (row, col)),
-                                     shape=(length,
-                                            length)).toarray()
-        return boundary_matrix
-
-    def homology_array(self, simplices):
-        '''outputs the ordered list of betti numbers,
-        e.g. [0,0,0,1,1,2,2,3] for b_0=3, b_1=2, b_2=2, b_3=1'''
-        bdry_mat = self._boundary_matrix(simplices)
-        comp_hom = self._compute_homology(bdry_mat)
-        homology_generators = np.array(list(map(len,
-                                                simplices)))[comp_hom] - 1
-        return homology_generators, self._reduce(bdry_mat)[1].T[comp_hom]
-
-    def phi_homology_array(self, simplices, list_proper):
-        '''outputs the ordered list of phi-betti numbers,
-        e.g. [0,0,0,1,1,2,2,3] for b_0=3, b_1=2, b_2=2, b_3=1'''
-        D = self._boundary_matrix(simplices)
-        comp_hom = self._compute_homology(*self.phi(D, list_proper))
-        homology_generators = np.array(list(map(len,
-                                                simplices)))[comp_hom] - 1
-        return homology_generators, self._reduce(D)[1].T[comp_hom]
-
-    def generators(self, simplices):
-        '''Input smplices and prints the genetarors and the
-        homology degree.'''
-        arr, vv = self.homology_array(simplices)
-        output = []
-        for index in range(len(arr)):
-            output.append((arr[index],
-                           list(np.array(simplices,
-                                         dtype=object)[(np.nonzero(vv[index])[0]
-                                                       )])))
-        return output
-
-    def phi_generators(self, simplices, list_proper):
-        '''Input smplices and the index list of proper ones,
-        Output: prints the genetarors and the
-        homology degree.'''
-        arr, vv = self.phi_homology_array(simplices, list_proper)
-        output = []
-        for index in range(len(arr)):
-            output.append((arr[index],
-                           list(np.array(simplices,
-                                         dtype=object)[(np.nonzero(vv[index])[0]
-                                                       )])))
-        return output
-
-    # top perversity
-    @staticmethod
-    def p(k):
-        if k < 1:
-            return 0
-        return k - 2
-
-    # zero perversity
-    @staticmethod
-    def q(k):
-        return 0
-
-    def check_stratum(self, sigma, maximal_cliques):
-        '''Define a stratum to be a set of simplices such that each simplex
-            in one stratum is:
-            1. it is a maximal d-dimensional clique
-            2. OR it is a d-dimensional simplex such that it belogs to
-                3 or more max simplices that
-                are not otherwise connected. If they are
-                otherwise connected, the graph of cthe max
-                simplices connectivity is computed and analysed.
-                If there are more than one connected components,
-                then we have a d-dimensional piece of a stratum
-            3. OR it belongs to exactly 2 max cliques of dimension
-                larger than d+2
-        '''
-        cond1 = sigma in maximal_cliques
-        num_inter = np.sum([set(sigma).issubset(set(max_spx))
-                            for max_spx in maximal_cliques])
-        cond2 = False
-        if num_inter > 2:
-            intersecting_max_spxs = [max_spx for max_spx in maximal_cliques
-                                     if set(sigma).issubset(set(max_spx))]
-            G_new = nx.Graph()
-            for s in intersecting_max_spxs:
-                proto_edges = np.nonzero([len(set(s).intersection(ss)) >
-                                          len(sigma) for ss
-                                          in intersecting_max_spxs])[0]
-                # print("proto_edges", proto_edges)
-                edges_of_G = list(combinations(proto_edges, 2))
-                # print("edges_of_G", edges_of_G)
-                G_new.add_edges_from(edges_of_G)
-            num = number_connected_components(G_new)
-            cond2 = (num > 1)
-        cond3 = False
-        if num_inter == 2:
-            cond3 = np.sum([set(sigma).issubset(set(max_spx)) *
-                           (len(max_spx) - len(sigma)) > 1 for max_spx
-                           in maximal_cliques]) == 2
-        return cond1 or cond2 or cond3
-
-    def perv_list(self, perversity):
-        '''This function produces the list of
-        proper simplices, to be
-        inputted directly to the `fit` method.
-        '''
-        n=max(map(len, self.simplices))
-        # there is the option to get the coarse
-        # stratification with the next line.
-        if self.coarse_strata:
-            strata = compute_coarsest_strata(self.simplices)
-        perverse_list = []
-        for i, sigma in enumerate(self.simplices):
-            check = 0
-            for k in range(n):
-
-                stratum = [s for s in self.simplices if len(s) == n-k
-                           and self.check_stratum(s, self.maximal_cliques)]
-                # print(stratum, sigma)
-                if self.coarse_strata:
-                    stratum = strata[n-k-1]
-                else:
-                    stratum = [s for s in self.simplices if len(s) == n-k
-                               and self.check_stratum(s,
-                                                      self.maximal_cliques)]
-                # print(stratum, sigma)
-                if len(stratum)>0:
-                    intersection_temp = \
-                        [set(sigma).intersection(set(simplex_stratum))
-                         for simplex_stratum in stratum]
-                    top_dim = max(map(len, intersection_temp)) - 1
-                    if top_dim == -1:
-                        top_dim = -np.inf
-                    # print("top dim: ", top_dim)
-                    if top_dim <= len(sigma) - 1 - k + perversity(k):
-                        check += 1
-                        # print("ok, ", k)
-                else:
-                    # print("ok, ", k)
-                    check += 1
-            if check == n:
-                # print("appending ",sigma)
-                perverse_list.append(i)
-        return perverse_list
-
-
-def _powerset(iterable):
-    "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
-    s = list(iterable)
-    return list(map(list, chain.from_iterable(combinations(s, r)
-                                              for r in range(1, len(s)+1,
-                                                             1))))
-
-
-def _star(s, simplices):
-    return [x for x in simplices if s in x]
-
-
-def _intersect_lists(ls1, ls2):
-    '''intersect lists of lists'''
-    return [s for s in ls1 if s in ls2]
-
-
-def compute_coarsest_strata(simplices):
-    d = max(map(len, simplices))
-    X_d = [s for s in simplices if len(s) <= d]
-    Sj = [[] for k in range(d)]
-    while d >= 1:
-        for x in [s for s in X_d if len(s) == d]:
-            Bx = _star(x, simplices)
-            # print("Bx:",Bx)
-            set_to_check = _intersect_lists(Bx, X_d)
-            # print("stc:", set_to_check)
-            for w in set_to_check:
-                for y in set_to_check:
-                    By = _star(y, simplices)
-                    Bw = _star(w, simplices)
-                    # y>=w i.e. By <= Bw
-                    if IntersectionHomology.is_sublist(By, Bw):
-                        Sj[d - 1].append(x)
-            if len(Bx) == 0:
-                Sj[d - 1].append(x)
-        # print("Sj:",Sj)
-        Sj_cpx = list(chain.from_iterable(map(_powerset, Sj[d - 1])))
-        # print("cpx:",Sj_cpx)
-        lista = [s for s in X_d if s not in Sj_cpx]
-        # print("lista:",lista)
-        if len(lista) == 0:
-            d -= 1
-            X_d = [s for s in X_d if len(s) <= d and s not in Sj_cpx]
-        else:
-            d = max(map(len, lista))
-            X_d = [s for s in lista if len(s) <= d and s not in Sj_cpx]
-        # print("X_d:",X_d)
-    return Sj
diff --git a/gdeep/utility/utils.py b/gdeep/utility/utils.py
index fb4fd686..9998aa4e 100644
--- a/gdeep/utility/utils.py
+++ b/gdeep/utility/utils.py
@@ -1,5 +1,5 @@
 from ctypes import Union
-import imp
+#import imp
 from IPython import get_ipython  # type: ignore
 
 import base64
diff --git a/requirements.txt b/requirements.txt
index 72f2096f..4ed34965 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,6 +5,7 @@ plotly
 matplotlib
 matplotlib_venn
 scikit-learn
+spacy
 scipy
 networkx
 torch==1.10.0


Test patch:
diff --git a/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py b/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py
index eb5e9c74..462b219d 100644
--- a/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py
+++ b/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py
@@ -1,16 +1,16 @@
 """Testing for decision_boundary_calculator."""
 # License: GNU AGPLv3
 from .. import GradientFlowDecisionBoundaryCalculator
-from ....data import *
+from gdeep.data import *
 from ....models import FFNet
 import torch
 
+from gdeep.utility import DEVICE
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def test_gfdbc_2_dim():
 
-    circle_detect_nn = FFNet([2, 2]).to(device)
+    circle_detect_nn = FFNet([2, 2]).to(DEVICE)
 
     g = GradientFlowDecisionBoundaryCalculator(
             model=circle_detect_nn,
@@ -24,7 +24,7 @@ def test_gfdbc_2_dim():
 
 def test_gfdbc_multiclass():
 
-    circle_nn_3d = FFNet([3, 3]).to(device)
+    circle_nn_3d = FFNet([3, 3]).to(DEVICE)
 
     g = GradientFlowDecisionBoundaryCalculator(model=
                                                circle_nn_3d,
diff --git a/gdeep/data/datasets/tests/__init__.py b/gdeep/data/datasets/tests/__init__.py
new file mode 100644
index 00000000..66247349
--- /dev/null
+++ b/gdeep/data/datasets/tests/__init__.py
@@ -0,0 +1,1 @@
+#from .data_cloud import DataCloud
\ No newline at end of file
diff --git a/gdeep/data/datasets/tests/img_data/img1.png b/gdeep/data/datasets/tests/img_data/img1.png
new file mode 100644
index 00000000..a7a51fb4
Binary files /dev/null and b/gdeep/data/datasets/tests/img_data/img1.png differ
diff --git a/gdeep/data/datasets/tests/img_data/img2.png b/gdeep/data/datasets/tests/img_data/img2.png
new file mode 100644
index 00000000..24366aa3
Binary files /dev/null and b/gdeep/data/datasets/tests/img_data/img2.png differ
diff --git a/gdeep/data/datasets/tests/img_data/img3.png b/gdeep/data/datasets/tests/img_data/img3.png
new file mode 100644
index 00000000..75363e75
Binary files /dev/null and b/gdeep/data/datasets/tests/img_data/img3.png differ
diff --git a/gdeep/data/datasets/tests/img_data/labels.csv b/gdeep/data/datasets/tests/img_data/labels.csv
new file mode 100644
index 00000000..162e88c8
--- /dev/null
+++ b/gdeep/data/datasets/tests/img_data/labels.csv
@@ -0,0 +1,4 @@
+"Filename","Label"
+img1.png,0
+img2.png,1
+img3.png,2
diff --git a/gdeep/data/tests/test_data_cloud.py b/gdeep/data/datasets/tests/test_data_cloud.py
similarity index 98%
rename from gdeep/data/tests/test_data_cloud.py
rename to gdeep/data/datasets/tests/test_data_cloud.py
index 520230d6..7d505cfe 100644
--- a/gdeep/data/tests/test_data_cloud.py
+++ b/gdeep/data/datasets/tests/test_data_cloud.py
@@ -1,20 +1,20 @@
 # %%
-from gdeep.data import _DataCloud
+from os import remove, makedirs, environ
+from os.path import join, exists
+import logging
 
 import google  # type: ignore
 from google.cloud import storage  # type: ignore
 from google.cloud.storage import Bucket  # type: ignore
 from google.auth.exceptions import DefaultCredentialsError  # type: ignore
 import hashlib
-import logging
-from os import remove, makedirs, environ
-from os.path import join, exists
 import pytest
 import random
 from shutil import rmtree
 
+from gdeep.data.datasets import _DataCloud
 from gdeep.utility.utils import get_checksum
-from gdeep.utility.constants import DATASET_BUCKET_NAME
+from gdeep.utility import DATASET_BUCKET_NAME
 
 LOGGER = logging.getLogger(__name__)
 
diff --git a/gdeep/data/tests/test_dataset_cloud.py b/gdeep/data/datasets/tests/test_dataset_cloud.py
similarity index 99%
rename from gdeep/data/tests/test_dataset_cloud.py
rename to gdeep/data/datasets/tests/test_dataset_cloud.py
index 3a6ad417..413d332a 100644
--- a/gdeep/data/tests/test_dataset_cloud.py
+++ b/gdeep/data/datasets/tests/test_dataset_cloud.py
@@ -1,4 +1,4 @@
-from gdeep.data import DatasetCloud, dataset_cloud
+from gdeep.data.datasets import DatasetCloud, dataset_cloud
 
 import hashlib
 import logging
diff --git a/gdeep/data/datasets/tests/test_dataset_for_image.py b/gdeep/data/datasets/tests/test_dataset_for_image.py
new file mode 100644
index 00000000..f0039714
--- /dev/null
+++ b/gdeep/data/datasets/tests/test_dataset_for_image.py
@@ -0,0 +1,41 @@
+from gdeep.data.datasets import ImageClassificationFromFiles, \
+    DataLoaderBuilder, DlBuilderFromDataCloud
+
+from gdeep.data import TransformingDataset
+from gdeep.data.preprocessors import ToTensorImage
+import os
+from os.path import join
+import numpy as np
+import torch
+import logging
+from google.auth.exceptions import DefaultCredentialsError  # type: ignore
+
+LOGGER = logging.getLogger(__name__)
+
+
+
+def test_images_from_file():
+    """test DatasetImageClassificationFromFiles"""
+    file_path = os.path.dirname(os.path.realpath(__file__))
+    transform = ToTensorImage((32,32))  # this is already fitted
+    ds = ImageClassificationFromFiles(
+        os.path.join(file_path,"img_data"),
+        os.path.join(file_path,"img_data","labels.csv"))
+
+    tds = TransformingDataset(ds, transform)
+    dl, *_ = DataLoaderBuilder((tds,)).build(({"batch_size" : 2},))
+    assert len(next(iter(dl))[0].shape) == 4
+
+
+def test_dlbuilderfromdatacloud():
+    dataset_name = "SmallDataset"
+    download_directory = join("examples", "data", "DatasetCloud")
+
+    dl_cloud_builder = DlBuilderFromDataCloud(dataset_name,
+                                    download_directory)
+
+    train_dataloader, val_dataloader, test_dataloader = \
+        dl_cloud_builder.build(({"batch_size" : 10},))
+    x, y = next(iter(train_dataloader))
+    assert x.shape == torch.Size([10, 5])
+    assert y.shape == torch.Size([10])
\ No newline at end of file
diff --git a/gdeep/data/datasets/tests/test_dataset_from_array.py b/gdeep/data/datasets/tests/test_dataset_from_array.py
new file mode 100644
index 00000000..3ef619e7
--- /dev/null
+++ b/gdeep/data/datasets/tests/test_dataset_from_array.py
@@ -0,0 +1,36 @@
+from gdeep.data.datasets import ImageClassificationFromFiles, \
+    DataLoaderBuilder, FromArray, DlBuilderFromDataCloud
+
+from gdeep.data.preprocessors import ToTensorImage
+import os
+from os.path import join
+import numpy as np
+import torch
+import logging
+from google.auth.exceptions import DefaultCredentialsError  # type: ignore
+
+LOGGER = logging.getLogger(__name__)
+
+
+
+def test_array():
+    """test class DatasetFromArray"""
+    X = np.random.rand(10,4)
+    y = np.random.randint(3, size=10)
+    ds = FromArray(X, y)
+    dl, *_ = DataLoaderBuilder((ds,)).build(({"batch_size" : 1},))
+    item = next(iter(dl))
+    assert torch.norm(item[0] - torch.tensor(X[0])) < 1e-6
+    assert item[1] == torch.tensor(y[0])
+    assert len(dl.dataset) == 10
+
+def test_array_tensor():
+    """test class DatasetFromArray"""
+    X = torch.rand(10,4)
+    y = torch.randint(3, size=(10,))
+    ds = FromArray(X, y)
+    dl, *_ = DataLoaderBuilder((ds,)).build()
+    item = next(iter(dl))
+    assert torch.norm(item[0] - X[0]) < 1e-6
+    assert item[1] == y[0]
+
diff --git a/gdeep/data/tests/test_parallel_orbit.py b/gdeep/data/datasets/tests/test_parallel_orbit.py
similarity index 99%
rename from gdeep/data/tests/test_parallel_orbit.py
rename to gdeep/data/datasets/tests/test_parallel_orbit.py
index fa8031cf..faa0291a 100644
--- a/gdeep/data/tests/test_parallel_orbit.py
+++ b/gdeep/data/datasets/tests/test_parallel_orbit.py
@@ -1,6 +1,9 @@
-from ..parallel_orbit import OrbitsGenerator
+
 import numpy as np
 
+from ..parallel_orbit import OrbitsGenerator
+
+
 parameters = (1.0, 2.0, 3.0)
 num_orbits_per_class = 10
 num_pts_per_orbit = 15
diff --git a/gdeep/data/preprocessors/tests/__init__.py b/gdeep/data/preprocessors/tests/__init__.py
new file mode 100644
index 00000000..24fa198e
--- /dev/null
+++ b/gdeep/data/preprocessors/tests/__init__.py
@@ -0,0 +1,1 @@
+#from ..data_cloud import DataCloud
diff --git a/gdeep/data/preprocessors/tests/img_data/img1.png b/gdeep/data/preprocessors/tests/img_data/img1.png
new file mode 100644
index 00000000..a7a51fb4
Binary files /dev/null and b/gdeep/data/preprocessors/tests/img_data/img1.png differ
diff --git a/gdeep/data/preprocessors/tests/img_data/img2.png b/gdeep/data/preprocessors/tests/img_data/img2.png
new file mode 100644
index 00000000..75363e75
Binary files /dev/null and b/gdeep/data/preprocessors/tests/img_data/img2.png differ
diff --git a/gdeep/data/preprocessors/tests/img_data/img3.png b/gdeep/data/preprocessors/tests/img_data/img3.png
new file mode 100644
index 00000000..24366aa3
Binary files /dev/null and b/gdeep/data/preprocessors/tests/img_data/img3.png differ
diff --git a/gdeep/data/preprocessors/tests/img_data/labels.csv b/gdeep/data/preprocessors/tests/img_data/labels.csv
new file mode 100644
index 00000000..162e88c8
--- /dev/null
+++ b/gdeep/data/preprocessors/tests/img_data/labels.csv
@@ -0,0 +1,4 @@
+"Filename","Label"
+img1.png,0
+img2.png,1
+img3.png,2
diff --git a/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py b/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py
new file mode 100644
index 00000000..d84fc833
--- /dev/null
+++ b/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py
@@ -0,0 +1,29 @@
+from sklearn import preprocessing
+from gdeep.data.datasets import ImageClassificationFromFiles
+from gdeep.data.preprocessors import ToTensorImage
+import os
+from ...preprocessing_pipeline import PreprocessingPipeline
+from gdeep.data.preprocessors import Normalization
+from ...transforming_dataset import TransformingDataset
+
+
+def test_preeprocessing_pipeline():
+    """test preprocessing pipeline"""
+    normalize = Normalization()
+    to_tensor = ToTensorImage([32,32])
+    file_path = os.path.dirname(os.path.realpath(__file__))
+
+    image_dataset = ImageClassificationFromFiles(
+        os.path.join(file_path, "img_data"),
+        os.path.join(file_path, "img_data", "labels.csv"))
+    # define preprocessing pipeline
+    preprocessing_pipeline = PreprocessingPipeline((to_tensor, normalize))
+
+    # fit preprocessing pipeline to dataset
+    preprocessing_pipeline.fit_to_dataset(image_dataset)
+    
+    # apply preprocessing pipeline to dataset
+    preprocessed_dataset = preprocessing_pipeline.attach_transform_to_dataset(image_dataset)
+    
+    assert len(preprocessed_dataset[0][0].shape) == 3
+    
diff --git a/gdeep/data/tests/__init__.py b/gdeep/data/tests/__init__.py
deleted file mode 100644
index 9974dc42..00000000
--- a/gdeep/data/tests/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-#from ..data_cloud import DataCloud
\ No newline at end of file
diff --git a/gdeep/data/tests/test_torch_datasets.py b/gdeep/data/tests/test_torch_datasets.py
deleted file mode 100644
index d2f1113d..00000000
--- a/gdeep/data/tests/test_torch_datasets.py
+++ /dev/null
@@ -1,56 +0,0 @@
-from gdeep.data import DataLoaderFromArray, TorchDataLoader, DlBuilderFromDataCloud
-import os
-from os.path import join
-import numpy as np
-import torch
-import logging
-from google.auth.exceptions import DefaultCredentialsError  # type: ignore
-
-LOGGER = logging.getLogger(__name__)
-
-
-
-def test_array():
-    """test class DataLoaderFromArray"""
-    X = np.random.rand(10,4)
-    y = np.random.randint(3, size=10)
-    X1 = np.random.rand(8, 5)
-    y1 = np.random.randint(3, size=8)
-    dl = DataLoaderFromArray(X, y)
-    dl2 = DataLoaderFromArray(X, y, X1, y1)
-    dl.build_dataloaders()
-    v1, v2, v3 = dl2.build_dataloaders()
-    assert torch.norm(next(iter(v1))[0] - torch.tensor(X[0])) < 1e-6
-    assert next(iter(v1))[1] == torch.tensor(y[0])
-
-def test_array_tensor():
-    """test class DataLoaderFromArray"""
-    X = torch.rand(10,4)
-    y = torch.randint(3, size=(10,))
-    X1 = torch.rand(8, 5)
-    y1 = torch.randint(3, size=(8,))
-    dl = DataLoaderFromArray(X, y)
-    dl2 = DataLoaderFromArray(X, y, X1, y1)
-    dl.build_dataloaders()
-    v1, v2, v3 = dl2.build_dataloaders()
-    assert torch.norm(next(iter(v1))[0] - torch.tensor(X[0])) < 1e-6
-    assert next(iter(v1))[1] == torch.tensor(y[0])
-
-def test_torchdataloader():
-    dl = TorchDataLoader(name="DoubleTori")
-    # train_indices = list(range(160))
-    dl_tr, dl_val = dl.build_dataloaders(batch_size=23)
-        
-        
-def test_dlbuilderfromdatacloud():
-    dataset_name = "SmallDataset"
-    download_directory = join("examples", "data", "DatasetCloud")
-
-    dl_cloud_builder = DlBuilderFromDataCloud(dataset_name,
-                                    download_directory)
-
-    train_dataloader, val_dataloader, test_dataloader = \
-        dl_cloud_builder.build_dataloaders(batch_size=10)
-    x, y = next(iter(train_dataloader))
-    assert x.shape == torch.Size([10, 5])
-    assert y.shape == torch.Size([10])
\ No newline at end of file
diff --git a/gdeep/search/tests/test_persformer_hyperparamter_search.py b/gdeep/search/tests/test_persformer_hyperparamter_search.py
index f95100d9..9d2973d5 100644
--- a/gdeep/search/tests/test_persformer_hyperparamter_search.py
+++ b/gdeep/search/tests/test_persformer_hyperparamter_search.py
@@ -5,7 +5,7 @@
 from os.path import join
 from shutil import rmtree
 
-from gdeep.utility.constants import ROOT_DIR
+from gdeep.utility import ROOT_DIR
 
 from gdeep.search import clean_up_files
 
diff --git a/gdeep/pipeline/tests/test_pipeline.py b/gdeep/trainer/tests/test_trainer.py
similarity index 70%
rename from gdeep/pipeline/tests/test_pipeline.py
rename to gdeep/trainer/tests/test_trainer.py
index 9ca429cd..eca84d62 100644
--- a/gdeep/pipeline/tests/test_pipeline.py
+++ b/gdeep/trainer/tests/test_trainer.py
@@ -1,12 +1,13 @@
 from torch import nn
 from torch.optim import SGD
 from sklearn.model_selection import StratifiedKFold
-from gdeep.pipeline import Pipeline
+import numpy as np
+
+from gdeep.trainer import Trainer
 from gdeep.models import FFNet
-from gdeep.data import TorchDataLoader
 from gdeep.search import GiottoSummaryWriter
-import numpy as np
-from gdeep.data import DataLoaderFromArray, TorchDataLoader
+from gdeep.data.datasets import FromArray, DataLoaderBuilder, \
+    get_dataset
 from gdeep.search import clean_up_files
 
 
@@ -14,25 +15,26 @@
 def test_pipe_1():
     # model
     class model1(nn.Module):
+
         def __init__(self):
             super(model1, self).__init__()
             self.seqmodel = nn.Sequential(nn.Flatten(), FFNet(arch=[3, 5, 4]))
+
         def forward(self, x):
             return self.seqmodel(x).reshape(-1,2,2)
 
     model = model1()
     # dataloaders
-    X = np.random.rand(100,3)
-    y = np.random.randint(2,size=100*2).reshape(-1,2)
-    dl = DataLoaderFromArray(X, y)
-    dl_tr, dl_val, dl_ts = dl.build_dataloaders(batch_size=23)
+    X = np.array(np.random.rand(100, 3), dtype=np.float32)
+    y = np.array(np.random.randint(2, size=100*2).reshape(-1, 2), dtype=np.int64)
+    dl_tr, *_ = DataLoaderBuilder((FromArray(X, y),)).build(({"batch_size" : 23},))
 
     # loss function
     loss_fn = nn.CrossEntropyLoss()
     # tb writer
     writer = GiottoSummaryWriter()
     # pipeline
-    pipe = Pipeline(model, [dl_tr, None],
+    pipe = Trainer(model, [dl_tr, None],
                     loss_fn, writer)#,StratifiedKFold(5, shuffle=True))
     # then one needs to train the model using the pipeline!
     pipe.train(SGD, 2, True, {"lr": 0.001}, n_accumulated_grads=2)
diff --git a/gdeep/utility/intersection_homology/tests/__init__.py b/gdeep/utility/intersection_homology/tests/__init__.py
deleted file mode 100644
index 8b137891..00000000
--- a/gdeep/utility/intersection_homology/tests/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-
diff --git a/gdeep/utility/intersection_homology/tests/test_intersection_homology.py b/gdeep/utility/intersection_homology/tests/test_intersection_homology.py
deleted file mode 100644
index 55b83e1c..00000000
--- a/gdeep/utility/intersection_homology/tests/test_intersection_homology.py
+++ /dev/null
@@ -1,122 +0,0 @@
-import numpy as np
-import networkx as nx
-from ...intersection_homology import IntersectionHomology, \
-    compute_coarsest_strata
-
-
-def test_IntersectionHomology():
-    '''check if the class is consistent.
-    '''
-    ih = IntersectionHomology()
-    X = np.random.random((7, 3))
-    ih.fit(X)
-    ih.transform(X)
-    ih.plot()
-    ih.transform(X, [0])
-
-
-def test_IntersectionHomology2():
-    '''check if the class is consistent.
-    '''
-    ih = IntersectionHomology()
-    G = nx.Graph()
-    edges = [(0, 1), (1, 2), (0, 3), (2, 3), (0, 4),
-             (0, 5), (4, 6), (5, 6), (1, 3), (0, 2),
-             (0, 6), (4, 5)]
-    G.add_edges_from(edges)
-    ih.fit(G)
-    ih.transform(G)
-    ih.plot()
-    ih.transform(G, [0])
-
-
-def test_IntersectionHomology3():
-    '''check if the class is consistent.
-    '''
-    ih = IntersectionHomology(perversity_check=False)
-    G = nx.Graph()
-    edges = [(0, 1), (1, 2), (0, 3), (2, 3), (0, 4),
-             (0, 5), (4, 6), (5, 6), (1, 3), (0, 2),
-             (0, 6), (4, 5)]
-    G.add_edges_from(edges)
-    ih.fit(G)
-    ih.transform(G)
-    ih.plot()
-    ih.transform(G, [0])
-
-
-def test_use_of_coarse_strata():
-    '''check if the option of using the coarsest
-    stratification works.
-    '''
-    ih = IntersectionHomology(coarse_strata=True)
-    G = nx.Graph()
-    edges = [(0, 1), (1, 2), (0, 3), (2, 3), (0, 4),
-             (0, 5), (4, 6), (5, 6), (1, 3), (0, 2),
-             (0, 6), (4, 5)]
-    G.add_edges_from(edges)
-    ih.fit(G)
-    ih.transform(G)
-
-
-def test_IntersectionHomology4():
-    '''check if the class is consistent.
-    '''
-    ih = IntersectionHomology(perversity_check=False)
-    ih2 = IntersectionHomology(perversity_check=False)
-    G = nx.Graph()
-    edges = [(0, 1), (1, 2), (0, 3), (2, 3), (0, 4),
-             (0, 5), (4, 6), (5, 6), (1, 3), (0, 2),
-             (0, 6), (4, 5)]
-    G.add_edges_from(edges)
-    ih.fit(G)
-    res = ih.transform(G)
-    assert res == ih2.fit_transform(G)
-
-
-def test_low_():
-    # unit tests
-    M = np.array([[0, 0, 0, 1, 0, 1, 0],
-                  [0, 0, 0, 0, 1, 1, 0],
-                  [0, 0, 0, 1, 1, 0, 0],
-                  [0, 0, 1, 0, 0, 0, 1],
-                  [0, 1, 1, 0, 0, 0, 1],
-                  [0, 1, 0, 0, 0, 0, 1],
-                  [0, 0, 0, 1, 1, 1, 0]])
-
-    assert IntersectionHomology._low_(M, 0) == 0
-    assert IntersectionHomology._low_(M, 1) == 5
-    assert IntersectionHomology._low_(M, 2) == 4
-    assert IntersectionHomology._low_(M, 3) == 6
-
-
-def test_reduce():
-    # unit test
-    mat_paper = np.array([
-        [1, 0, 0, 0],
-        [0, 1, 0, 0],
-        [0, 0, 1, 0],
-        [0, 0, 0, 1],
-        [1, 0, 0, 1],
-        [1, 0, 1, 0],
-        [0, 1, 1, 0],
-        [0, 1, 0, 1]
-    ])
-    mm, _ = IntersectionHomology._reduce(mat_paper)
-    assert (mm == np.array([[1, 0, 0, 1],
-           [0, 1, 0, 1],
-           [0, 0, 1, 1],
-           [0, 0, 0, 1],
-           [1, 0, 0, 0],
-           [1, 0, 1, 0],
-           [0, 1, 1, 0],
-           [0, 1, 0, 0]])).all()
-
-
-def test_coarsest_strata():
-    # unit test
-    simp = [[0], [1], [2], [3], [0, 1],
-            [1, 2], [1, 3], [2, 3], [1, 2, 3]]
-    assert compute_coarsest_strata(simp) == [[],
-                                             [[0, 1]],
-                                             [[1, 2, 3]]]