Repository: iKernels/transformers-lightning (language: python)

Issue:
Remove IterableDataset
`IterableDataset` allowed usage of large dataset without loading the entire set of data in memory. However it had many drawbacks, like problem with checkpoints loading and computation of the number of epochs and steps. Moreover, defining the `__len__()` methods was not a solution since `lightning` would recognize the dataset as a normal one instead of an iterable. 
Remove IterableDataset
`IterableDataset` allowed usage of large dataset without loading the entire set of data in memory. However it had many drawbacks, like problem with checkpoints loading and computation of the number of epochs and steps. Moreover, defining the `__len__()` methods was not a solution since `lightning` would recognize the dataset as a normal one instead of an iterable. 


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 055734f..160d388 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## 0.6.0
 
+- Added `CompressedDataModule` based on `CompressedDataset`
+
+- Added `CompressedDataset` based on [`CompressedDictionary`](https://github.com/lucadiliello/compressed-dictionary)
+
+- Removed `IterableDataset`
+
 - Metrics has been moved to the `torchmetrics` library ([#81](https://github.com/iKernels/transformers-lightning/issues/81))
 
 - Removed losses package because it has been empty for months.
diff --git a/README.md b/README.md
index 6fe27bd..1c9e630 100644
--- a/README.md
+++ b/README.md
@@ -37,9 +37,6 @@ pip install git+https://github.com/iKernels/transformers-lightning.git@0.5.4 --u
 ## Documentation
 
 The documentation of each component is described in the relative folder.
-Follows a general schema that clearly explains connections between components and data flow.
-
-<img src="./pictures/transformers_lightning.svg">
 
 
 <a name="main"></a>
diff --git a/pictures/transformers_lightning.svg b/pictures/transformers_lightning.svg
deleted file mode 100644
index 1bc3c41..0000000
--- a/pictures/transformers_lightning.svg
+++ /dev/null
@@ -1,157 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" fill-opacity="1" color-rendering="auto" color-interpolation="auto" text-rendering="auto" stroke="black" stroke-linecap="square" width="556" stroke-miterlimit="10" shape-rendering="auto" stroke-opacity="1" fill="black" stroke-dasharray="none" font-weight="normal" stroke-width="1" height="338" font-family="'Dialog'" font-style="normal" stroke-linejoin="miter" font-size="12px" stroke-dashoffset="0" image-rendering="auto">
-  <!--Generated by ySVG 2.5-->
-  <defs id="genericDefs"/>
-  <g>
-    <defs id="defs1">
-      <clipPath clipPathUnits="userSpaceOnUse" id="clipPath1">
-        <path d="M0 0 L556 0 L556 338 L0 338 L0 0 Z"/>
-      </clipPath>
-      <clipPath clipPathUnits="userSpaceOnUse" id="clipPath2">
-        <path d="M1115 440 L1671 440 L1671 778 L1115 778 L1115 440 Z"/>
-      </clipPath>
-    </defs>
-    <g fill="white" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="translate(-1115,-440)" stroke="white">
-      <rect x="1115" width="556" height="338" y="440" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g fill="rgb(204,255,204)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(204,255,204)">
-      <rect x="1278" width="106" height="27" y="456" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1278" width="106" height="27" y="456" clip-path="url(#clipPath2)"/>
-      <text x="1307.9521" xml:space="preserve" y="474.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Adapter</text>
-    </g>
-    <g fill="rgb(204,255,204)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(204,255,204)">
-      <rect x="1278" width="106" height="27" y="502" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1278" width="106" height="27" y="502" clip-path="url(#clipPath2)"/>
-      <text x="1307.9521" xml:space="preserve" y="520.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Adapter</text>
-    </g>
-    <g fill="rgb(204,255,204)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(204,255,204)">
-      <rect x="1278" width="106" height="27" y="548" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1278" width="106" height="27" y="548" clip-path="url(#clipPath2)"/>
-      <text x="1307.9521" xml:space="preserve" y="566.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Adapter</text>
-    </g>
-    <g fill="rgb(204,255,255)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(204,255,255)">
-      <rect x="1414" width="106" height="27" y="456" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1414" width="106" height="27" y="456" clip-path="url(#clipPath2)"/>
-      <text x="1444.9893" xml:space="preserve" y="474.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Dataset</text>
-    </g>
-    <g fill="rgb(204,255,255)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(204,255,255)">
-      <rect x="1414" width="106" height="27" y="502" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1414" width="106" height="27" y="502" clip-path="url(#clipPath2)"/>
-      <text x="1444.9893" xml:space="preserve" y="520.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Dataset</text>
-    </g>
-    <g fill="rgb(204,255,255)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(204,255,255)">
-      <rect x="1414" width="106" height="27" y="548" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1414" width="106" height="27" y="548" clip-path="url(#clipPath2)"/>
-      <text x="1444.9893" xml:space="preserve" y="566.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Dataset</text>
-    </g>
-    <g fill="rgb(153,204,255)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(153,204,255)">
-      <rect x="1550" width="106" height="27" y="502" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1550" width="106" height="27" y="502" clip-path="url(#clipPath2)"/>
-      <text x="1568.2041" xml:space="preserve" y="520.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">DataModule</text>
-    </g>
-    <g fill="rgb(255,153,204)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(255,153,204)">
-      <rect x="1550" width="106" height="27" y="690" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1550" width="106" height="27" y="690" clip-path="url(#clipPath2)"/>
-      <text x="1582.1816" xml:space="preserve" y="708.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Trainer</text>
-    </g>
-    <g fill="silver" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="silver">
-      <rect x="1375" width="106" height="27" y="690" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1375" width="106" height="27" y="690" clip-path="url(#clipPath2)"/>
-      <text x="1410.293" xml:space="preserve" y="708.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Model</text>
-    </g>
-    <g fill="lime" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="lime">
-      <rect x="1200" width="106" height="27" y="690" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1200" width="106" height="27" y="690" clip-path="url(#clipPath2)"/>
-      <text x="1231.9238" xml:space="preserve" y="708.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Metrics</text>
-    </g>
-    <g fill="rgb(255,153,0)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(255,153,0)">
-      <rect x="1200" width="106" height="27" y="736" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1200" width="106" height="27" y="736" clip-path="url(#clipPath2)"/>
-      <text x="1225.3379" xml:space="preserve" y="754.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Callbacks</text>
-    </g>
-    <g fill="yellow" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="yellow">
-      <rect x="1200" width="106" height="27" y="644" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <rect fill="none" x="1200" width="106" height="27" y="644" clip-path="url(#clipPath2)"/>
-      <text x="1233.5967" xml:space="preserve" y="662.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Losses</text>
-    </g>
-    <g fill="rgb(255,204,0)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(255,204,0)">
-      <path d="M1130 455.5 L1233 455.5 L1233 470.5 L1248 470.5 L1248 483.5 L1130 483.5 Z" stroke="none" clip-path="url(#clipPath2)"/>
-      <path fill="white" d="M1233 455.5 L1233 470.5 L1248 470.5 Z" clip-path="url(#clipPath2)" fill-rule="evenodd" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <path fill="none" d="M1130 455.5 L1233 455.5 L1233 470.5 L1248 470.5 L1248 483.5 L1130 483.5 Z" clip-path="url(#clipPath2)"/>
-      <path fill="none" d="M1233 455.5 L1233 470.5 L1248 470.5 Z" clip-path="url(#clipPath2)" fill-rule="evenodd"/>
-      <text x="1152.5195" xml:space="preserve" y="474.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Train set file</text>
-    </g>
-    <g fill="rgb(255,204,0)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(255,204,0)">
-      <path d="M1130 501.5 L1233 501.5 L1233 516.5 L1248 516.5 L1248 529.5 L1130 529.5 Z" stroke="none" clip-path="url(#clipPath2)"/>
-      <path fill="white" d="M1233 501.5 L1233 516.5 L1248 516.5 Z" clip-path="url(#clipPath2)" fill-rule="evenodd" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <path fill="none" d="M1130 501.5 L1233 501.5 L1233 516.5 L1248 516.5 L1248 529.5 L1130 529.5 Z" clip-path="url(#clipPath2)"/>
-      <path fill="none" d="M1233 501.5 L1233 516.5 L1248 516.5 Z" clip-path="url(#clipPath2)" fill-rule="evenodd"/>
-      <text x="1138.3574" xml:space="preserve" y="520.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Validation set file</text>
-    </g>
-    <g fill="rgb(255,204,0)" text-rendering="geometricPrecision" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke="rgb(255,204,0)">
-      <path d="M1130 547.5 L1233 547.5 L1233 562.5 L1248 562.5 L1248 575.5 L1130 575.5 Z" stroke="none" clip-path="url(#clipPath2)"/>
-      <path fill="white" d="M1233 547.5 L1233 562.5 L1248 562.5 Z" clip-path="url(#clipPath2)" fill-rule="evenodd" stroke="none"/>
-    </g>
-    <g text-rendering="geometricPrecision" stroke-miterlimit="1.45" shape-rendering="geometricPrecision" transform="matrix(1,0,0,1,-1115,-440)" stroke-linecap="butt">
-      <path fill="none" d="M1130 547.5 L1233 547.5 L1233 562.5 L1248 562.5 L1248 575.5 L1130 575.5 Z" clip-path="url(#clipPath2)"/>
-      <path fill="none" d="M1233 547.5 L1233 562.5 L1248 562.5 Z" clip-path="url(#clipPath2)" fill-rule="evenodd"/>
-      <text x="1155.1006" xml:space="preserve" y="566.0352" clip-path="url(#clipPath2)" font-family="sans-serif" stroke="none">Test set file</text>
-      <path fill="none" d="M1384 469.5 L1406 469.5" clip-path="url(#clipPath2)"/>
-      <path d="M1414 469.5 L1402 464.5 L1405 469.5 L1402 474.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1384 515.5 L1406 515.5" clip-path="url(#clipPath2)"/>
-      <path d="M1414 515.5 L1402 510.5 L1405 515.5 L1402 520.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1384 561.5 L1406 561.5" clip-path="url(#clipPath2)"/>
-      <path d="M1414 561.5 L1402 556.5 L1405 561.5 L1402 566.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1506.9131 548 L1555.5087 531.5632" clip-path="url(#clipPath2)"/>
-      <path d="M1563.0869 529 L1550.1176 528.1085 L1554.5614 531.8837 L1553.3215 537.5812 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1520 515.5 L1542 515.5" clip-path="url(#clipPath2)"/>
-      <path d="M1550 515.5 L1538 510.5 L1541 515.5 L1538 520.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1506.9131 483 L1555.5087 499.4368" clip-path="url(#clipPath2)"/>
-      <path d="M1563.0869 502 L1553.3215 493.4187 L1554.5614 499.1164 L1550.1176 502.8916 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1603 529 L1603 682" clip-path="url(#clipPath2)"/>
-      <path d="M1603 690 L1608 678 L1603 681 L1598 678 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1550 703.5 L1489 703.5" clip-path="url(#clipPath2)"/>
-      <path d="M1481 703.5 L1493 708.5 L1490 703.5 L1493 698.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1375 703.5 L1314 703.5" clip-path="url(#clipPath2)"/>
-      <path d="M1306 703.5 L1318 708.5 L1315 703.5 L1318 698.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1376.6414 717 L1312.0958 733.9662" clip-path="url(#clipPath2)"/>
-      <path d="M1304.3586 736 L1317.2355 737.7851 L1313.063 733.712 L1314.6934 728.1136 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1368.9042 687.9662 L1312.0958 673.0338" clip-path="url(#clipPath2)"/>
-      <path d="M1376.6414 690 L1366.3066 682.1136 L1367.937 687.712 L1363.7645 691.7851 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path d="M1304.3586 671 L1314.6934 678.8864 L1313.063 673.288 L1317.2355 669.2149 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1247.9702 469.5 L1270 469.5" clip-path="url(#clipPath2)"/>
-      <path d="M1278 469.5 L1266 464.5 L1269 469.5 L1266 474.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1247.9702 515.5 L1270 515.5" clip-path="url(#clipPath2)"/>
-      <path d="M1278 515.5 L1266 510.5 L1269 515.5 L1266 520.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-      <path fill="none" d="M1247.9702 561.5 L1270 561.5" clip-path="url(#clipPath2)"/>
-      <path d="M1278 561.5 L1266 556.5 L1269 561.5 L1266 566.5 Z" clip-path="url(#clipPath2)" stroke="none"/>
-    </g>
-  </g>
-</svg>
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..35a2a82
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,20 @@
+[build-system]
+requires = [
+    "setuptools",
+    "wheel",
+]
+
+[tool.black]
+# https://github.com/psf/black
+line-length = 120
+target-version = ["py38"]
+exclude = "(.eggs|.git|.hg|.mypy_cache|.nox|.tox|.venv|.svn|_build|buck-out|build|dist)"
+
+[tool.isort]
+known_first_party = [
+    "torchmetrics",
+    "tests",
+]
+skip_glob = []
+profile = "black"
+line_length = 120
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 12db96e..e032a62 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 pytest
 pytorch-lightning>=1.2.7
 transformers>=4.4.2
-torch>=1.8.1
\ No newline at end of file
+torch>=1.8.1
+compressed-dictionary>=1.2.1
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 785e2f7..f0af540 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,7 @@
-import setuptools
 import json
 
+import setuptools
+
 
 def load_long_description():
     with open("DESCRIPTION.md", "r") as fh:
diff --git a/transformers_lightning/__init__.py b/transformers_lightning/__init__.py
index e349dfd..e43578e 100644
--- a/transformers_lightning/__init__.py
+++ b/transformers_lightning/__init__.py
@@ -1,7 +1,7 @@
-import transformers_lightning.datasets
-import transformers_lightning.defaults
 import transformers_lightning.callbacks
-import transformers_lightning.utils
-import transformers_lightning.models
 import transformers_lightning.datamodules
+import transformers_lightning.datasets
+import transformers_lightning.defaults
 import transformers_lightning.language_modeling
+import transformers_lightning.models
+import transformers_lightning.utils
diff --git a/transformers_lightning/adapters/README.md b/transformers_lightning/adapters/README.md
index f68a41b..5ce8c88 100644
--- a/transformers_lightning/adapters/README.md
+++ b/transformers_lightning/adapters/README.md
@@ -6,37 +6,5 @@ An `Adapter` must override at least two methods of the superclass `SuperAdapter`
 
 The following example illustrates how to build an `Adapter` to read data from a `tsv` file and how to implement tokenization in the `preprocess_line` method.
 
-The `__iter__` method does not has arguments other than `self`. Paramters should be retrieved through `self.hparams`.
+The `__iter__` method does not has arguments other than `self`. Parameters should be retrieved through `self.hparams`.
 `preprocess_line` instead receives a line at a time and is strongly recommended to return a `dict`. `dict` improves readability and values are automagically concatenated by the `SuperDataModule` class.
-
-```python
-import csv
-from transformers_lightning import SuperAdapter
-
-
-class TransformersAdapter(SuperAdapter):
-
-    def __init__(self, *args, tokenizer=None, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.tokenizer = tokenizer
-
-    def __iter__(self) -> Iterable:
-        """ Return a generator of parsed lines. """
-        with open(self.filepath, "r") as fi:
-            reader = csv.reader(fi, delimiter='\t', quoting=csv.QUOTE_MINIMAL)
-            yield from reader
-
-    def preprocess_line(self, line: list) -> list:
-        """
-        Tokenize a single line. Suppose each line
-        containt two sentences and a label.
-        """
-        results = self.tokenizer.encode_plus(
-            line[0], line[1],
-            padding='max_length',
-            max_length=self.hparams.max_sequence_length,
-            truncation=True
-        )
-        res = { **results, 'label': line[2] }
-        return res
-```
\ No newline at end of file
diff --git a/transformers_lightning/adapters/__init__.py b/transformers_lightning/adapters/__init__.py
index 92feb2a..2ed1553 100644
--- a/transformers_lightning/adapters/__init__.py
+++ b/transformers_lightning/adapters/__init__.py
@@ -1,4 +1,4 @@
-from transformers_lightning.adapters.super_adapter import SuperAdapter
-from transformers_lightning.adapters.file_adapter import FileAdapter
 from transformers_lightning.adapters.csv_adapter import CSVAdapter
-from transformers_lightning.adapters.transformers_adapter import TransformersAdapter
\ No newline at end of file
+from transformers_lightning.adapters.file_adapter import FileAdapter
+from transformers_lightning.adapters.super_adapter import SuperAdapter
+from transformers_lightning.adapters.transformers_adapter import TransformersAdapter
diff --git a/transformers_lightning/adapters/csv_adapter.py b/transformers_lightning/adapters/csv_adapter.py
index f9c01a6..3358fe3 100644
--- a/transformers_lightning/adapters/csv_adapter.py
+++ b/transformers_lightning/adapters/csv_adapter.py
@@ -1,18 +1,20 @@
 import csv
+from argparse import Namespace
 from typing import Iterable
 
-from transformers_lightning.adapters import FileAdapter
+from transformers_lightning.adapters.file_adapter import FileAdapter
 from transformers_lightning.utils import strip_lines
 
 
 class CSVAdapter(FileAdapter):
     r"""
-    An Adapter to load data from a file. Only adds a `filepath` parameter.
-    It is still abstract and need to be subclassed.
+    An Adapter to load data from a CSV file. Only adds a `filepath` parameter.
     """
 
-    def __init__(self, *args, delimiter="\t", quoting=csv.QUOTE_MINIMAL, quotechar='"', **kwargs):
-        super().__init__(*args)
+    def __init__(
+        self, hparams: Namespace, filepath: str, delimiter="\t", quoting=csv.QUOTE_MINIMAL, quotechar='"', **kwargs
+    ):
+        super().__init__(hparams, filepath)
         self.csv_kwargs = {'delimiter': delimiter, 'quoting': quoting, 'quotechar': quotechar, **kwargs}
 
     def __iter__(self) -> Iterable:
diff --git a/transformers_lightning/adapters/file_adapter.py b/transformers_lightning/adapters/file_adapter.py
index 8e7af4b..75c7be3 100644
--- a/transformers_lightning/adapters/file_adapter.py
+++ b/transformers_lightning/adapters/file_adapter.py
@@ -1,5 +1,6 @@
 import os
 from argparse import Namespace
+
 from transformers_lightning.adapters.super_adapter import SuperAdapter
 
 
@@ -11,7 +12,8 @@ class FileAdapter(SuperAdapter):
 
     def __init__(self, hparams: Namespace, filepath: str) -> None:
         r"""
-        :param filepath: path of the file that should be loaded
+        Args:
+            filepath: path of the file that should be loaded
         """
         super().__init__(hparams)
 
diff --git a/transformers_lightning/adapters/transformers_adapter.py b/transformers_lightning/adapters/transformers_adapter.py
index 3317e51..c7ae464 100644
--- a/transformers_lightning/adapters/transformers_adapter.py
+++ b/transformers_lightning/adapters/transformers_adapter.py
@@ -1,6 +1,9 @@
-from typing import Union, List
+from argparse import Namespace
+from typing import List, Union
+
 import transformers
-from transformers_lightning.adapters import CSVAdapter
+
+from transformers_lightning.adapters.csv_adapter import CSVAdapter
 
 
 class TransformersAdapter(CSVAdapter):
@@ -9,16 +12,18 @@ class TransformersAdapter(CSVAdapter):
     transformers sentence in sequences of ids. It inherits from CSV because most of the NLP datasets
     are given in this format.
 
-    An example of `preprocess_line` for NLP applications with the `transfomers` library is given below:
-    >>> def preprocess_line(self, line: list) -> list:
-    >>>     results = self.tokenizer.encode_plus(
-    >>>         line[1],
-    >>>         padding='max_length',
-    >>>         max_length=self.hparams.max_sequence_length,
-    >>>         truncation=True
-    >>>     )
-    >>>     results['words_tails'] = self._convert_ids_to_word_tails(results['input_ids'])
-    >>>     return results
+    An example of `preprocess_line` for NLP applications with the `transfomers` library is given below.
+
+    Example:
+        >>> def preprocess_line(self, line: list) -> list:
+        >>>     results = self.tokenizer.encode_plus(
+        >>>         line[1],
+        >>>         padding='max_length',
+        >>>         max_length=self.hparams.max_sequence_length,
+        >>>         truncation=True
+        >>>     )
+        >>>     results['words_tails'] = self._convert_ids_to_word_tails(results['input_ids'])
+        >>>     return results
 
     Moreover, some useful tools are given as additional methods:
     - `_convert_ids_to_word_tails`: convert a sequence of ids in a list of boolean values where `True`
@@ -27,16 +32,17 @@ class TransformersAdapter(CSVAdapter):
 
     def __init__(
         self,
-        *args,
+        hparams: Namespace,
+        filepath: str,
         tokenizer: Union[transformers.PreTrainedTokenizer, transformers.PreTrainedTokenizerFast] = None,
         **kwargs
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__(hparams, filepath, **kwargs)
         self.tokenizer = tokenizer
 
     def _convert_ids_to_word_tails(self, ids: List[int]):
         r"""
-        Convert ids to a mask with True values representing tail words (starting with '##')
+        Convert ids to a mask with True values representing tail words (e.g. starting with '##' in BERT)
 
         Example:
         >>> sentence = "This is a syntactically correct sentence"
diff --git a/transformers_lightning/callbacks/transformers_model_checkpoint.py b/transformers_lightning/callbacks/transformers_model_checkpoint.py
index 9330ff6..f84134f 100644
--- a/transformers_lightning/callbacks/transformers_model_checkpoint.py
+++ b/transformers_lightning/callbacks/transformers_model_checkpoint.py
@@ -3,6 +3,7 @@
 
 from pytorch_lightning.callbacks.base import Callback
 from pytorch_lightning.utilities import rank_zero_only
+
 from transformers_lightning.utils import dump_json, is_simple
 
 PARAMS_FILENAME = "hparams.json"
diff --git a/transformers_lightning/datamodules/README.md b/transformers_lightning/datamodules/README.md
index 17fb570..2e77642 100644
--- a/transformers_lightning/datamodules/README.md
+++ b/transformers_lightning/datamodules/README.md
@@ -1,6 +1,6 @@
 # DataModules
 
-A `DataModule` is a collection of at most `3` dataset to provide data for both training, validation and testing in a flexible way. For more information about a `DataLoader`, see the original implementation in the [`pytorch-lightning`](https://pytorch-lightning.readthedocs.io/en/stable/datamodules.html) repository. This folder extends the original `DataLoader` with some simple utilities to:
+A `DataModule` is a collection of at most `3` dataset to provide data for both `training`, `validation` and `testing` in a flexible way. For more information about a `DataModule`, see the original implementation in the [`pytorch-lightning`](https://pytorch-lightning.readthedocs.io/en/stable/datamodules.html) repository. This folder extends the original `DataModule` with some simple utilities to:
 - Do training only then a `train_dataloader` method is defined.
 - Same for test.
 
diff --git a/transformers_lightning/datamodules/__init__.py b/transformers_lightning/datamodules/__init__.py
index 018eb75..f23d0cf 100644
--- a/transformers_lightning/datamodules/__init__.py
+++ b/transformers_lightning/datamodules/__init__.py
@@ -1,1 +1,3 @@
-from transformers_lightning.datamodules.superdatamodule import SuperDataModule
+from transformers_lightning.datamodules.adapter_datamodule import AdaptersDataModule
+from transformers_lightning.datamodules.compressed_datamodule import CompressedDataModule
+from transformers_lightning.datamodules.super_datamodule import SuperDataModule
diff --git a/transformers_lightning/datamodules/adapter_datamodule.py b/transformers_lightning/datamodules/adapter_datamodule.py
new file mode 100644
index 0000000..b9267bc
--- /dev/null
+++ b/transformers_lightning/datamodules/adapter_datamodule.py
@@ -0,0 +1,104 @@
+import multiprocessing
+from argparse import ArgumentParser, Namespace
+
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+
+from transformers_lightning import utils
+from transformers_lightning.adapters.super_adapter import SuperAdapter
+from transformers_lightning.datamodules.super_datamodule import SuperDataModule
+from transformers_lightning.datasets.map_dataset import MapDataset
+
+
+class AdaptersDataModule(SuperDataModule):
+    """
+    AdaptersDataModule should be used when you want to read and tokenizer data on-the-fly.
+    It implements some simple methods to check whether training, val or testing is required.
+    It work with adapters: you could define them inside the `__init__()` method or pass them
+    as arguments.
+    """
+
+    train_adapter: SuperAdapter = None
+    valid_adapter: SuperAdapter = None
+    test_adapter: SuperAdapter = None
+
+    def __init__(
+        self,
+        hparams: Namespace,
+        train_adapter: SuperAdapter = None,
+        valid_adapter: SuperAdapter = None,
+        test_adapter: SuperAdapter = None,
+    ):
+        super().__init__(hparams)
+        self.hparams = hparams
+
+        # instantiate eventual adapters passed from init method
+        if train_adapter is not None:
+            assert isinstance(train_adapter, SuperAdapter), f"Argument `train_adapter` must be of type `SuperAdapter`"
+            self.train_adapter = train_adapter
+
+        if valid_adapter is not None:
+            assert isinstance(valid_adapter, SuperAdapter), f"Argument `valid_adapter` must be of type `SuperAdapter`"
+            self.valid_adapter = valid_adapter
+
+        if test_adapter is not None:
+            assert (
+                isinstance(test_adapter, SuperAdapter) or isinstance(test_adapter, list)
+            ), f"Argument `test_adapter` must be of type `SuperAdapter` or List[SuperAdapter]"
+
+            if isinstance(test_adapter, list):
+                for adapter in test_adapter:
+                    assert isinstance(
+                        adapter, SuperAdapter
+                    ), (f"Argument `test_adapter` must be of type `SuperAdapter` or List[SuperAdapter]")
+            self.test_adapter = test_adapter
+        """
+        This space should be used to instantiate the Adapters it they were not passed through the kwargs
+
+        >>> self.train_adapter = TSVAdapter(self.hparams, "pre-training/train.tsv", delimiter="\t")
+        >>> self.valid_adapter = TSVAdapter(self.hparams, "pre-training/valid.tsv", delimiter="\t")
+        >>> self.test_adapter = TSVAdapter(self.hparams, "pre-training/test.tsv", delimiter="\t")
+        """
+
+    # Optional, called for every GPU/machine (assigning state is OK)
+    def setup(self, stage=None):
+        """
+        Load datasets only if respective Adapter are defined.
+        Finally check that is a 
+        This implementation should be enough for most subclasses.
+        """
+
+        if stage == 'fit' or stage is None:
+            if self.train_adapter is not None:
+                self.train_dataset = MapDataset(self.hparams, self.train_adapter, self.trainer)
+            if self.valid_adapter is not None:
+                self.valid_dataset = MapDataset(self.hparams, self.valid_adapter, self.trainer)
+
+            assert self.train_adapter is None or self.train_dataset is not None, (
+                f"Cannot specify `train_adapter` and then `train_dataset` is None: "
+                f"{self.train_adapter} and {self.train_dataset}"
+            )
+            assert self.valid_adapter is None or self.valid_dataset is not None, (
+                f"Cannot specify `valid_adapter` and then `valid_dataset` is None: "
+                f"{self.valid_adapter} and {self.valid_dataset}"
+            )
+
+        elif stage == 'test' or stage is None:
+            if self.test_adapter is not None:
+                if isinstance(self.test_adapter, SuperAdapter):
+                    self.test_adapter = [self.test_adapter]
+                self.test_dataset = [MapDataset(self.hparams, adapter, self.trainer) for adapter in self.test_adapter]
+
+            assert self.test_adapter is None or self.test_dataset is not None, (
+                f"Cannot specify `test_adapter` and then `test_dataset` is None: "
+                f"{self.test_adapter} and {self.test_dataset}"
+            )
+
+    def do_train(self):
+        return self.train_adapter is not None
+
+    def do_validation(self):
+        return self.valid_adapter is not None
+
+    def do_test(self):
+        return self.test_adapter is not None
diff --git a/transformers_lightning/datamodules/compressed_datamodule.py b/transformers_lightning/datamodules/compressed_datamodule.py
new file mode 100644
index 0000000..4d4c1df
--- /dev/null
+++ b/transformers_lightning/datamodules/compressed_datamodule.py
@@ -0,0 +1,84 @@
+import multiprocessing
+import os
+from argparse import ArgumentParser
+from typing import Callable
+
+import pytorch_lightning as pl
+from pytorch_lightning import _logger as logger
+from torch.utils.data import DataLoader
+
+from transformers_lightning import utils
+from transformers_lightning.datamodules.super_datamodule import SuperDataModule
+from transformers_lightning.datasets.compressed_dataset import CompressedDataset
+
+
+class CompressedDataModule(SuperDataModule):
+    r"""
+    CompressedDataModule implements some simple methods to check whether training, val or testing is required.
+    It uses fast `CompressedDictionary`s to store data and return them by decompressing on-the-fly.
+    """
+
+    train_filepath: str = None
+    valid_filepath: str = None
+    test_filepath: str = None
+
+    def __init__(self, hparams, collate_fn: Callable = utils.collate_single_fn):
+        super().__init__(hparams, collate_fn)
+
+        # instantiate eventual adapters passed from init method
+        if hparams.train_filepath is not None:
+            if not os.path.isfile(hparams.train_filepath):
+                raise ValueError(f"Argument `train_filepath` is not a valid file")
+            self.train_filepath = hparams.train_filepath
+
+        if hparams.valid_filepath is not None:
+            if not os.path.isfile(hparams.valid_filepath):
+                raise ValueError(f"Argument `valid_filepath` is not a valid file")
+            self.valid_filepath = hparams.valid_filepath
+
+        if hparams.test_filepath is not None:
+            for test_file in hparams.test_filepath:
+                if not os.path.isfile(test_file):
+                    raise ValueError(f"file `{test_file}` is not a valid test file")
+            self.test_filepath = hparams.test_filepath
+
+    # Optional, called for every GPU/machine (assigning state is OK)
+    def setup(self, stage=None):
+        r"""
+        Load datasets only if respective file is defined.
+        This implementation should be enough for most subclasses.
+        """
+
+        if stage == 'fit' or stage is None:
+            if self.do_train():
+                logger.info("Loading training dataset from CompressedDictionary...")
+                self.train_dataset = CompressedDataset(self.hparams, self.train_filepath)
+            if self.do_validation():
+                logger.info("Loading validation dataset from CompressedDictionary...")
+                self.valid_dataset = CompressedDataset(self.hparams, self.valid_filepath)
+
+        elif stage == 'test' or stage is None:
+            if self.do_test():
+                logger.info("Loading test dataset from CompressedDictionary...")
+                self.test_dataset = [CompressedDataset(self.hparams, filepath) for filepath in self.test_filepath]
+
+    def do_train(self):
+        return self.train_filepath is not None
+
+    def do_validation(self):
+        return self.valid_filepath is not None
+
+    def do_test(self):
+        return len(self.test_filepath) > 0
+
+    @staticmethod
+    def add_datamodule_specific_args(parser: ArgumentParser):
+        parser = super(CompressedDataModule, CompressedDataModule).add_datamodule_specific_args(parser)
+
+        parser.add_argument('--train_filepath', type=str, required=False, default=None, help="Path to training file")
+        parser.add_argument('--valid_filepath', type=str, required=False, default=None, help="Path to validation file")
+        parser.add_argument(
+            '--test_filepath', type=str, required=False, default=[], nargs='+', help="Path to test file(s)"
+        )
+
+        return parser
diff --git a/transformers_lightning/datamodules/super_datamodule.py b/transformers_lightning/datamodules/super_datamodule.py
new file mode 100644
index 0000000..6a23d85
--- /dev/null
+++ b/transformers_lightning/datamodules/super_datamodule.py
@@ -0,0 +1,94 @@
+import multiprocessing
+from abc import abstractmethod
+from argparse import ArgumentParser, Namespace
+from typing import Callable
+
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader, Dataset
+
+from transformers_lightning import utils
+
+
+class SuperDataModule(pl.LightningDataModule):
+    """
+    SuperDataModule should be the superclass of all the DataModule in your project.
+    It implements some simple methods to check whether training, val or testing is required.
+    Moreover, it adds to the command line parameters the basic arguments used by Dataset,
+    like `batch_size`, `val_batch_size`, `test_batch_size` and `num_workers`.
+
+    Example:
+
+    >>> if datamodule.do_train():
+    >>>     trainer.fit(model, datamodule=datamodule)
+    
+    >>> if datamodule.do_test():
+    >>>     trainer.test(model, datamodule=datamodule)
+    """
+
+    train_dataset: Dataset = None
+    valid_dataset: Dataset = None
+    test_dataset: Dataset = None
+
+    def __init__(self, hparams: Namespace, collate_fn: Callable = utils.collate_single_fn):
+        super().__init__()
+        self.hparams = hparams
+        self.collate_fn = collate_fn
+
+    @abstractmethod
+    def do_train(self) -> bool:
+        """ Whether to do training. """
+
+    @abstractmethod
+    def do_validation(self) -> bool:
+        """ Whether to do validation. """
+
+    @abstractmethod
+    def do_test(self):
+        """ Whether to do testing. """
+
+    def default_dataloader(self, dataset: Dataset, batch_size: int, shuffle: bool = True):
+        """ Return a dataloader with all usual default parameters. """
+        return DataLoader(
+            dataset,
+            batch_size=batch_size,
+            num_workers=self.hparams.num_workers,
+            pin_memory=True,
+            collate_fn=self.collate_fn,
+            shuffle=shuffle
+        )
+
+    def train_dataloader(self):
+        """ Return the training dataloader. """
+        if self.do_train():
+            return self.default_dataloader(self.train_dataset, self.hparams.batch_size, shuffle=True)
+        return None
+
+    def val_dataloader(self):
+        """ Return the validation dataloader. """
+        if self.do_validation():
+            return self.default_dataloader(self.valid_dataset, self.hparams.val_batch_size, shuffle=False)
+        return None
+
+    def test_dataloader(self):
+        """ Return the test dataloader. """
+        if self.do_test():
+            return [
+                self.default_dataloader(dataset, self.hparams.test_batch_size, shuffle=False)
+                for dataset in self.test_dataset
+            ]
+        return None
+
+    @staticmethod
+    def add_datamodule_specific_args(parser: ArgumentParser) -> ArgumentParser:
+        parser.add_argument(
+            '--num_workers',
+            required=False,
+            default=multiprocessing.cpu_count(),
+            type=int,
+            help='Number of workers to be used to load datasets'
+        )
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--val_batch_size', type=int, default=256)
+        parser.add_argument('--test_batch_size', type=int, default=256)
+
+        return parser
diff --git a/transformers_lightning/datamodules/superdatamodule.py b/transformers_lightning/datamodules/superdatamodule.py
deleted file mode 100644
index 8a57803..0000000
--- a/transformers_lightning/datamodules/superdatamodule.py
+++ /dev/null
@@ -1,169 +0,0 @@
-from argparse import ArgumentParser
-import multiprocessing
-from typing import Iterable
-
-from torch.utils.data.dataset import IterableDataset
-from transformers_lightning.adapters import SuperAdapter
-
-import pytorch_lightning as pl
-from torch.utils.data import DataLoader
-from transformers_lightning import utils
-from transformers_lightning.datasets import (TransformersIterableDataset, TransformersMapDataset)
-
-
-class SuperDataModule(pl.LightningDataModule):
-    """
-    SuperDataModule should be the parent class of all `LightingDataModules` in your project.
-    It implements some simple methods to check whether training, val or testing is required.
-    Example:
-
-    >>> if datamodule.do_train():
-    >>>     trainer.fit(model, datamodule=datamodule)
-    
-    >>> if datamodule.do_test():
-    >>>     trainer.test(model, datamodule=datamodule)
-    """
-
-    train_adapter: SuperAdapter = None
-    valid_adapter: SuperAdapter = None
-    test_adapter: SuperAdapter = None
-
-    def __init__(
-        self, hparams, train_adapter=None, valid_adapter=None, test_adapter=None, collate_fn=utils.collate_single_fn
-    ):
-        super().__init__()
-        self.hparams = hparams
-        self.collate_fn = collate_fn
-
-        # instantiate eventual adapters passed from init method
-        if train_adapter is not None:
-            assert isinstance(train_adapter, SuperAdapter), f"Argument `train_adapter` must be of type `SuperAdapter`"
-            self.train_adapter = train_adapter
-
-        if valid_adapter is not None:
-            assert isinstance(valid_adapter, SuperAdapter), f"Argument `valid_adapter` must be of type `SuperAdapter`"
-            self.valid_adapter = valid_adapter
-
-        if test_adapter is not None:
-            assert (
-                isinstance(test_adapter, SuperAdapter) or isinstance(test_adapter, list)
-            ), f"Argument `test_adapter` must be of type `SuperAdapter` or List[SuperAdapter]"
-
-            if isinstance(test_adapter, list):
-                for adapter in test_adapter:
-                    assert isinstance(
-                        adapter, SuperAdapter
-                    ), (f"Argument `test_adapter` must be of type `SuperAdapter` or List[SuperAdapter]")
-            self.test_adapter = test_adapter
-        """
-        This space should be used to instantiate the Adapters it they were not passed through the kwargs
-
-        >>> self.train_adapter = TSVAdapter(self.hparams, "pre-training/train.tsv", delimiter="\t")
-        >>> self.valid_adapter = TSVAdapter(self.hparams, "pre-training/valid.tsv", delimiter="\t")
-        >>> self.test_adapter = TSVAdapter(self.hparams, "pre-training/test.tsv", delimiter="\t")
-        """
-
-    # Optional, called for every GPU/machine (assigning state is OK)
-    def setup(self, stage=None):
-        """
-        Load datasets only if respective Adapter are defined.
-        Finally check that is a 
-        This implementation should be enough for most subclasses.
-        """
-        if self.hparams.iterable_datasets:
-            dataset_class = TransformersIterableDataset
-        else:
-            dataset_class = TransformersMapDataset
-
-        if stage == 'fit' or stage is None:
-            if self.train_adapter is not None:
-                kwargs = {'start_from_step': self.hparams.skip_in_training} if self.hparams.iterable_datasets else {}
-                self.train_dataset = dataset_class(self.hparams, self.train_adapter, self.trainer, **kwargs)
-            if self.valid_adapter is not None:
-                self.valid_dataset = dataset_class(self.hparams, self.valid_adapter, self.trainer)
-
-            assert self.train_adapter is None or self.train_dataset is not None, (
-                f"Cannot specify `train_adapter` and then `train_dataset` is None: "
-                f"{self.train_adapter} and {self.train_dataset}"
-            )
-            assert self.valid_adapter is None or self.valid_dataset is not None, (
-                f"Cannot specify `valid_adapter` and then `valid_dataset` is None: "
-                f"{self.valid_adapter} and {self.valid_dataset}"
-            )
-
-        elif stage == 'test' or stage is None:
-            if self.test_adapter is not None:
-                if isinstance(self.test_adapter, SuperAdapter):
-                    self.test_dataset = dataset_class(self.hparams, self.test_adapter, self.trainer)
-                else:
-                    self.test_dataset = [
-                        dataset_class(self.hparams, adapter, self.trainer) for adapter in self.test_adapter
-                    ]
-
-            assert self.test_adapter is None or self.test_dataset is not None, (
-                f"Cannot specify `test_adapter` and then `test_dataset` is None: "
-                f"{self.test_adapter} and {self.test_dataset}"
-            )
-
-    def do_train(self):
-        return self.train_adapter is not None
-
-    def do_validation(self):
-        return self.valid_adapter is not None
-
-    def do_test(self):
-        return self.test_adapter is not None
-
-    def default_dataloader(self, dataset, batch_size, shuffle=False):
-        """ Return a dataloader with all usual default parameters. """
-        return DataLoader(
-            dataset,
-            batch_size=batch_size,
-            num_workers=self.hparams.num_workers,
-            pin_memory=True,
-            collate_fn=self.collate_fn,
-            shuffle=shuffle
-        )
-
-    def train_dataloader(self):
-        if self.train_dataset:
-            shuffle = not isinstance(self.train_dataset, IterableDataset)
-            return self.default_dataloader(self.train_dataset, self.hparams.batch_size, shuffle=shuffle)
-        return None
-
-    def val_dataloader(self):
-        if self.valid_adapter:
-            return self.default_dataloader(self.valid_dataset, self.hparams.val_batch_size, shuffle=False)
-        return None
-
-    def test_dataloader(self):
-        if self.test_adapter:
-            if isinstance(self.test_dataset, list):
-                return [
-                    self.default_dataloader(dataset, self.hparams.test_batch_size, shuffle=False)
-                    for dataset in self.test_dataset
-                ]
-            return self.default_dataloader(self.test_dataset, self.hparams.test_batch_size, shuffle=False)
-        return None
-
-    @staticmethod
-    def add_datamodule_specific_args(parser: ArgumentParser):
-        parser.add_argument(
-            '--num_workers',
-            required=False,
-            default=multiprocessing.cpu_count(),
-            type=int,
-            help='Number of workers to be used to load datasets'
-        )
-        parser.add_argument('--skip_in_training', type=int, default=None, required=False)
-        parser.add_argument('--batch_size', type=int, default=32)
-        parser.add_argument('--val_batch_size', type=int, default=256)
-        parser.add_argument('--test_batch_size', type=int, default=256)
-        parser.add_argument('--iterable_datasets', action="store_true")
-
-        tmp_args, _ = parser.parse_known_args()
-        assert tmp_args.skip_in_training is None or tmp_args.iterable_datasets, (
-            f"At the moment, `--skip_in_training <steps>` can be used only with `--iterable_datasets`"
-        )
-
-        return parser
diff --git a/transformers_lightning/datasets/README.md b/transformers_lightning/datasets/README.md
index b0d0978..97b6c22 100644
--- a/transformers_lightning/datasets/README.md
+++ b/transformers_lightning/datasets/README.md
@@ -1,75 +1,29 @@
 # Datasets
 
-`Datasets` are a convinient way to manage dataset and process and return the right data only when needed.
-This folder contains two main datasets: (i) a simple `MapDataset` that behaves exactly like a list and can be indexed and (ii) an `IterableDataset` that has no well-defined length and continues to preprocess and provide data until the corresponding adapter is completely emptied. The library provides also a simple `StackDataset` to emulate the `zip` command of `python`.
+`Datasets` are a convinient way to manage data and return the right row only when needed.
+This folder contains two main datasets: (i) a simple `MapDataset` that behaves exactly like a list and can be indexed and (ii) a `CompressedDataset` that store data in a `CompressedDictionary` and uses much less memory. It also can be indexes and the length is well-defined. It substitutes the old `IterableDataset` that cause many troubles. The library provides also a simple `StackDataset` to emulate the `zip` command of `python`.
+
+To preprocess each entry of the dataset, the `MapDataset` will call `adapter.preprocess_line` by default. `CompressedDataset`  is instead usually used with already-preprocessed data that are stored in a `CompressedDictionary`.
 
-To preprocess each entry of the dataset, the `Dataset` will call `adapter.preprocess_line` by default.
 
 ## MapDataset
 
-A `MapDataset` by default expects only an `Adapter` as input. I will completely read the `adapter` into memory and then it will provide primitives to read the dataset length and for integer indexing. When doing distributed training, the `Sampler` added by `PyTorch Lightning` will index the right data on each node, making life of the user extremely simple.
+A `MapDataset` by default expects only an `Adapter` as input. I will completely read the `adapter` into memory and then it will provide primitives to read the dataset length and for indexing. When doing distributed training, the `Sampler` added by `PyTorch Lightning` will index the right data on each node, making life of the user extremely simple.
 
 To create a MapDataset, you have only to to the following:
 ```python
+trainer = pl.Trainer(...)
 adapter = ExampleAdapter(hparams)
-dataset = MapDataset(hparams, adapter)
+dataset = MapDataset(hparams, adapter, trainer=trainer)
 ```
 
 Normally, you should not implement `Datasets` directly since they are automagically added by [`SuperDataModule`](/transformers-lightning/datamodules).
 
 
-## IterableDataset
-
-An `IterableDataset` is far more interesting. It contains the logic to read and preprocess a dataset on the fly while doing training by distributing the workload on many processes (by default as many processes as the number of logic threads of your CPU). Since it is not possible to know the dataset length in advance (as with generators and iterators), it includes the logic to provide different samples to all the different nodes of a distributed execution.
-
-In particular, the `IterableDataset` return the data in the following way (index are shown):
-
-Single node training (no 'ddp' neither 'dp'), `num_workers = 0` or `num_workers = 1`, check [`this`](https://pytorch.org/docs/stable/data.html) for more info about workers.
-```bash
->>> worker_pid 0 or 1: 0, 1, 2, 3, 4, 5, ..., length-1
-```
-
-Single node with many workers to preprocess data (`num_workers = 4`):
-```bash
->>> worker_pid 0: 0, 4, 8, 12, ...
->>> worker_pid 1: 1, 5, 9, 13, ...
->>> worker_pid 2: 2, 6, 10, 14, ...
->>> worker_pid 3: 3, 7, 11, 15, ...
-```
-
-Distributed training with `world_size = 2` and `num_workers = 4` for each node:
-```bash
->>> proc 0: 0, 2, 4, 6, 8, 10, ...
->>> # distributed in the following way across workers:
->>> proc 0, worker_pid 0: 0, 8, 16, 24, ...
->>> proc 0, worker_pid 1: 2, 10, 18, 26, ...
->>> proc 0, worker_pid 2: 4, 12, 20, 28, ...
->>> proc 0, worker_pid 3: 6, 14, 22, 30, ...
-
->>> proc 1: 1, 3, 5, 7, 9, 11, ...
->>> # distributed in the following way across workers:
->>> proc 1, worker_pid 0: 1, 9, 17, 25, ...
->>> proc 1, worker_pid 1: 3, 11, 19, 27, ...
->>> proc 1, worker_pid 2: 5, 13, 21, 29, ...
->>> proc 1, worker_pid 3: 7, 15, 23, 31, ...
-```
-
-Moreover, one must ensure that each node receives exactly the same number of data.
-This is not allowed and may lead to a crash in the distributed training:
-```
->>> proc 0: 0, 2, 4, 6, 8
->>> proc 1: 1, 3, 5, 7, 9, 11
-```
-This can be solved by reading at least `world_size` (2 in this case) elements for each iteration from the adapter. So be warned: if you provide an Adapter to the distributed session with a number of elements that is not and exact multiple of the `world_size`, a small quantity of data in the last batches will be dropped!
-
-The `IterableDataset` offers also a last **experimental** feature: jump the first `start_from_step` steps to resume exactly from a checkpoint. This will compute the effective `batch_size` by taking into account the number of nodes in distributed training and the `accumulate_grad_batches` parameter and will skip data read in the first `start_from_step` steps.
-
-```python
-adapter = ExampleAdapter(hparams)
-dataset = IterableDataset(hparams, adapter, start_from_step=10)
-```
+## CompressedDataset
 
-Please do not provide a `start_from_step` longer than an epoch!
+A `CompressedDataset` reads a `CompressedDictionary` from the disk and then behaves exactly like a `MapDataset` (indexing, length, ...).
+See [here](https://github.com/lucadiliello/compressed-dictionary) for more details.
 
 
 ## StackDataset
@@ -91,4 +45,4 @@ for data_1, data_2 in dataset:
     # data_2 will come from dataset_2
 ```
 
-`StackDataset` does not inherit from `SuperTransformersDataset`!
\ No newline at end of file
+`StackDataset` does not inherit from `SuperDataset`!
\ No newline at end of file
diff --git a/transformers_lightning/datasets/__init__.py b/transformers_lightning/datasets/__init__.py
index f4406ac..3ffd7cb 100644
--- a/transformers_lightning/datasets/__init__.py
+++ b/transformers_lightning/datasets/__init__.py
@@ -1,5 +1,5 @@
 # triples of name as key, yaml config file and data file as values
-from transformers_lightning.datasets.superdataset import SuperTransformersDataset
-from transformers_lightning.datasets.map_dataset import TransformersMapDataset
-from transformers_lightning.datasets.iterable_dataset import TransformersIterableDataset
-from transformers_lightning.datasets.stack_dataset import StackDataset
\ No newline at end of file
+from transformers_lightning.datasets.compressed_dataset import CompressedDataset
+from transformers_lightning.datasets.map_dataset import MapDataset
+from transformers_lightning.datasets.stack_dataset import StackDataset
+from transformers_lightning.datasets.super_dataset import SuperDataset
diff --git a/transformers_lightning/datasets/compressed_dataset.py b/transformers_lightning/datasets/compressed_dataset.py
new file mode 100644
index 0000000..a0ebc67
--- /dev/null
+++ b/transformers_lightning/datasets/compressed_dataset.py
@@ -0,0 +1,28 @@
+from compressed_dictionary import CompressedDictionary
+
+from transformers_lightning.datasets.super_dataset import SuperDataset
+
+
+class CompressedDataset(SuperDataset):
+    r"""
+    Superclass of all dataset using new `CompressedDictionary`s.
+    Dataset is completely read into memory in a compressed way and decompression is done on-the-fly.
+    """
+
+    def __init__(self, hparams, filepath: str):
+        super().__init__(hparams)
+        self.data = CompressedDictionary.load(filepath)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx) -> dict:
+        """ Get dict of data at a given position. """
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError("absolute value of index should not exceed dataset length")
+            idx = len(self) + idx
+
+        assert 0 <= idx < len(self), (f"Received index out of range {idx}, range: {0} <= idx < {len(self)}")
+
+        return self.data[idx]
\ No newline at end of file
diff --git a/transformers_lightning/datasets/iterable_dataset.py b/transformers_lightning/datasets/iterable_dataset.py
deleted file mode 100644
index 6b31cf7..0000000
--- a/transformers_lightning/datasets/iterable_dataset.py
+++ /dev/null
@@ -1,129 +0,0 @@
-from pytorch_lightning import _logger as logger
-import torch
-from types import GeneratorType
-
-from torch.utils.data import IterableDataset
-from transformers_lightning import utils
-from transformers_lightning.datasets import SuperTransformersDataset
-
-
-class TransformersIterableDataset(SuperTransformersDataset, IterableDataset):
-    r"""
-    Superclass of all iterable datasets. Pre-processing is performed on the fly.
-    Dataset is read on the fly from disk to save memory.
-
-    When doing distributed training, Lightning does not add a specific sampler 
-    for IterableDatasets: this means that we should implement the logic to
-    let each node have the right portion of data.
-
-    Example
-    no distributed training (idx)
-    >>> pid default: 0, 1, 2, 3, 4, 5, ..., get_length-1
-
-    with num_workers > 0: example with num_workers=4
-    >>> pid 0: 0, 4, 8, 12, ...
-    >>> pid 1: 1, 5, 9, 13, ...
-    >>> pid 2: 2, 6, 10, 14, ...
-    >>> pid 3: 3, 7, 11, 15, ...
-
-    in distributed training with world_size=2 and num_workers=4
-    >>> proc 0: 0, 2, 4, 6, 8, 10, ...
-    >>> proc 0, worker_pid 0: 0, 8, 16, 24, ...
-    >>> proc 0, worker_pid 1: 2, 10, 18, 26, ...
-    >>> proc 0, worker_pid 2: 4, 12, 20, 28, ...
-    >>> proc 0, worker_pid 3: 6, 14, 22, 30, ...
-
-    >>> proc 1: 1, 3, 5, 7, 9, 11, ...
-    >>> proc 1, worker_pid 0: 1, 9, 17, 25, ...
-    >>> proc 1, worker_pid 1: 3, 11, 19, 27, ...
-    >>> proc 1, worker_pid 2: 5, 13, 21, 29, ...
-    >>> proc 1, worker_pid 3: 7, 15, 23, 31, ...
-
-    Moreover, one must ensure that each node receives exactly the same number of data.
-    This is not allowed and may lead to a crash in the distributed training:
-    >>> proc 0: 0, 2, 4, 6, 8
-    >>> proc 1: 1, 3, 5, 7, 9, 11
-    This can be solved by reading at least `world_size` (2 in this case) elements for
-    each iteration from the adapter.
-    """
-
-    @property
-    def length(self):
-        r"""
-        Even if this is an IterableDataset, length may be computed by scrolling the document
-        without pre-processing in a fast way. However, this cannot be set as __len__ attribute because
-        in that way it may be misleaded for a normal index-based MapDataset.
-        """
-        if not hasattr(self, "_length"):
-            self._length = sum(1 for _ in iter(self.adapter))
-        return self._length
-
-    def __init__(self, *args, start_from_step=None):
-        super().__init__(*args)
-        r"""
-        If `start_from_step` is provided, this dataset will return data
-        relative to the `start_from_step`+1 effective step. This is not a problem
-        since the training algorithm does not know in advance the total dataset length.
-        This applied only to the first epoch, from the following all the data are provided.
-        Do not provide a `start_from_step` higher than the number of elements in this
-        dataset or higher than the total number of max_steps
-        """
-        self.start_from_step = None
-        if start_from_step is not None:
-            assert isinstance(start_from_step, int), (f"`start_from` must be integer, found {start_from_step}")
-
-            total_devices = utils.get_total_devices(trainer=self.trainer)
-            effective_batch_size = self.hparams.batch_size * self.hparams.accumulate_grad_batches * total_devices
-            self.start_from_step = effective_batch_size * start_from_step
-
-            logger.warning(
-                f"IterableDataset starting from step {start_from_step}. If this is the correct "
-                f"behaviour, please ignore this warning."
-            )
-
-    def __iter__(self):
-        r"""
-        Return the iterable by nesting different generators, each of which does a different
-        filtering based on the process id when in distributed training and on the worker id
-        if using also parallel loading in the dataloader.
-
-        0) eventually skip the first `start_from_step` rows to restart from a precise point.
-        1) utils.batch_filter simply ensures that at least `world_size` elements are read at a time
-        2) utils.filter_generator on distributed training to keep one element every `world_size`
-        3) utils.filter_generator on parallel workers processing to keep one element every `num_workers`
-        """
-        self.reader = iter(self.adapter)
-
-        # skip first steps if needed
-        if self.start_from_step is not None:
-            self.reader = utils.filter_generator(self.reader, step=1, offset=self.start_from_step)
-            self.start_from_step = None
-
-        # add distributed training logic
-        if torch.distributed.is_initialized():
-
-            world_size = torch.distributed.get_world_size()
-            rank = torch.distributed.get_rank()
-
-            self.reader = utils.batch_filter(self.reader, size=world_size)
-
-            self.reader = utils.filter_generator(self.reader, step=world_size, offset=rank)
-
-        # add parallel processing with workers logic
-        if torch.utils.data.get_worker_info() is not None:
-            self.reader = utils.filter_generator(
-                self.reader,
-                torch.utils.data.get_worker_info().num_workers,
-                torch.utils.data.get_worker_info().id
-            )
-
-        # pre-process data and return
-        for line in self.reader:
-
-            preprocessed_data = self.adapter.preprocess_line(line)
-
-            # if the adapter return a generator, yield an element per time
-            if isinstance(preprocessed_data, GeneratorType):
-                yield from preprocessed_data
-            else:
-                yield preprocessed_data
diff --git a/transformers_lightning/datasets/map_dataset.py b/transformers_lightning/datasets/map_dataset.py
index 26b4be0..cee18b2 100644
--- a/transformers_lightning/datasets/map_dataset.py
+++ b/transformers_lightning/datasets/map_dataset.py
@@ -1,16 +1,20 @@
-from torch.utils.data import Dataset
-from transformers_lightning.datasets import SuperTransformersDataset
+from pytorch_lightning import Trainer
 
+from transformers_lightning.adapters.super_adapter import SuperAdapter
+from transformers_lightning.datasets.super_dataset import SuperDataset
 
-class TransformersMapDataset(SuperTransformersDataset, Dataset):
+
+class MapDataset(SuperDataset):
     r"""
     Superclass of all map datasets. Tokenization is performed on the fly.
-    Dataset is completely read into memory. If you need some lighter way of doing training, see `IterableDataset`
+    Dataset is completely read into memory.
     """
 
-    def __init__(self, *args):
-        super().__init__(*args)
-        self.data = list(iter(self.adapter))
+    def __init__(self, hparams, adapter: SuperAdapter, trainer: Trainer):
+        super().__init__(hparams)
+        self.trainer = trainer
+        self.adapter = adapter
+        self.data = list(iter(adapter))
 
     def __len__(self):
         return len(self.data)
diff --git a/transformers_lightning/datasets/stack_dataset.py b/transformers_lightning/datasets/stack_dataset.py
index 698c74d..ea6c767 100644
--- a/transformers_lightning/datasets/stack_dataset.py
+++ b/transformers_lightning/datasets/stack_dataset.py
@@ -1,11 +1,12 @@
 import torch
 from pytorch_lightning.utilities.distributed import rank_zero_warn
+
 from transformers_lightning import utils
 
 
 class StackDataset(torch.utils.data.Dataset):
     r"""Dataset as a stack of multiple datasets. (parallel)
-    This class is useful to assemble different existing datasets.
+    This class is useful to assemble different existing datasets in parallel.
     Arguments:
         datasets (sequence): List of datasets to be stacked
     """
diff --git a/transformers_lightning/datasets/super_dataset.py b/transformers_lightning/datasets/super_dataset.py
new file mode 100644
index 0000000..1c42023
--- /dev/null
+++ b/transformers_lightning/datasets/super_dataset.py
@@ -0,0 +1,12 @@
+from argparse import Namespace
+
+from transformers_lightning.adapters.super_adapter import SuperAdapter
+
+
+class SuperDataset:
+    r"""
+    Superclass of all datasets. Only implements hparams storage. 
+    """
+
+    def __init__(self, hparams: Namespace):
+        self.hparams = hparams
diff --git a/transformers_lightning/datasets/superdataset.py b/transformers_lightning/datasets/superdataset.py
deleted file mode 100644
index f1e6c1a..0000000
--- a/transformers_lightning/datasets/superdataset.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from argparse import Namespace
-
-from pytorch_lightning import Trainer
-from transformers_lightning.adapters.super_adapter import SuperAdapter
-
-
-class SuperTransformersDataset:
-    r"""
-    Superclass of all datasets. Only implements default fields filling. 
-    """
-    hparams: Namespace = None
-    adapter: SuperAdapter = None
-    trainer: Trainer = None
-
-    def __init__(self, hparams: Namespace, adapter: SuperAdapter, trainer: Trainer):
-        r"""
-        Doing type check to avoid stupid errors related to arguments inversion
-        """
-        assert isinstance(hparams, Namespace), f"Argument `hparams` must be of type `Namespace`"
-        assert isinstance(adapter, SuperAdapter), f"Argument `adapter` must be of type `SuperAdapter`"
-        assert isinstance(trainer, Trainer), f"Argument `trainer` must be of type `Trainer`"
-
-        self.hparams = hparams
-        self.adapter = adapter
-        self.trainer = trainer
diff --git a/transformers_lightning/defaults/__init__.py b/transformers_lightning/defaults/__init__.py
index 646f50b..bde8529 100644
--- a/transformers_lightning/defaults/__init__.py
+++ b/transformers_lightning/defaults/__init__.py
@@ -9,7 +9,6 @@ class DefaultConfig:
     """
 
     output_dir: str = "outputs"
-    results_dir: str = "results"
     pre_trained_dir: str = "pre_trained_models"
     tensorboard_dir: str = "tensorboard"
     checkpoints_dir: str = "checkpoints"
diff --git a/transformers_lightning/language_modeling/masked_language_modeling.py b/transformers_lightning/language_modeling/masked_language_modeling.py
index 95674ef..20dd447 100644
--- a/transformers_lightning/language_modeling/masked_language_modeling.py
+++ b/transformers_lightning/language_modeling/masked_language_modeling.py
@@ -1,6 +1,6 @@
 from typing import Tuple
-import torch
 
+import torch
 import transformers
 
 from transformers_lightning.language_modeling import IGNORE_IDX, LanguageModel
diff --git a/transformers_lightning/language_modeling/random_token_substitution.py b/transformers_lightning/language_modeling/random_token_substitution.py
index 63b13ce..fd177e6 100644
--- a/transformers_lightning/language_modeling/random_token_substitution.py
+++ b/transformers_lightning/language_modeling/random_token_substitution.py
@@ -1,6 +1,6 @@
 from typing import Tuple
-import torch
 
+import torch
 import transformers
 
 from transformers_lightning.language_modeling import IGNORE_IDX, LanguageModel
diff --git a/transformers_lightning/language_modeling/utils.py b/transformers_lightning/language_modeling/utils.py
index be3adcd..d32d843 100644
--- a/transformers_lightning/language_modeling/utils.py
+++ b/transformers_lightning/language_modeling/utils.py
@@ -1,5 +1,4 @@
 import torch
-
 from transformers import PreTrainedTokenizer
 
 
diff --git a/transformers_lightning/models/transformers_model.py b/transformers_lightning/models/transformers_model.py
index 2df0a6b..675b0d0 100644
--- a/transformers_lightning/models/transformers_model.py
+++ b/transformers_lightning/models/transformers_model.py
@@ -3,11 +3,11 @@
 import torch
 from pytorch_lightning import LightningModule
 from pytorch_lightning import _logger as logger
+from pytorch_lightning.utilities.distributed import rank_zero_warn
 from transformers import AdamW
 
 from transformers_lightning import utils
-from transformers_lightning.schedulers.linear_scheduler_with_warmup import \
-    LinearSchedulerWithWarmup
+from transformers_lightning.schedulers.linear_scheduler_with_warmup import LinearSchedulerWithWarmup
 
 
 class TransformersModel(LightningModule):
@@ -97,6 +97,6 @@ def add_model_specific_args(parser: ArgumentParser):
 
         tmp_args, _ = parser.parse_known_args()
         if tmp_args.learning_rate > 1:
-            logger.warning(f"You specified a huge learning rate! Learning rate: {tmp_args.learning_rate}")
+            rank_zero_warn(f"You specified a huge learning rate! Learning rate: {tmp_args.learning_rate}")
 
         return parser
diff --git a/transformers_lightning/schedulers/__init__.py b/transformers_lightning/schedulers/__init__.py
index a50b91a..f81e406 100644
--- a/transformers_lightning/schedulers/__init__.py
+++ b/transformers_lightning/schedulers/__init__.py
@@ -1,7 +1,9 @@
-from transformers_lightning.schedulers.linear_scheduler_with_warmup import LinearSchedulerWithWarmup
-from transformers_lightning.schedulers.linear_scheduler import LinearScheduler
 from transformers_lightning.schedulers.constant_scheduler import ConstantScheduler
 from transformers_lightning.schedulers.constant_scheduler_with_warmup import ConstantSchedulerWithWarmup
 from transformers_lightning.schedulers.cosine_scheduler_with_warmup import CosineSchedulerWithWarmup
-from transformers_lightning.schedulers.cosine_scheduler_with_warmup_and_hard_restart import CosineSchedulerWithWarmupAndHardRestart
+from transformers_lightning.schedulers.cosine_scheduler_with_warmup_and_hard_restart import (
+    CosineSchedulerWithWarmupAndHardRestart,
+)
 from transformers_lightning.schedulers.layerwise_decay_scheduler import PolynomialLayerwiseDecaySchedulerWithWarmup
+from transformers_lightning.schedulers.linear_scheduler import LinearScheduler
+from transformers_lightning.schedulers.linear_scheduler_with_warmup import LinearSchedulerWithWarmup
diff --git a/transformers_lightning/schedulers/layerwise_decay_scheduler.py b/transformers_lightning/schedulers/layerwise_decay_scheduler.py
index 88299b6..3537479 100644
--- a/transformers_lightning/schedulers/layerwise_decay_scheduler.py
+++ b/transformers_lightning/schedulers/layerwise_decay_scheduler.py
@@ -19,8 +19,8 @@
 rates for different layers of the network.
 """
 import math
-
 import warnings
+
 from torch.optim.lr_scheduler import _LRScheduler
 
 
diff --git a/transformers_lightning/utils/__init__.py b/transformers_lightning/utils/__init__.py
index a22058d..6eede6d 100644
--- a/transformers_lightning/utils/__init__.py
+++ b/transformers_lightning/utils/__init__.py
@@ -1,7 +1,7 @@
+from transformers_lightning.utils.distributed import *
 from transformers_lightning.utils.functional import *
+from transformers_lightning.utils.generators import *
+from transformers_lightning.utils.gradient import *
 from transformers_lightning.utils.inspectors import *
+from transformers_lightning.utils.language_modeling import *
 from transformers_lightning.utils.readers import *
-from transformers_lightning.utils.gradient import *
-from transformers_lightning.utils.generators import *
-from transformers_lightning.utils.distributed import *
-from transformers_lightning.utils.language_modeling import *
\ No newline at end of file
diff --git a/transformers_lightning/utils/distributed.py b/transformers_lightning/utils/distributed.py
index f97db8c..ee247b0 100644
--- a/transformers_lightning/utils/distributed.py
+++ b/transformers_lightning/utils/distributed.py
@@ -2,8 +2,10 @@
 from argparse import Namespace
 
 import torch
-from pytorch_lightning import _logger as logger
 from pytorch_lightning import Trainer
+from pytorch_lightning import _logger as logger
+from pytorch_lightning.utilities.distributed import rank_zero_warn
+
 from transformers_lightning import utils
 
 
@@ -35,28 +37,13 @@ def compute_max_steps(hparams: Namespace, trainer: Trainer) -> int:
         return hparams.max_steps
 
     if not hasattr(trainer, 'datamodule'):
-        logger.warning(
-            "You tried to fix max_steps but didn't provide a datamodule to "
-            "the trainer.fit function, returning max_steps=None"
+        rank_zero_warn(
+            "You tried to fix `max_steps` but didn't provide a datamodule to "
+            "the trainer.fit function. Returning `max_steps=None`"
         )
         return None
 
-    # if cannot retrieve len of the dataset, skip
-    # this can happen with iterabledatasets
-    if not (
-        hasattr(trainer.datamodule.train_dataset, '__len__') or hasattr(trainer.datamodule.train_dataset, 'length')
-    ):
-        logger.warning("Cannot infer dataset length from datamodule.train_dataset, returning max_steps=None")
-        return None
-
-    try:
-        dataset_len = len(trainer.datamodule.train_dataset)
-    except:
-        try:
-            dataset_len = trainer.datamodule.train_dataset.length
-        except:
-            return None
-
+    dataset_len = len(trainer.datamodule.train_dataset)
     total_devices = utils.get_total_devices(trainer=trainer)
 
     num_training_batches = math.ceil(dataset_len / hparams.batch_size)
@@ -64,6 +51,6 @@ def compute_max_steps(hparams: Namespace, trainer: Trainer) -> int:
     steps_per_epoch = math.ceil(training_batches_per_epoch / hparams.accumulate_grad_batches)
     steps = hparams.max_epochs * steps_per_epoch
 
-    logger.warning(f"Automagically computed max_steps={steps}. If it appears to be OK, ignore this warning")
+    rank_zero_warn(f"Automagically computed max_steps={steps}. If it appears to be OK, ignore this warning")
 
     return steps


Test patch:
diff --git a/tests/adapters/test_csv_adapter.py b/tests/adapters/test_csv_adapter.py
index 2145097..f7a5290 100644
--- a/tests/adapters/test_csv_adapter.py
+++ b/tests/adapters/test_csv_adapter.py
@@ -1,7 +1,9 @@
 from argparse import Namespace
-from transformers_lightning.adapters import CSVAdapter
+
 import pytest
 
+from transformers_lightning.adapters import CSVAdapter
+
 
 class ExampleAdapter(CSVAdapter):
 
diff --git a/tests/adapters/test_transformers_adapter.py b/tests/adapters/test_transformers_adapter.py
index 63fdb26..a34820f 100644
--- a/tests/adapters/test_transformers_adapter.py
+++ b/tests/adapters/test_transformers_adapter.py
@@ -1,8 +1,9 @@
 from argparse import Namespace
 
+import pytest
 import transformers
+
 from transformers_lightning.adapters import TransformersAdapter
-import pytest
 
 
 class ExampleAdapter(TransformersAdapter):
diff --git a/tests/callbacks/test_ckpt_callback.py b/tests/callbacks/test_ckpt_callback.py
index 655279c..1c2d775 100644
--- a/tests/callbacks/test_ckpt_callback.py
+++ b/tests/callbacks/test_ckpt_callback.py
@@ -1,23 +1,23 @@
 import multiprocessing
-from argparse import Namespace
 import os
-from transformers_lightning.callbacks.transformers_model_checkpoint import TransformersModelCheckpointCallback
-from transformers_lightning.datamodules import SuperDataModule
-
-from transformers import BertTokenizer
 import shutil
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from tests.datamodules.test_utils import SimpleTransformerLikeModel, ExampleAdapter
+from transformers import BertTokenizer
+
+from tests.datamodules.test_utils import ExampleAdapter, SimpleTransformerLikeModel
+from transformers_lightning.callbacks.transformers_model_checkpoint import TransformersModelCheckpointCallback
+from transformers_lightning.datamodules import AdaptersDataModule
 
 n_cpus = multiprocessing.cpu_count()
-OUTPUT_DIR = "tests/output"
+OUTPUT_DIR = "/tmp/tests"
 if os.path.isdir(OUTPUT_DIR):
     shutil.rmtree(OUTPUT_DIR)
 
 
-class ExampleDataModule(SuperDataModule):
+class ExampleDataModule(AdaptersDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -67,7 +67,6 @@ def test_datamodule_cpu(epochs, accumulate_grad_batches, batch_size, callback_in
         max_steps=None,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         skip_in_training=None,
         checkpoint_interval=callback_interval,
         no_val_checkpointing=not val_callback,
diff --git a/tests/datamodules/test_datamodule_cpu.py b/tests/datamodules/test_datamodule_cpu.py
index 0d44408..1ac0b38 100644
--- a/tests/datamodules/test_datamodule_cpu.py
+++ b/tests/datamodules/test_datamodule_cpu.py
@@ -1,53 +1,36 @@
 import multiprocessing
-from argparse import Namespace
 import os
-from transformers import BertTokenizer
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from .test_utils import SimpleTransformerLikeModel, ExampleDataModule
+from transformers import BertTokenizer
+
+from .test_utils import ExampleDataModule, SimpleTransformerLikeModel
 
 n_cpus = multiprocessing.cpu_count()
 
 
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "gpus", "epochs"],
+    ["num_workers", "gpus", "epochs"],
     [
-
-    # ITER dataset
-    # test different num_workers in single node on cpu
-        ['iter', 0, 0, 1],
-        ['iter', 1, 0, 1],
-        ['iter', 2, 0, 1],
-        ['iter', n_cpus, 0, 1],
-
-    # num_workers through epochs
-        ['iter', 0, 0, 1],
-        ['iter', 0, 0, 4],
-        ['iter', 0, 0, 2],
-        ['iter', 0, 0, 10],
-        ['iter', 2, 0, 1],
-        ['iter', 2, 0, 2],
-        ['iter', 2, 0, 4],
-
-    # MAP dataset
     # test different num_workers in single node on cpu
-        ['map', 0, 0, 1],
-        ['map', 1, 0, 1],
-        ['map', 2, 0, 1],
-        ['map', n_cpus, 0, 1],
+        [0, 0, 1],
+        [1, 0, 1],
+        [2, 0, 1],
+        [n_cpus, 0, 1],
 
     # num_workers through epochs
-        ['map', 0, 0, 1],
-        ['map', 0, 0, 2],
-        ['map', 0, 0, 4],
-        ['map', 2, 0, 1],
-        ['map', 2, 0, 2],
-        ['map', 2, 0, 4],
+        [0, 0, 1],
+        [0, 0, 2],
+        [0, 0, 4],
+        [2, 0, 1],
+        [2, 0, 2],
+        [2, 0, 4],
     ]
 )
-def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
+def test_datamodule_cpu(num_workers, gpus, epochs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
@@ -62,7 +45,6 @@ def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
         max_steps=None,
         max_sequence_length=10,
         gpus=gpus,
-        iterable_datasets=ds_type == 'iter',
         skip_in_training=None
     )
 
diff --git a/tests/datamodules/test_datamodule_ddp.py b/tests/datamodules/test_datamodule_ddp.py
index 63b39a8..5dce13a 100644
--- a/tests/datamodules/test_datamodule_ddp.py
+++ b/tests/datamodules/test_datamodule_ddp.py
@@ -1,41 +1,32 @@
 import multiprocessing
-from argparse import Namespace
-import time
 import os
+import time
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
 import torch
-from .test_utils import SimpleTransformerLikeModel, ExampleDataModule
 from transformers import BertTokenizer
 
+from .test_utils import ExampleDataModule, SimpleTransformerLikeModel
+
 n_cpus = multiprocessing.cpu_count()
 
 
 # Test iter dataset work correctly with dp
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "distributed_backend", "gpus", "epochs", "dataset_idx"],
+    ["num_workers", "distributed_backend", "gpus", "epochs", "dataset_idx"],
     [
-
-    # ITER dataset
-    # num_workers with ddp
-        ['iter', 0, 'ddp', 2, 1, 1],
-        ['iter', 1, 'ddp', 2, 2, 1],
-        ['iter', 2, 'ddp', 2, 2, 1],
-        ['iter', 0, 'ddp', 2, 1, 2],
-        ['iter', n_cpus, 'ddp', 2, 5, 2],
-
-    # MAP dataset
     # num_workers with ddp
-        ['map', 0, 'ddp', 2, 2, 1],
-        ['map', 1, 'ddp', 2, 2, 1],
-        ['map', 2, 'ddp', 2, 2, 1],
-        ['map', 0, 'ddp', 2, 1, 2],
-        ['map', n_cpus, 'ddp', 2, 5, 3],
+        [0, 'ddp', 2, 2, 1],
+        [1, 'ddp', 2, 2, 1],
+        [2, 'ddp', 2, 2, 1],
+        [0, 'ddp', 2, 1, 2],
+        [n_cpus, 'ddp', 2, 5, 3],
     ]
 )
-def test_datamodule_gpu_ddp_only(ds_type, num_workers, distributed_backend, gpus, epochs, dataset_idx):
+def test_datamodule_gpu_ddp_only(num_workers, distributed_backend, gpus, epochs, dataset_idx):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
     time.sleep(2)    # sleep for 5 second to be sure area is clean
@@ -51,7 +42,6 @@ def test_datamodule_gpu_ddp_only(ds_type, num_workers, distributed_backend, gpus
         max_steps=None,
         max_sequence_length=10,
         gpus=gpus,
-        iterable_datasets=ds_type == 'iter',
         accelerator=distributed_backend,
         skip_in_training=None
     )
diff --git a/tests/datamodules/test_datamodule_dp.py b/tests/datamodules/test_datamodule_dp.py
index b2d4768..407952f 100644
--- a/tests/datamodules/test_datamodule_dp.py
+++ b/tests/datamodules/test_datamodule_dp.py
@@ -1,13 +1,14 @@
 import multiprocessing
-from argparse import Namespace
-import time
 import os
-from transformers import BertTokenizer
+import time
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
 import torch
-from .test_utils import SimpleTransformerLikeModel, ExampleDataModule
+from transformers import BertTokenizer
+
+from .test_utils import ExampleDataModule, SimpleTransformerLikeModel
 
 n_cpus = multiprocessing.cpu_count()
 
@@ -15,33 +16,20 @@
 # Test iter dataset work correctly with dp
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "distributed_backend", "gpus", "epochs"],
+    ["num_workers", "distributed_backend", "gpus", "epochs"],
     [
-
-    # ITER dataset
-    # num_workers with dp
-        ['iter', 0, 'dp', 2, 2],
-        ['iter', 1, 'dp', 2, 2],
-        ['iter', 2, 'dp', 2, 2],
-        ['iter', n_cpus, 'dp', 2, 2],
-        ['iter', 0, 'dp', 2, 1],
-        ['iter', 1, 'dp', 2, 2],
-        ['iter', 2, 'dp', 2, 4],
-        ['iter', n_cpus, 'dp', 2, 5],
-
-    # MAP dataset
     # num_workers with dp
-        ['map', 0, 'dp', 2, 2],
-        ['map', 1, 'dp', 2, 2],
-        ['map', 2, 'dp', 2, 2],
-        ['map', n_cpus, 'dp', 2, 2],
-        ['map', 0, 'dp', 2, 1],
-        ['map', 1, 'dp', 2, 2],
-        ['map', 2, 'dp', 2, 4],
-        ['map', n_cpus, 'dp', 2, 5]
+        [0, 'dp', 2, 2],
+        [1, 'dp', 2, 2],
+        [2, 'dp', 2, 2],
+        [n_cpus, 'dp', 2, 2],
+        [0, 'dp', 2, 1],
+        [1, 'dp', 2, 2],
+        [2, 'dp', 2, 4],
+        [n_cpus, 'dp', 2, 5]
     ]
 )
-def test_datamodule_gpu_dp(ds_type, num_workers, distributed_backend, gpus, epochs):
+def test_datamodule_gpu_dp(num_workers, distributed_backend, gpus, epochs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
     time.sleep(2)    # sleep for 5 second to be sure area is clean
@@ -57,7 +45,6 @@ def test_datamodule_gpu_dp(ds_type, num_workers, distributed_backend, gpus, epoc
         max_steps=None,
         max_sequence_length=10,
         gpus=gpus,
-        iterable_datasets=ds_type == 'iter',
         accelerator=distributed_backend,
         skip_in_training=None
     )
diff --git a/tests/datamodules/test_datamodule_val_test.py b/tests/datamodules/test_datamodule_val_test.py
index 52ddfeb..fd62e1c 100644
--- a/tests/datamodules/test_datamodule_val_test.py
+++ b/tests/datamodules/test_datamodule_val_test.py
@@ -1,17 +1,19 @@
 import multiprocessing
-from argparse import Namespace
 import os
-from transformers import BertTokenizer
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning import datamodules
-from .test_utils import SimpleTransformerLikeModel, ExampleAdapter
+from transformers import BertTokenizer
+
+from transformers_lightning.datamodules import AdaptersDataModule
+
+from .test_utils import ExampleAdapter, SimpleTransformerLikeModel
 
 n_cpus = multiprocessing.cpu_count()
 
 
-class ExampleDataModule(datamodules.SuperDataModule):
+class ExampleDataModule(AdaptersDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -29,29 +31,18 @@ def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
 
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "gpus", "epochs"],
+    ["num_workers", "gpus", "epochs"],
     [
-
-    # ITER dataset
-    # test different num_workers in single node on cpu
-        ['iter', 0, 0, 1],
-        ['iter', n_cpus, 0, 1],
-
-    # num_workers through epochs
-        ['iter', 0, 0, 1],
-        ['iter', 0, 0, 10],
-
-    # MAP dataset
     # test different num_workers in single node on cpu
-        ['map', 0, 0, 1],
-        ['map', n_cpus, 0, 1],
+        [0, 0, 1],
+        [n_cpus, 0, 1],
 
     # num_workers through epochs
-        ['map', 0, 0, 1],
-        ['map', 2, 0, 1],
+        [0, 0, 1],
+        [2, 0, 1],
     ]
 )
-def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
+def test_datamodule_cpu(num_workers, gpus, epochs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
@@ -66,7 +57,6 @@ def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
         max_steps=None,
         max_sequence_length=10,
         gpus=gpus,
-        iterable_datasets=ds_type == 'iter',
         skip_in_training=None
     )
 
diff --git a/tests/datamodules/test_utils.py b/tests/datamodules/test_utils.py
index d0373f8..ae5b4e2 100644
--- a/tests/datamodules/test_utils.py
+++ b/tests/datamodules/test_utils.py
@@ -1,11 +1,13 @@
 import torch
-from transformers.models.bert.modeling_bert import (BertConfig, BertForSequenceClassification)
 from transformers import AdamW
+from transformers.models.bert.modeling_bert import BertConfig, BertForSequenceClassification
+
 from transformers_lightning.adapters.csv_adapter import CSVAdapter
-from transformers_lightning import datamodules, models
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 
 
-class SimpleTransformerLikeModel(models.TransformersModel):
+class SimpleTransformerLikeModel(TransformersModel):
 
     def __init__(self, hparams, do_ids_check=True):
         super().__init__(hparams)
@@ -101,7 +103,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class ExampleDataModule(datamodules.SuperDataModule):
+class ExampleDataModule(AdaptersDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
diff --git a/tests/datasets/test_start_from_steps.py b/tests/datasets/test_start_from_steps.py
deleted file mode 100644
index 521da6f..0000000
--- a/tests/datasets/test_start_from_steps.py
+++ /dev/null
@@ -1,148 +0,0 @@
-import multiprocessing
-from argparse import Namespace
-
-from transformers import AdamW
-from transformers_lightning import utils
-from tests import adapters
-
-import pytest
-import pytorch_lightning as pl
-import torch
-from transformers_lightning import datamodules, models, adapters
-
-n_cpus = multiprocessing.cpu_count()
-N = 100
-
-
-class SimpleTransformerLikeModel(models.TransformersModel):
-
-    def __init__(self, hparams):
-        super().__init__(hparams)
-        self.lin = torch.nn.Linear(10, 10)
-
-    def configure_optimizers(self):
-        self.computed_max_steps = utils.compute_max_steps(self.hparams, self.trainer)
-        return AdamW(self.lin.parameters())
-
-    def training_step(self, batch, batch_idx):
-        results = self.lin(batch['data']).mean()
-        return {'loss': results, 'ids': batch['ids']}
-
-    def training_epoch_end(self, outputs):
-        self.all_received_ids = torch.cat([output['ids'] for output in outputs])
-
-        if self.trainer.distributed_backend == "ddp":
-            gather_ids = [torch.ones_like(self.all_received_ids) for _ in range(torch.distributed.get_world_size())]
-            torch.distributed.all_gather(gather_ids, self.all_received_ids)
-            self.all_received_ids = torch.cat(gather_ids, dim=0)
-
-        self.all_received_ids = self.all_received_ids.tolist()
-
-
-class ExampleAdapter(adapters.SuperAdapter):
-
-    def __init__(self, hparams):
-        super().__init__(hparams)
-
-    def __iter__(self):
-        for i in range(N):
-            yield (i, [1.0] * 10, torch.LongTensor(1))
-
-    def preprocess_line(self, line: list) -> list:
-        res = {'ids': line[0], 'data': line[1], 'labels': line[2]}
-        return res
-
-
-class ExampleDataModule(datamodules.SuperDataModule):
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.train_adapter = ExampleAdapter(self.hparams)
-
-
-# Test if skip_steps works correctly
-@pytest.mark.parametrize(
-    ["max_epochs", "accumulate_grad_batches", "batch_size", "skip", "expected_ids"], [
-        [1, 1, 4, 4, list(range(16, 100))],
-        [1, 3, 8, 3, list(range(72, 100))],
-        [1, 2, 5, 5, list(range(50, 100))],
-    ]
-)
-def test_skip_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, skip, expected_ids):
-
-    hparams = Namespace(
-        batch_size=batch_size,
-        accumulate_grad_batches=accumulate_grad_batches,
-        num_workers=5,
-        output_dir='output',
-        max_epochs=max_epochs,
-        max_sequence_length=10,
-        gpus=None,
-        max_steps=None,
-        iterable_datasets=True,
-        skip_in_training=skip
-    )
-
-    # instantiate PL trainer
-    trainer = pl.Trainer.from_argparse_args(
-        hparams,
-        profiler='simple',
-        logger=None,
-        callbacks=[],
-    )
-
-    # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)
-
-    # Datasets
-    datamodule = ExampleDataModule(hparams)
-
-    # Train!
-    trainer.fit(model, datamodule=datamodule)
-    assert sorted(model.all_received_ids
-                 ) == sorted(expected_ids), f"{sorted(model.all_received_ids)} != {sorted(expected_ids)}"
-
-
-# Test if skip_steps works correctly in ddp
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
-@pytest.mark.parametrize(
-    ["max_epochs", "accumulate_grad_batches", "batch_size", "skip", "expected_ids"], [
-        [1, 1, 4, 4, list(range(32, 100))],
-        [1, 3, 8, 2, list(range(96, 100))],
-        [1, 2, 7, 2, list(range(28, 100))],
-    ]
-)
-def test_skip_steps_gpu(max_epochs, accumulate_grad_batches, batch_size, skip, expected_ids):
-
-    hparams = Namespace(
-        batch_size=batch_size,
-        accumulate_grad_batches=accumulate_grad_batches,
-        num_workers=5,
-        output_dir='output',
-        max_epochs=max_epochs,
-        max_sequence_length=10,
-        gpus=2,
-        max_steps=None,
-        iterable_datasets=True,
-        skip_in_training=skip,
-        accelerator='ddp'
-    )
-
-    # instantiate PL trainer
-    trainer = pl.Trainer.from_argparse_args(
-        hparams,
-        profiler='simple',
-        logger=None,
-        callbacks=[],
-    )
-
-    # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)
-
-    # Datasets
-    datamodule = ExampleDataModule(hparams)
-
-    # Train!
-    trainer.fit(model, datamodule=datamodule)
-    assert sorted(model.all_received_ids
-                 ) == sorted(expected_ids), f"{sorted(model.all_received_ids)} != {sorted(expected_ids)}"
diff --git a/tests/language_models/test_mlm.py b/tests/language_models/test_mlm.py
index ca18af9..940d45e 100644
--- a/tests/language_models/test_mlm.py
+++ b/tests/language_models/test_mlm.py
@@ -1,8 +1,8 @@
+import pytest
 import torch
 from transformers import BertTokenizer
 
-import pytest
-from transformers_lightning.language_modeling.masked_language_modeling import MaskedLanguageModeling, IGNORE_IDX
+from transformers_lightning.language_modeling.masked_language_modeling import IGNORE_IDX, MaskedLanguageModeling
 from transformers_lightning.language_modeling.utils import whole_word_tails_mask
 
 tok = BertTokenizer.from_pretrained('bert-base-cased')
diff --git a/tests/language_models/test_rts.py b/tests/language_models/test_rts.py
index 15680ae..80995fd 100644
--- a/tests/language_models/test_rts.py
+++ b/tests/language_models/test_rts.py
@@ -1,7 +1,7 @@
+import pytest
 import torch
 from transformers import BertTokenizer
 
-import pytest
 from transformers_lightning.language_modeling import IGNORE_IDX
 from transformers_lightning.language_modeling.random_token_substitution import RandomTokenSubstitution
 from transformers_lightning.language_modeling.utils import whole_word_tails_mask
diff --git a/tests/models/test_fix_max_steps.py b/tests/models/test_fix_max_steps.py
index 2818f17..2b07c69 100644
--- a/tests/models/test_fix_max_steps.py
+++ b/tests/models/test_fix_max_steps.py
@@ -1,21 +1,23 @@
 import multiprocessing
-from argparse import Namespace
-
-from transformers import AdamW
-from tests import adapters
 import time
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
 import torch
-from transformers_lightning import datamodules, models, adapters, utils
-from transformers import BertTokenizer
-from transformers.models.bert.modeling_bert import (BertConfig, BertForSequenceClassification)
+from transformers import AdamW, BertTokenizer
+from transformers.models.bert.modeling_bert import BertConfig, BertForSequenceClassification
+
+from tests import adapters
+from transformers_lightning import adapters, datamodules, models, utils
+from transformers_lightning.adapters import CSVAdapter
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 
 n_cpus = multiprocessing.cpu_count()
 
 
-class SimpleTransformerLikeModel(models.TransformersModel):
+class SimpleTransformerLikeModel(TransformersModel):
 
     def __init__(self, hparams):
         super().__init__(hparams)
@@ -34,7 +36,7 @@ def training_step(self, batch, batch_idx):
         return results.loss
 
 
-class ExampleAdapter(adapters.CSVAdapter):
+class ExampleAdapter(CSVAdapter):
 
     def __init__(self, *args, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -54,7 +56,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class ExampleDataModule(datamodules.SuperDataModule):
+class ExampleDataModule(AdaptersDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -65,18 +67,18 @@ def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
 
 # Test if max_steps fix works correctly
 @pytest.mark.parametrize(
-    ["max_epochs", "accumulate_grad_batches", "batch_size", "expected_max_steps", "ds_type"], [
-        [1, 1, 4, 10, 'map'],
-        [1, 3, 8, 2, 'map'],
-        [4, 2, 12, 8, 'map'],
-        [4, 4, 16, 4, 'map'],
-        [1, 1, 4, 10, 'iter'],
-        [1, 3, 8, 2, 'iter'],
-        [4, 4, 16, 4, 'iter'],
-        [4, 2, 12, 8, 'iter'],
+    ["max_epochs", "accumulate_grad_batches", "batch_size", "expected_max_steps"], [
+        [1, 1, 4, 10],
+        [1, 3, 8, 2],
+        [4, 2, 12, 8],
+        [4, 4, 16, 4],
+        [1, 1, 4, 10],
+        [1, 3, 8, 2],
+        [4, 4, 16, 4],
+        [4, 2, 12, 8],
     ]
 )
-def test_fix_max_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, expected_max_steps, ds_type):
+def test_fix_max_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, expected_max_steps):
 
     hparams = Namespace(
         batch_size=batch_size,
@@ -87,7 +89,6 @@ def test_fix_max_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, expe
         max_steps=None,
         max_sequence_length=10,
         gpus=None,
-        iterable_datasets=ds_type == 'iter',
         skip_in_training=None
     )
 
@@ -139,7 +140,6 @@ def test_fix_max_steps_gpu(max_epochs, accumulate_grad_batches, batch_size, dist
         max_steps=None,
         max_sequence_length=10,
         gpus=2,
-        iterable_datasets=False,
         skip_in_training=None,
         accelerator=distributed_backend
     )
diff --git a/tests/schedulers/test_constant_scheduler.py b/tests/schedulers/test_constant_scheduler.py
index 2486e77..65a5471 100644
--- a/tests/schedulers/test_constant_scheduler.py
+++ b/tests/schedulers/test_constant_scheduler.py
@@ -1,14 +1,15 @@
-from argparse import Namespace
 import os
-import torch
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning.models import TransformersModel
+import torch
+from transformers import AdamW
+
 from transformers_lightning.adapters import SuperAdapter
-from transformers_lightning.datamodules import SuperDataModule
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 from transformers_lightning.schedulers import ConstantScheduler
-from transformers import AdamW
 
 
 class FakeModel(TransformersModel):
@@ -55,7 +56,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class FakeDataModule(SuperDataModule):
+class FakeDataModule(AdaptersDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -80,7 +81,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
         last_epoch=last_epoch,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         log_every_n_steps=1,
         learning_rate=1.0,
     )
diff --git a/tests/schedulers/test_constant_scheduler_with_warmup.py b/tests/schedulers/test_constant_scheduler_with_warmup.py
index 20d76fa..865edbc 100644
--- a/tests/schedulers/test_constant_scheduler_with_warmup.py
+++ b/tests/schedulers/test_constant_scheduler_with_warmup.py
@@ -1,14 +1,15 @@
-from argparse import Namespace
 import os
-import torch
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning.models import TransformersModel
+import torch
+from transformers import AdamW
+
 from transformers_lightning.adapters import SuperAdapter
-from transformers_lightning.datamodules import SuperDataModule
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 from transformers_lightning.schedulers import ConstantSchedulerWithWarmup
-from transformers import AdamW
 
 
 class FakeModel(TransformersModel):
@@ -55,7 +56,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class FakeDataModule(SuperDataModule):
+class FakeDataModule(AdaptersDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -79,7 +80,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
         last_epoch=last_epoch,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         log_every_n_steps=1,
         learning_rate=1.0,
     )
diff --git a/tests/schedulers/test_cosine_scheduler_with_warmup.py b/tests/schedulers/test_cosine_scheduler_with_warmup.py
index 17cc839..f97eea7 100644
--- a/tests/schedulers/test_cosine_scheduler_with_warmup.py
+++ b/tests/schedulers/test_cosine_scheduler_with_warmup.py
@@ -1,14 +1,15 @@
-from argparse import Namespace
 import os
-import torch
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning.models import TransformersModel
+import torch
+from transformers import AdamW
+
 from transformers_lightning.adapters import SuperAdapter
-from transformers_lightning.datamodules import SuperDataModule
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 from transformers_lightning.schedulers import CosineSchedulerWithWarmup
-from transformers import AdamW
 
 
 class FakeModel(TransformersModel):
@@ -61,7 +62,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class FakeDataModule(SuperDataModule):
+class FakeDataModule(AdaptersDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -93,7 +94,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
         last_epoch=last_epoch,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         log_every_n_steps=1,
         learning_rate=1.0,
     )
diff --git a/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py b/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py
index 532ea3a..0e5ba2a 100644
--- a/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py
+++ b/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py
@@ -1,14 +1,15 @@
-from argparse import Namespace
 import os
-import torch
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning.models import TransformersModel
+import torch
+from transformers import AdamW
+
 from transformers_lightning.adapters import SuperAdapter
-from transformers_lightning.datamodules import SuperDataModule
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 from transformers_lightning.schedulers import CosineSchedulerWithWarmupAndHardRestart
-from transformers import AdamW
 
 
 class FakeModel(TransformersModel):
@@ -61,7 +62,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class FakeDataModule(SuperDataModule):
+class FakeDataModule(AdaptersDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -93,7 +94,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
         last_epoch=last_epoch,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         log_every_n_steps=1,
         learning_rate=1.0,
     )
diff --git a/tests/schedulers/test_linear_scheduler.py b/tests/schedulers/test_linear_scheduler.py
index ab623ac..3d23730 100644
--- a/tests/schedulers/test_linear_scheduler.py
+++ b/tests/schedulers/test_linear_scheduler.py
@@ -1,14 +1,15 @@
-from argparse import Namespace
 import os
-import torch
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning.models import TransformersModel
+import torch
+from transformers import AdamW
+
 from transformers_lightning.adapters import SuperAdapter
-from transformers_lightning.datamodules import SuperDataModule
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 from transformers_lightning.schedulers import LinearScheduler
-from transformers import AdamW
 
 
 class FakeModel(TransformersModel):
@@ -56,7 +57,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class FakeDataModule(SuperDataModule):
+class FakeDataModule(AdaptersDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -88,7 +89,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
         last_epoch=last_epoch,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         log_every_n_steps=1,
         learning_rate=1.0,
     )
diff --git a/tests/schedulers/test_linear_scheduler_with_warmup.py b/tests/schedulers/test_linear_scheduler_with_warmup.py
index 79ceb18..9ae4998 100644
--- a/tests/schedulers/test_linear_scheduler_with_warmup.py
+++ b/tests/schedulers/test_linear_scheduler_with_warmup.py
@@ -1,14 +1,15 @@
-from argparse import Namespace
 import os
-import torch
+from argparse import Namespace
 
 import pytest
 import pytorch_lightning as pl
-from transformers_lightning.models import TransformersModel
+import torch
+from transformers import AdamW
+
 from transformers_lightning.adapters import SuperAdapter
-from transformers_lightning.datamodules import SuperDataModule
+from transformers_lightning.datamodules import AdaptersDataModule
+from transformers_lightning.models import TransformersModel
 from transformers_lightning.schedulers import LinearSchedulerWithWarmup
-from transformers import AdamW
 
 
 class FakeModel(TransformersModel):
@@ -59,7 +60,7 @@ def preprocess_line(self, line: list) -> list:
         return res
 
 
-class FakeDataModule(SuperDataModule):
+class FakeDataModule(AdaptersDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -89,7 +90,6 @@ def test_datamodule_cpu(num_warmup_steps, num_training_steps, last_epoch, expect
         last_epoch=last_epoch,
         max_sequence_length=10,
         gpus=0,
-        iterable_datasets=False,
         log_every_n_steps=1,
         learning_rate=1.0,
     )