Repository: litus-ai/classy (language: python)

Issue:
T5 support and bug fixes



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 180f6ced..9fc30248 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -24,8 +24,7 @@ jobs:
 
       - name: Generate pdoc output & Build Docusaurus website
         run: |
-          pip install -r requirements.txt
-          pip install -r <(python classy/optional_deps.py)
+          pip install -e .[all]
           pdoc -f --template-dir docs/pdoc/templates -o docs/docs classy
           mv docs/docs/classy docs/docs/api
           python docs/pdoc/pdoc_postprocess.py
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 2fe6968c..fb348f0b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -8,7 +8,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.3.0
+    rev: v4.1.0
     hooks:
       - id: check-yaml
       - id: check-case-conflict
@@ -31,13 +31,13 @@ repos:
         files: ^classy/
         args: ["--profile", "black"]
   - repo: https://github.com/psf/black
-    rev: 21.12b0
+    rev: 22.3.0
     hooks:
       - id: black
         name: Format code
         files: ^classy/
   - repo: https://github.com/asottile/blacken-docs
-    rev: v1.12.0
+    rev: v1.12.1
     hooks:
       - id: blacken-docs
         name: Format code in docs
diff --git a/Dockerfile b/Dockerfile
index 84587837..dbdd00ea 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -22,7 +22,7 @@ RUN conda update -y conda && conda init
 WORKDIR /classy
 COPY . .
 RUN \
-    bash -c "source ~/miniconda3/etc/profile.d/conda.sh && printf 'classy\n3.8\n1.9.0\n11.3\nN\n' | bash setup.sh"
+    bash -c "source ~/miniconda3/etc/profile.d/conda.sh && printf 'classy\n3.8\n1.10.2\n11.3\nN\n' | bash setup.sh"
 
 # standard cmd
 CMD [ "/bin/bash" ]
diff --git a/README.md b/README.md
index f6341c7a..cf3e1486 100644
--- a/README.md
+++ b/README.md
@@ -19,7 +19,7 @@
     </a>
 -->
     <a href="">
-        <img alt="Python" src="https://img.shields.io/badge/Python 3.7--3.9-blue?style=for-the-badge&logo=python&logoColor=white">
+        <img alt="Python" src="https://img.shields.io/badge/Python 3.8+-blue?style=for-the-badge&logo=python&logoColor=white">
     </a>
     <a href="https://pypi.org/project/classy-core/">
         <img alt="PyPI" src="https://img.shields.io/pypi/v/classy-core?style=for-the-badge&logo=pypi">
@@ -28,7 +28,7 @@
         <img alt="PyTorch" src="https://img.shields.io/badge/PyTorch 1.8+-ee4c2c?style=for-the-badge&logo=pytorch&logoColor=white">
     </a>
     <a href="https://pytorchlightning.ai/">
-        <img alt="Lightning" src="https://img.shields.io/badge/Lightning 1.4.9-792ee5?style=for-the-badge&logo=pytorchlightning&logoColor=white">
+        <img alt="Lightning" src="https://img.shields.io/badge/Lightning 1.5-792ee5?style=for-the-badge&logo=pytorchlightning&logoColor=white">
     </a>
     <a href="https://hydra.cc/">
         <img alt="Config: hydra" src="https://img.shields.io/badge/config-hydra 1.1.1-89b8cd?style=for-the-badge&labelColor=gray">
diff --git a/classy/data/data_drivers.py b/classy/data/data_drivers.py
index 6ea39ddf..c24be642 100644
--- a/classy/data/data_drivers.py
+++ b/classy/data/data_drivers.py
@@ -1,13 +1,26 @@
 import functools
 import json
 from pathlib import Path
-from typing import Any, Dict, Generator, Iterator, List, Optional, Tuple, Union
+from typing import Any, Dict, Generator, Iterator, List, Optional, Tuple, Type, Union
 
 from classy.utils.log import get_project_logger
 
 logger = get_project_logger(__name__)
 
 
+# TASK TYPES
+SEQUENCE = "sequence"
+SENTENCE_PAIR = "sentence-pair"
+TOKEN = "token"
+QA = "qa"
+GENERATION = "generation"
+
+
+# FILE EXTENSIONS
+TSV = "tsv"
+JSONL = "jsonl"
+
+
 class ClassySample:
     reference_annotation: Optional
     predicted_annotation: Optional
@@ -300,6 +313,29 @@ def input(self) -> str:
 
 
 class DataDriver:
+    _registry: Dict[Tuple[str, str], Type["DataDriver"]] = {}
+
+    @classmethod
+    def register(cls, task: str, extension: str):
+        def _register(data_driver: Type["DataDriver"]):
+            if (task, extension) in DataDriver._registry:
+                raise ValueError(
+                    f"Pair ({task}, {extension}) already registered as a DataDriver: found {DataDriver._registry[task, extension]}"
+                )
+            DataDriver._registry[task, extension] = data_driver
+            return data_driver
+
+        return _register
+
+    @classmethod
+    def from_task_and_extension(cls, task, extension):
+        if (task, extension) not in DataDriver._registry:
+            raise ValueError(
+                f"Extension '{extension}' does not appear to be supported for task {task}. "
+                f"Supported extensions are: {[e for t, e in DataDriver._registry.keys() if t == task]}"
+            )
+        return DataDriver._registry[task, extension]()
+
     def dataset_exists_at_path(self, path: str) -> bool:
         return Path(path).exists()
 
@@ -388,6 +424,7 @@ def save(
         raise NotImplementedError
 
 
+@DataDriver.register(SENTENCE_PAIR, TSV)
 class TSVSentencePairDataDriver(SentencePairDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[SentencePairSample]:
         for line in lines:
@@ -416,6 +453,7 @@ def save(
                 f.write(f"{sample.sentence1}\t{sample.sentence2}\t{used_annotation}\n")
 
 
+@DataDriver.register(SENTENCE_PAIR, JSONL)
 class JSONLSentencePairDataDriver(SentencePairDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[SentencePairSample]:
         for line in lines:
@@ -447,6 +485,7 @@ def save(
                 )
 
 
+@DataDriver.register(SEQUENCE, TSV)
 class TSVSequenceDataDriver(SequenceDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[SequenceSample]:
         for line in lines:
@@ -475,6 +514,7 @@ def save(
                 f.write(f"{sample.sequence}\t{used_annotation}\n")
 
 
+@DataDriver.register(SEQUENCE, JSONL)
 class JSONLSequenceDataDriver(SequenceDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[SequenceSample]:
         for line in lines:
@@ -505,6 +545,7 @@ def save(
                 )
 
 
+@DataDriver.register(TOKEN, TSV)
 class TSVTokensDataDriver(TokensDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[TokensSample]:
         for line in lines:
@@ -540,6 +581,7 @@ def save(
                 )
 
 
+@DataDriver.register(TOKEN, JSONL)
 class JSONLTokensDataDriver(TokensDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[TokensSample]:
         for line in lines:
@@ -578,6 +620,7 @@ def save(
                 f.write(json.dumps(d) + "\n")
 
 
+@DataDriver.register(QA, TSV)
 class TSVQADataDriver(QADataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[QASample]:
         for i, line in enumerate(lines):
@@ -614,6 +657,7 @@ def save(
                 f.write("\n")
 
 
+@DataDriver.register(QA, JSONL)
 class JSONLQADataDriver(QADataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[QASample]:
         for line in lines:
@@ -649,6 +693,7 @@ def save(
                 )
 
 
+@DataDriver.register(GENERATION, TSV)
 class TSVGenerationDataDriver(GenerationDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[GenerationSample]:
         previous_line_parts = None
@@ -692,6 +737,7 @@ def save(
                     f.write(f"{sample.source_sequence}\t{used_annotation}\n")
 
 
+@DataDriver.register(GENERATION, JSONL)
 class JSONLGenerationDataDriver(GenerationDataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[GenerationSample]:
         for line in lines:
@@ -724,40 +770,6 @@ def save(
                 )
 
 
-# TASK TYPES
-SEQUENCE = "sequence"
-SENTENCE_PAIR = "sentence-pair"
-TOKEN = "token"
-QA = "qa"
-GENERATION = "generation"
-
-# FILE EXTENSIONS
-TSV = "tsv"
-JSONL = "jsonl"
-
-READERS_DICT = {
-    (SEQUENCE, TSV): TSVSequenceDataDriver,
-    (SENTENCE_PAIR, TSV): TSVSentencePairDataDriver,
-    (SEQUENCE, JSONL): JSONLSequenceDataDriver,
-    (SENTENCE_PAIR, JSONL): JSONLSentencePairDataDriver,
-    (TOKEN, JSONL): JSONLTokensDataDriver,
-    (TOKEN, TSV): TSVTokensDataDriver,
-    (QA, TSV): TSVQADataDriver,
-    (QA, JSONL): JSONLQADataDriver,
-    (GENERATION, TSV): TSVGenerationDataDriver,
-    (GENERATION, JSONL): JSONLGenerationDataDriver,
-}
-
-
 @functools.lru_cache(maxsize=1_000)
-def get_data_driver(task_type: str, file_extension: str) -> DataDriver:
-    reader_identifier = (task_type, file_extension)
-    if reader_identifier not in READERS_DICT:
-        logger.info(
-            f"No reader available for task {task_type} and extension {file_extension}."
-        )
-    assert reader_identifier in READERS_DICT, (
-        f"Extension '{file_extension}' does not appear to be supported for task {task_type}. "
-        f"Supported extensions are: {[e for t, e in READERS_DICT.keys() if t == task_type]}"
-    )
-    return READERS_DICT[reader_identifier]()
+def get_data_driver(task: str, extension: str) -> DataDriver:
+    return DataDriver.from_task_and_extension(task, extension)
diff --git a/classy/data/data_modules.py b/classy/data/data_modules.py
index be68da1c..071e94e9 100644
--- a/classy/data/data_modules.py
+++ b/classy/data/data_modules.py
@@ -1,16 +1,19 @@
 import itertools
 import os
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
 import hydra.utils
 import omegaconf
 import pytorch_lightning as pl
+from omegaconf import OmegaConf
 from torch.utils.data import DataLoader
 
 from classy.data.data_drivers import DataDriver, get_data_driver
 from classy.utils.data import create_data_dir, shuffle_and_store_dataset, split_dataset
 from classy.utils.log import get_project_logger
+from classy.utils.train_coordinates import load_bundle
 from classy.utils.vocabulary import Vocabulary
 
 logger = get_project_logger(__name__)
@@ -20,12 +23,171 @@ def path_if_exists(path: str, data_driver: DataDriver) -> Optional[str]:
     return path if data_driver.dataset_exists_at_path(path) else None
 
 
+@dataclass
+class TrainCoordinates:
+    main_file_extension: str
+    main_data_driver: DataDriver
+    train_bundle: Dict[str, DataDriver]
+    validation_bundle: Optional[Dict[str, DataDriver]]
+    test_bundle: Optional[Dict[str, DataDriver]]
+
+
+def load_coordinates(coordinates_path: str, task: str) -> TrainCoordinates:
+    """
+    Computes the train coordinates of a training process.
+    Args:
+        coordinates_path:
+            a path to
+                - a file containing the training coordinates (check the documentation for more info)
+                - a single file containing the whole dataset to be split in train / dev
+                - directory containing two (three) files for train, validation (and test)
+        task:
+            one of the supported tasks in classy (e.g. sentence-pair)
+    Returns:
+        train_coordinates (TrainCoordinates): the train_coordinates containing
+         all the info on the datasets involved in the training.
+    """
+
+    train_coordinates = TrainCoordinates(None, None, None, None, None)
+
+    # If the "data" directory exists (and so this a resume from a previous run)
+    # or the "coordinates_path" point to a directory we have to retrieve the data
+    # from there.
+    if Path("data/").exists() or Path(coordinates_path).is_dir():
+
+        def scan_dir_for_file(
+            dir_path: str, file_name: str
+        ) -> Optional[Tuple[str, str, DataDriver]]:
+
+            if not Path(dir_path).is_dir():
+                return None
+
+            files_in_dir = os.listdir(dir_path)
+            matching_files = [fp for fp in files_in_dir if file_name in fp]
+
+            if len(matching_files) == 1:
+                matching_file = matching_files[0]
+                file_extension = matching_file.split(".")[-1]
+                data_driver = get_data_driver(task, file_extension)
+                return f"{dir_path}/{matching_file}", file_extension, data_driver
+
+            return None
+
+        previous_run_dir = "data/"
+
+        # check if a previous run stored a train file
+        train_scan_output = scan_dir_for_file(previous_run_dir, "train")
+        # if not retrieve it from the directory at the coordinates path
+        if train_scan_output is None:
+            train_scan_output = scan_dir_for_file(coordinates_path, "train")
+
+        assert train_scan_output is not None, "The training file could not be found."
+
+        # Load training data
+        train_file, train_file_extension, train_data_driver = train_scan_output
+        train_coordinates.main_file_extension = train_file_extension
+        train_coordinates.main_data_driver = train_data_driver
+        train_coordinates.train_bundle = {train_file: train_data_driver}
+
+        # check if the previous run stored a validation file
+        validation_scan_output = scan_dir_for_file(previous_run_dir, "validation")
+        # if not retrieve it from the directory at the coordinates path
+        if validation_scan_output is None:
+            validation_scan_output = scan_dir_for_file(coordinates_path, "validation")
+        if validation_scan_output is not None:
+            validation_file, _, validation_data_driver = validation_scan_output
+            train_coordinates.validation_bundle = {
+                validation_file: validation_data_driver
+            }
+
+        # check if the previous run stored a test file
+        test_scan_output = scan_dir_for_file(previous_run_dir, "test")
+        # if not retrieve it from the directory at the coordinates path
+        if test_scan_output is None:
+            test_scan_output = scan_dir_for_file(coordinates_path, "test")
+        if test_scan_output is not None:
+            test_file, _, test_data_driver = test_scan_output
+            train_coordinates.test_bundle = {test_file: test_data_driver}
+
+    # If the coordinates_path points to a file, it is either a yaml
+    # file containing the paths the datasets or a single file that
+    # we have to split in train and dev.
+    if Path(coordinates_path).is_file():
+
+        if coordinates_path.split(".")[-1] == "yaml":
+
+            coordinates_dict = OmegaConf.load(coordinates_path)
+
+            # this is not a resume from a previous run
+            if train_coordinates.train_bundle is None:
+
+                assert (
+                    "train_dataset" in coordinates_dict
+                ), "The coordinates file must specify the 'train_dataset' field"
+
+                # assign the main file extension if specified in the config
+                train_coordinates.main_file_extension = coordinates_dict.get(
+                    "main_file_extension", None
+                )
+
+                # train_bundle
+                if train_coordinates.main_file_extension is None:
+                    logger.info(
+                        "No main extension declared in coordinates file,"
+                        " inferring it from training datasets"
+                    )
+                    (
+                        train_coordinates.train_bundle,
+                        train_coordinates.main_file_extension,
+                    ) = load_bundle(
+                        coordinates_dict.get("train_dataset"),
+                        task,
+                        compute_main_extension=True,
+                    )
+                    logger.info(
+                        f"Main extension found: {train_coordinates.main_file_extension}"
+                    )
+                else:
+                    train_coordinates.train_bundle = load_bundle(
+                        coordinates_dict.get("train_dataset"),
+                        task,
+                    )
+                train_coordinates.main_data_driver = get_data_driver(
+                    task, train_coordinates.main_file_extension
+                )
+
+            if train_coordinates.validation_bundle is None:
+                # validation_bundle
+                train_coordinates.validation_bundle = load_bundle(
+                    coordinates_dict.get("validation_dataset"), task
+                )
+
+            if train_coordinates.test_bundle is None:
+                # test_bundle
+                train_coordinates.test_bundle = load_bundle(
+                    coordinates_dict.get("test_dataset"), task
+                )
+
+        else:
+            # just one file that will later be split in train and dev
+            if train_coordinates.train_bundle is None:
+                train_coordinates.main_file_extension = coordinates_path.split(".")[-1]
+                train_coordinates.main_data_driver = get_data_driver(
+                    task, train_coordinates.main_file_extension
+                )
+                train_coordinates.train_bundle = {
+                    coordinates_path: train_coordinates.main_data_driver
+                }
+
+    return train_coordinates
+
+
 class ClassyDataModule(pl.LightningDataModule):
     def __init__(
         self,
         task: str,
         dataset_path: str,
-        train_dataset: omegaconf.DictConfig,
+        dataset: omegaconf.DictConfig,
         validation_dataset: Optional[omegaconf.DictConfig] = None,
         test_dataset: Optional[omegaconf.DictConfig] = None,
         validation_split_size: Optional[float] = None,
@@ -37,18 +199,16 @@ def __init__(
         super().__init__()
         self.task = task
         self.dataset_path = dataset_path
-        self.file_extension = None
-        self.data_driver = None
-
-        self.train_path, self.validation_path, self.test_path = None, None, None
-        self.train_dataset, self.validation_dataset, self.test_dataset = (
-            None,
-            None,
-            None,
-        )
-        self.train_dataset_conf = train_dataset
-        self.validation_dataset_conf = validation_dataset or train_dataset
-        self.test_dataset_conf = test_dataset or validation_dataset or train_dataset
+
+        self.train_dataset_conf = dataset
+        self.validation_dataset_conf = validation_dataset or dataset
+        self.test_dataset_conf = test_dataset or validation_dataset or dataset
+
+        self.train_coordinates: TrainCoordinates = None
+
+        self.train_dataset = None
+        self.validation_dataset = None
+        self.test_dataset = None
 
         self.validation_split_size = validation_split_size
         self.test_split_size = test_split_size
@@ -58,220 +218,123 @@ def __init__(
         self.external_vocabulary_path = external_vocabulary_path
         self.vocabulary = None
 
-        # If the data directory exists this run is being resumed from a previous one.
-        # If the previous run stored some data we have to load them instead of recomputing
-        # the splits or the shuffling.
-        if Path("data/").exists():
-            files_in_dir = os.listdir("data/")
-
-            # check if the previous run stored a train file
-            possible_train_paths = [fp for fp in files_in_dir if "train" in fp]
-            if len(possible_train_paths) == 1:
-                self.train_path = possible_train_paths[0]
-
-            # check if the previous run stored a validation file
-            possible_validation_paths = [
-                fp for fp in files_in_dir if "validation" in fp
-            ]
-            if len(possible_validation_paths) == 1:
-                self.validation_path = possible_validation_paths[0]
-
-            # check if the previous run stored a test file
-            possible_test_paths = [fp for fp in files_in_dir if "test" in fp]
-            if len(possible_test_paths) == 1:
-                self.test_path = possible_test_paths[0]
-
     def get_examples(self, n: int) -> Tuple[str, List]:
-        source = self.test_path or self.validation_path
+        source = (
+            self.train_coordinates.test_bundle
+            or self.train_coordinates.validation_bundle
+        )
+        dataset_path, data_driver = list(source.items())[0]
         assert source is not None
-        return "test" if self.test_path is not None else "validation", list(
-            itertools.islice(self.data_driver.read_from_path(source), n)
+        return (
+            "test" if self.train_coordinates.test_bundle is not None else "validation",
+            list(itertools.islice(data_driver.read_from_path(dataset_path), n)),
         )
 
+    def build_vocabulary(self) -> None:
+        # todo: can we improve it?
+        if Path("vocabulary/").exists():
+            logger.info("Loading vocabulary from previous run")
+            self.vocabulary = Vocabulary.from_folder("vocabulary/")
+        elif self.external_vocabulary_path is not None:
+            logger.info("Loading vocabulary from external passed directory")
+            self.vocabulary = Vocabulary.from_folder(self.external_vocabulary_path)
+            self.vocabulary.save("vocabulary")
+        else:
+            self.vocabulary = hydra.utils.instantiate(
+                self.train_dataset_conf,
+                path=self.train_coordinates.train_bundle,
+            ).vocabulary
+            if self.vocabulary is not None:
+                self.vocabulary.save("vocabulary")
+
     def prepare_data(self) -> None:
+        train_coordinates = load_coordinates(self.dataset_path, self.task)
 
-        # TODO: we should improve the flow of this code
-        if (
-            self.train_path is not None
-            and self.validation_path is not None
-            and self.test_path is not None
-        ):
-            logger.info(
-                "Using train dev and test splits produced by the run being resumed"
+        shuffled_train_dataset_path = (
+            f"data/train.shuffled.{train_coordinates.main_file_extension}"
+        )
+        must_shuffle_dataset = (
+            self.shuffle_dataset
+            and not train_coordinates.main_data_driver.dataset_exists_at_path(
+                shuffled_train_dataset_path
             )
-        elif Path(
-            self.dataset_path
-        ).is_dir():  # the user provided a directory containing the datasets
-            dir_train_files = [
-                fp for fp in os.listdir(self.dataset_path) if "train" in fp
-            ]
-
-            assert (
-                len(dir_train_files) == 1
-            ), f"Expected one file with 'train' in its name, but {len(dir_train_files)} were found in {self.dataset_path}: {dir_train_files}"
-
-            train_file = dir_train_files[0]
-            self.file_extension = train_file.split(".")[-1]
-            self.data_driver = get_data_driver(self.task, self.file_extension)
-
-            if (
-                self.train_path is None
-            ):  # does not belong to the train shuffling of a resume run
-                self.train_path = path_if_exists(
-                    os.path.join(self.dataset_path, f"train.{self.file_extension}"),
-                    self.data_driver,
-                )
-
-            if self.validation_path is None:
-                self.validation_path = path_if_exists(
-                    os.path.join(
-                        self.dataset_path, f"validation.{self.file_extension}"
-                    ),
-                    self.data_driver,
-                )
+        )
 
-            self.test_path = path_if_exists(
-                os.path.join(self.dataset_path, f"test.{self.file_extension}"),
-                self.data_driver,
+        if must_shuffle_dataset:
+            # create data folder
+            create_data_dir()
+            logger.info(
+                f"Shuffling training dataset. The shuffled dataset "
+                f"will be stored at: {os.getcwd()}/{shuffled_train_dataset_path}"
             )
-
-            assert (
-                self.train_path is not None
-            ), f"Cannot find the training file '{self.train_path}'"
-
-            must_shuffle_dataset = (
-                self.shuffle_dataset
-                and not self.data_driver.dataset_exists_at_path(
-                    f"data/train.shuffled.{self.file_extension}"
-                )
-                and not self.data_driver.dataset_exists_at_path(
-                    f"data/dataset.shuffled.{self.file_extension}"
-                )
+            shuffle_and_store_dataset(
+                train_coordinates.train_bundle,
+                train_coordinates.main_data_driver,
+                output_path=shuffled_train_dataset_path,
             )
+            train_coordinates.train_bundle = {
+                shuffled_train_dataset_path: train_coordinates.main_data_driver
+            }
 
-            if must_shuffle_dataset:
-                # create data folder
-                create_data_dir()
-                # shuffle input dataset
-                shuffled_dataset_path = f"data/train.shuffled.{self.file_extension}"
-                logger.info(
-                    f"Shuffling training dataset. The shuffled dataset "
-                    f"will be stored at: {os.getcwd()}/{shuffled_dataset_path}"
-                )
-                shuffle_and_store_dataset(
-                    self.train_path, self.data_driver, output_path=shuffled_dataset_path
-                )
-                self.train_path = shuffled_dataset_path
-
-            if self.validation_path is None:
-                logger.info(
-                    f"Validation dataset not found: splitting the training dataset "
-                    f"(split_size: {1 - self.validation_split_size} / {self.validation_split_size})"
-                    f"enforcing a maximum of {self.max_nontrain_split_size} instances on validation dataset"
-                )
-
-                # if we must split the shuffled train dataset in two, then we must change its name
-                if must_shuffle_dataset:
-                    shuffled_dataset_path = (
-                        f"data/dataset.shuffled.{self.file_extension}"
-                    )
-                    os.rename(self.train_path, shuffled_dataset_path)
-                    self.train_path = shuffled_dataset_path  # will be modified in the next lines of code
-
-                self.train_path, self.validation_path, _ = split_dataset(
-                    self.train_path,
-                    self.data_driver,
-                    "data/",  # hydra takes care of placing this folder within the appropriate folder
-                    validation_split_size=self.validation_split_size,
-                    data_max_split=self.max_nontrain_split_size,
-                    shuffle=False,
-                )
-                logger.info(
-                    f"Storing the newly created datasets at '{self.train_path}' and '{self.validation_path}'"
-                )
-
-        else:  # the user provided just one file that must be split in train, dev and test
-            self.file_extension = self.dataset_path.split(".")[-1]
-            self.data_driver = get_data_driver(self.task, self.file_extension)
-
-            if self.shuffle_dataset and not self.data_driver.dataset_exists_at_path(
-                f"data/dataset.shuffled.{self.file_extension}"
-            ):
-                # create data folder
-                create_data_dir()
-                # shuffle training dataset
-                shuffled_dataset_path = f"data/dataset.shuffled.{self.file_extension}"
-                logger.info(
-                    f"Shuffling input dataset. The shuffled dataset "
-                    f"will be stored at: {os.getcwd()}/{shuffled_dataset_path}"
-                )
-                shuffle_and_store_dataset(
-                    self.dataset_path,
-                    self.data_driver,
-                    output_path=shuffled_dataset_path,
-                )
-                self.dataset_path = shuffled_dataset_path
-
-            # splitting dataset in train, validation and test
+        if train_coordinates.validation_bundle is None:
             logger.info(
-                "Splitting the dataset in train, validation and test. "
-                f"(split_size: {1 - self.validation_split_size - self.test_split_size} "
-                f"/ {self.validation_split_size}, {self.test_split_size}) "
-                f"enforcing a maximum of {self.max_nontrain_split_size} instances on non-train splits"
+                f"Validation dataset not found: splitting the training dataset "
+                f"(split_size: {1 - self.validation_split_size} / {self.validation_split_size})"
+                f"enforcing a maximum of {self.max_nontrain_split_size} instances on validation dataset"
             )
 
-            self.train_path, self.validation_path, self.test_path = split_dataset(
-                self.dataset_path,
-                self.data_driver,
+            # if we must split the shuffled train dataset in two, then we must change its name
+            if must_shuffle_dataset:
+                shuffled_dataset_path = (
+                    f"data/dataset.shuffled.{train_coordinates.main_file_extension}"
+                )
+                os.rename(shuffled_train_dataset_path, shuffled_dataset_path)
+                train_coordinates.train_bundle = {
+                    shuffled_dataset_path: train_coordinates.main_data_driver
+                }
+
+            # split and assign
+            (
+                train_coordinates.train_bundle,
+                train_coordinates.validation_bundle,
+                _,
+            ) = split_dataset(
+                train_coordinates.train_bundle,
+                train_coordinates.main_data_driver,
+                train_coordinates.main_file_extension,
                 "data/",  # hydra takes care of placing this folder within the appropriate folder
                 validation_split_size=self.validation_split_size,
-                test_split_size=self.test_split_size,
                 data_max_split=self.max_nontrain_split_size,
                 shuffle=False,
             )
-
             logger.info(
-                f"Storing the newly created datasets at '{self.train_path}', "
-                f"'{self.validation_path}'and '{self.test_path}'"
+                f"Storing the newly created datasets at"
+                f" '{list(train_coordinates.train_bundle.keys())[0]}'"
+                f" and '{list(train_coordinates.validation_bundle.keys())[0]}'"
             )
 
-        # todo: can we improve it?
-        if Path("vocabulary/").exists():
-            logger.info("Loading vocabulary from previous run")
-            self.vocabulary = Vocabulary.from_folder("vocabulary/")
-        elif self.external_vocabulary_path is not None:
-            logger.info("Loading vocabulary from external passed directory")
-            self.vocabulary = Vocabulary.from_folder(self.external_vocabulary_path)
-            self.vocabulary.save("vocabulary")
-        else:
-            self.vocabulary = hydra.utils.instantiate(
-                self.train_dataset_conf,
-                path=self.train_path,
-                data_driver=self.data_driver,
-            ).vocabulary
-            if self.vocabulary is not None:
-                self.vocabulary.save("vocabulary")
+        self.train_coordinates = train_coordinates
+
+        # build vocabulary
+        self.build_vocabulary()
 
     def setup(self, stage: Optional[str] = None) -> None:
 
         if stage == "fit":
             self.train_dataset = hydra.utils.instantiate(
                 self.train_dataset_conf,
-                path=self.train_path,
-                data_driver=self.data_driver,
+                path=self.train_coordinates.train_bundle,
                 vocabulary=self.vocabulary,
             )
             self.validation_dataset = hydra.utils.instantiate(
                 self.validation_dataset_conf,
-                path=self.validation_path,
-                data_driver=self.data_driver,
+                path=self.train_coordinates.validation_bundle,
                 vocabulary=self.vocabulary,
             )
         if stage == "test":
             self.test_dataset = hydra.utils.instantiate(
                 self.test_dataset_conf,
-                path=self.test_path,
-                data_driver=self.data_driver,
+                path=self.train_coordinates.test_bundle,
                 vocabulary=self.vocabulary,
             )
 
diff --git a/classy/data/dataset/base.py b/classy/data/dataset/base.py
index 0ad44c7a..827c007e 100644
--- a/classy/data/dataset/base.py
+++ b/classy/data/dataset/base.py
@@ -1,19 +1,15 @@
+import copy
+import itertools
+import re
 from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, Union
 
 import numpy as np
 import torch
+from omegaconf import DictConfig
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import IterableDataset
 
-from classy.data.data_drivers import (
-    ClassySample,
-    DataDriver,
-    GenerationSample,
-    QASample,
-    SentencePairSample,
-    SequenceSample,
-    TokensSample,
-)
+from classy.data.data_drivers import ClassySample, DataDriver
 from classy.utils.commons import add_noise_to_value, chunks, flatten
 from classy.utils.log import get_project_logger
 from classy.utils.vocabulary import Vocabulary
@@ -44,19 +40,54 @@ def requires_vocab() -> bool:
     def fit_vocabulary(samples: Iterator[ClassySample]) -> Vocabulary:
         raise NotImplementedError
 
+    @classmethod
+    def adapt_dataset_from(cls, training_dataset: DictConfig, setting: str):
+        if setting == "validation":
+            validation_dataset = copy.deepcopy(training_dataset)
+            validation_dataset["materialize"] = True
+            validation_dataset["for_inference"] = True
+            return validation_dataset
+        elif setting == "prediction":
+            prediction_dataset = copy.deepcopy(training_dataset)
+            prediction_dataset["_target_"] = re.sub(
+                ".from_file$", ".from_samples", prediction_dataset["_target_"]
+            )
+            prediction_dataset["min_length"] = -1
+            prediction_dataset["max_length"] = -1
+            prediction_dataset["for_inference"] = True
+            return prediction_dataset
+        else:
+            raise ValueError(
+                f"Setting {setting} not supported. Choose between [validation, prediction] or change config."
+            )
+
     @classmethod
     def from_file(
-        cls, path: str, data_driver: DataDriver, vocabulary: Vocabulary = None, **kwargs
+        cls,
+        path: Union[str, Dict[str, DataDriver]],
+        data_driver: Optional[DataDriver] = None,
+        vocabulary: Vocabulary = None,
+        **kwargs,
     ) -> "BaseDataset":
 
+        dataset_bundle: Dict[str, Any] = (
+            {path: data_driver} if type(path) == str else path
+        )
+
         if vocabulary is None and cls.requires_vocab():
             # vocabulary fitting here
             logger.info("Fitting vocabulary")
-            vocabulary = cls.fit_vocabulary(data_driver.read_from_path(path))
+            vocabulary = cls.fit_vocabulary(
+                itertools.chain(
+                    *[dd.read_from_path(p) for p, dd in dataset_bundle.items()]
+                )
+            )
             logger.info("Vocabulary fitting completed")
 
         return cls(
-            samples_iterator=lambda: data_driver.read_from_path(path),
+            samples_iterator=lambda: itertools.chain(
+                *[dd.read_from_path(p) for p, dd in dataset_bundle.items()]
+            ),
             vocabulary=vocabulary,
             **kwargs,
         )
diff --git a/classy/data/dataset/hf/base.py b/classy/data/dataset/hf/base.py
index 1a5e4a08..5128ca22 100644
--- a/classy/data/dataset/hf/base.py
+++ b/classy/data/dataset/hf/base.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from transformers import AutoTokenizer
 
@@ -15,10 +15,14 @@ def __init__(
         self,
         transformer_model: str,
         additional_special_tokens: Optional[List[str]] = None,
+        truncation: Union[bool, str] = False,
+        max_length: int = -1,
         **kwargs,
     ):
         if "tokenizer" not in self._shared_state:
-            self._shared_state["tokenizer"] = AutoTokenizer.from_pretrained(
+            self._shared_state[
+                "tokenizer", tuple(additional_special_tokens or [])
+            ] = AutoTokenizer.from_pretrained(
                 transformer_model,
                 use_fast=True,
                 add_prefix_space=True,
@@ -26,7 +30,12 @@ def __init__(
                 if additional_special_tokens is not None
                 else None,
             )
-        self.tokenizer = self._shared_state["tokenizer"]
+        self.tokenizer = self._shared_state[
+            "tokenizer", tuple(additional_special_tokens or [])
+        ]
+        self.transformer_model = transformer_model
+        self.truncation = truncation
+        self.additional_special_tokens = additional_special_tokens
         batching_fields = (
             kwargs.pop("batching_fields")
             if "batching_fields" in kwargs
@@ -35,6 +44,9 @@ def __init__(
         super().__init__(
             fields_batchers=self.init_fields_batcher(),
             batching_fields=batching_fields,
+            max_length=max_length
+            if max_length != -1
+            else self.tokenizer.model_max_length,
             **kwargs,
         )
 
diff --git a/classy/data/dataset/hf/classification.py b/classy/data/dataset/hf/classification.py
index b92d7d40..5dc503f8 100644
--- a/classy/data/dataset/hf/classification.py
+++ b/classy/data/dataset/hf/classification.py
@@ -36,9 +36,15 @@ def init_fields_batcher(self) -> Dict:
     def dataset_iterator_func(self) -> Iterable[Dict[str, Any]]:
 
         for sequence_sample in self.samples_iterator():
-            input_ids = self.tokenizer(sequence_sample.sequence, return_tensors="pt")[
-                "input_ids"
-            ][0]
+            input_ids = self.tokenizer(
+                sequence_sample.sequence,
+                return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
+            )["input_ids"][0]
             elem_dict = {
                 "input_ids": input_ids,
                 "attention_mask": torch.ones_like(input_ids),
@@ -87,14 +93,15 @@ def dataset_iterator_func(self) -> Iterable[Dict[str, Any]]:
                 "token_offsets": token_offsets,
             }
             if token_sample.reference_annotation is not None:
-                elem_dict["labels"] = torch.tensor(
-                    [
-                        self.vocabulary.get_idx(
-                            k="labels", elem=token_sample.reference_annotation[idx]
-                        )
-                        for idx in token_sample.target
-                    ]
-                )
+                # -100 == cross entropy ignore index
+                labels = [-100] * len(token_sample.reference_annotation)
+
+                for target_idx in token_sample.target:
+                    labels[target_idx] = self.vocabulary.get_idx(
+                        k="labels", elem=token_sample.reference_annotation[target_idx]
+                    )
+
+                elem_dict["labels"] = torch.tensor(labels, dtype=input_ids.dtype)
 
             elem_dict["samples"] = token_sample
             yield elem_dict
@@ -103,7 +110,14 @@ def tokenize(
         self, tokens: List[str]
     ) -> Optional[Tuple[torch.Tensor, List[Tuple[int, int]]]]:
         tok_encoding = self.tokenizer.encode_plus(
-            tokens, return_tensors="pt", is_split_into_words=True
+            tokens,
+            return_tensors="pt",
+            is_split_into_words=True,
+            **(
+                {"truncation": True, "max_length": self.max_length}
+                if self.truncation
+                else {}
+            ),
         )
         try:
             return tok_encoding.input_ids.squeeze(0), [
@@ -128,6 +142,11 @@ def dataset_iterator_func(self) -> Iterable[Dict[str, Any]]:
                 sequence_sample.sentence1,
                 sequence_sample.sentence2,
                 return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
             )
 
             elem_dict = {
@@ -170,6 +189,11 @@ def dataset_iterator_func(self) -> Iterable[Dict[str, Any]]:
                 qa_sample.context,
                 return_offsets_mapping=True,
                 return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
             )
 
             elem_dict = {
diff --git a/classy/data/dataset/hf/generation.py b/classy/data/dataset/hf/generation.py
index 8581ef00..81205dfd 100644
--- a/classy/data/dataset/hf/generation.py
+++ b/classy/data/dataset/hf/generation.py
@@ -1,37 +1,21 @@
 import collections
-import re
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Generator,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Union,
-)
+from typing import Any, Callable, Dict, Generator, Iterator, List, Optional
 
 import torch
-from transformers import AutoTokenizer, MBartTokenizerFast
+from omegaconf import DictConfig
+from transformers import MBartTokenizerFast
 from transformers.models.mbart.tokenization_mbart_fast import FAIRSEQ_LANGUAGE_CODES
 
-from classy.data.data_drivers import (
-    ClassySample,
-    GenerationSample,
-    QASample,
-    SentencePairSample,
-    SequenceSample,
-    TokensSample,
-)
-from classy.data.dataset.base import BaseDataset, batchify
+from classy.data.data_drivers import ClassySample
+from classy.data.dataset.base import batchify
+from classy.data.dataset.hf.base import HFBaseDataset
 from classy.utils.log import get_project_logger
 from classy.utils.vocabulary import Vocabulary
 
 logger = get_project_logger(__name__)
 
 
-class HFGenerationDataset(BaseDataset):
+class HFGenerationBaseDataset(HFBaseDataset):
     @staticmethod
     def requires_vocab() -> bool:
         return False
@@ -40,130 +24,50 @@ def requires_vocab() -> bool:
     def fit_vocabulary(samples: Iterator[ClassySample]) -> Vocabulary:
         raise NotImplementedError
 
-    def __init__(
-        self,
-        samples_iterator: Callable[[], Iterator[ClassySample]],
-        vocabulary: Vocabulary,
-        transformer_model: str,
-        tokens_per_batch: int,
-        max_batch_size: Optional[int],
-        section_size: int,
-        prebatch: bool,
-        materialize: bool,
-        min_length: int,
-        max_length: int,
-        for_inference: bool,
-        additional_special_tokens: Optional[List[str]] = None,
-    ):
-        self.sample_encoder = HFGenerationSampleEncoder.from_transformer_model(
-            transformer_model, additional_special_tokens
-        )
+    @classmethod
+    def adapt_dataset_from(cls, training_dataset: DictConfig, setting: str):
+        dataset = super().adapt_dataset_from(training_dataset, setting)
+        if setting == "prediction":
+            dataset["teacher_forcing"] = False
+        return dataset
+
+    def __init__(self, teacher_forcing: Optional[bool] = True, **kwargs):
+        self.teacher_forcing = teacher_forcing
         super().__init__(
-            samples_iterator=samples_iterator,
-            vocabulary=vocabulary,
-            tokens_per_batch=tokens_per_batch,
-            max_batch_size=max_batch_size,
-            fields_batchers=self.sample_encoder.get_fields_batcher(),
-            section_size=section_size,
-            prebatch=prebatch,
-            materialize=materialize,
-            min_length=min_length,
-            max_length=max_length
-            if max_length != -1
-            else self.sample_encoder.get_tokenizer_max_length(),
-            for_inference=for_inference,
-            batching_fields=self.sample_encoder.get_batching_fields(for_inference),
+            batching_fields=self.get_batching_fields(self.teacher_forcing), **kwargs
         )
 
-    def dataset_iterator_func(self) -> Iterable[Dict[str, Any]]:
-        for generation_sample in self.samples_iterator():
-            yield self.sample_encoder.sample2elem_dict(
-                generation_sample, inference_mode=self.for_inference
-            )
+    def get_batching_fields(self, teacher_forcing: bool):
+        raise NotImplementedError
 
     def materialize_batches(
         self, dataset_elements: List[Dict[str, Any]]
     ) -> Generator[Dict[str, Any], None, None]:
-        for group in self.sample_encoder.group_elements_on_materializations(
-            dataset_elements, inference_mode=self.for_inference
+        for group in self.group_elements_on_materializations(
+            dataset_elements, teacher_forcing=self.teacher_forcing
         ):
             yield from super().materialize_batches(group)
 
-
-class HFGenerationSampleEncoder:
-    @classmethod
-    def from_transformer_model(
-        cls, transformer_model: str, additional_special_tokens: Optional[List[str]]
-    ):
-        if re.fullmatch("facebook/bart-(base|large)", transformer_model):
-            return BartHFGenerationSampleEncoder(
-                transformer_model, additional_special_tokens
-            )
-        elif transformer_model.startswith("gpt2"):
-            return GPT2HFGenerationSampleEcnoder(
-                transformer_model, additional_special_tokens
-            )
-        elif re.fullmatch("facebook/mbart-large-(cc25|50)", transformer_model):
-            return MBARTHFGenerationSampleEncoder(
-                transformer_model, additional_special_tokens
-            )
-        else:
-            raise ValueError
-
-    _shared_state = {}
-
-    def __init__(
-        self, transformer_model: str, additional_special_tokens: Optional[List[str]]
-    ):
-        if "tokenizer" not in self._shared_state:
-            self._shared_state["tokenizer"] = {}
-        if transformer_model not in self._shared_state["tokenizer"]:
-            self._shared_state["tokenizer"][
-                transformer_model
-            ] = AutoTokenizer.from_pretrained(
-                transformer_model,
-                additional_special_tokens=list(additional_special_tokens)
-                if additional_special_tokens is not None
-                else None,
-                use_fast=True,
-                add_prefix_space=True,
-            )
-        self.tokenizer = self._shared_state["tokenizer"][transformer_model]
-
-    def get_batching_fields(self, inference_mode: bool) -> List[str]:
-        raise NotImplementedError
-
-    def get_tokenizer_max_length(self) -> int:
-        return self.tokenizer.model_max_length
-
-    def get_fields_batcher(self) -> Dict[str, Callable]:
-        raise NotImplementedError
-
-    def sample2elem_dict(
-        self, sample: GenerationSample, inference_mode: bool
-    ) -> Dict[str, Any]:
-        raise NotImplementedError
-
     def group_elements_on_materializations(
-        self, dataset_elements: List[Dict[str, Any]], inference_mode: bool
+        self, dataset_elements: List[Dict[str, Any]], teacher_forcing: bool
     ) -> List[List[Dict[str, Any]]]:
         return [dataset_elements]
 
 
-class EncDecHFGenerationSampleEncoder(HFGenerationSampleEncoder):
-    def get_batching_fields(self, inference_mode: bool) -> List[str]:
-        return ["input_ids", "labels"] if not inference_mode else ["input_ids"]
+class EncDecHFGenerationBaseDataset(HFGenerationBaseDataset):
+    def get_batching_fields(self, teacher_forcing: bool) -> List[str]:
+        return ["input_ids", "labels"] if teacher_forcing else ["input_ids"]
 
 
-class DecHFGenerationSampleEncoder(HFGenerationSampleEncoder):
-    def get_batching_fields(self, inference_mode: bool) -> List[str]:
+class DecHFGenerationBaseDataset(HFGenerationBaseDataset):
+    def get_batching_fields(self, teacher_forcing: bool) -> List[str]:
         return ["input_ids"]
 
     def group_elements_on_materializations(
-        self, dataset_elements: List[Dict[str, Any]], inference_mode: bool
+        self, dataset_elements: List[Dict[str, Any]], teacher_forcing: bool
     ) -> List[List[Dict[str, Any]]]:
 
-        if not inference_mode:
+        if teacher_forcing:
             return [dataset_elements]
 
         groups = collections.defaultdict(list)
@@ -175,8 +79,8 @@ def group_elements_on_materializations(
         return [group for group_len, group in groups.items()]
 
 
-class BartHFGenerationSampleEncoder(EncDecHFGenerationSampleEncoder):
-    def get_fields_batcher(self) -> Dict[str, Callable]:
+class BartHFGenerationDataset(EncDecHFGenerationBaseDataset):
+    def init_fields_batcher(self) -> Dict:
         return {
             "input_ids": lambda lst: batchify(
                 lst, padding_value=self.tokenizer.pad_token_id
@@ -187,102 +91,133 @@ def get_fields_batcher(self) -> Dict[str, Callable]:
                 lst, padding_value=-100
             ),  # -100 == cross entropy ignore index
             "decoder_attention_mask": lambda lst: batchify(lst, padding_value=0),
-            "decoder_start": lambda lst: batchify(
+            "decoder_start_token_id": lambda lst: batchify(
                 lst, padding_value=-1
             ),  # -1 is to force the model to crash (padding should never happen)
         }
 
-    def sample2elem_dict(
-        self, sample: GenerationSample, inference_mode: bool
-    ) -> Dict[str, Any]:
-
-        assert (
-            sample.source_language is None and sample.target_language is None
-        ), f"BartHFGenerationSampleEncoder does not support language specification"
-
-        tokenization_output = self.tokenizer(
-            sample.source_sequence, return_tensors="pt"
-        )
-        elem_dict = {
-            "input_ids": tokenization_output["input_ids"].squeeze(),
-            "attention_mask": tokenization_output["attention_mask"].squeeze(),
-        }
-
-        if not inference_mode:
-            if sample.reference_annotation is not None:
-                tokenization_output = self.tokenizer(
-                    sample.reference_annotation, return_tensors="pt"
-                )
-                elem_dict.update(
-                    **{
-                        "labels": tokenization_output["input_ids"].squeeze(),
-                        "decoder_attention_mask": tokenization_output[
-                            "attention_mask"
-                        ].squeeze(),
-                    }
+    def dataset_iterator_func(self):
+        for sample in self.samples_iterator():
+
+            assert (
+                sample.source_language is None and sample.target_language is None
+            ), f"BartHFGenerationSampleEncoder does not support language specification"
+
+            tokenization_output = self.tokenizer(
+                sample.source_sequence,
+                return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
+            )
+            elem_dict = {
+                "input_ids": tokenization_output["input_ids"].squeeze(),
+                "attention_mask": tokenization_output["attention_mask"].squeeze(),
+            }
+
+            if self.teacher_forcing:
+                if sample.reference_annotation is not None:
+                    tokenization_output = self.tokenizer(
+                        sample.reference_annotation,
+                        return_tensors="pt",
+                        **(
+                            {"truncation": True, "max_length": self.max_length}
+                            if self.truncation
+                            else {}
+                        ),
+                    )
+                    elem_dict.update(
+                        **{
+                            "labels": tokenization_output["input_ids"].squeeze(),
+                            "decoder_attention_mask": tokenization_output[
+                                "attention_mask"
+                            ].squeeze(),
+                        }
+                    )
+            else:
+                elem_dict["decoder_start_token_id"] = torch.tensor(
+                    [self.tokenizer.eos_token_id]
                 )
-        else:
-            elem_dict["decoder_start"] = torch.tensor([self.tokenizer.eos_token_id])
 
-        elem_dict["samples"] = sample
-        return elem_dict
+            elem_dict["samples"] = sample
+            yield elem_dict
 
 
-class MBARTHFGenerationSampleEncoder(BartHFGenerationSampleEncoder):
+class MBartHFGenerationDataset(BartHFGenerationDataset):
 
     mbart_l2l_code = {}
     for l_code in FAIRSEQ_LANGUAGE_CODES:
         l = l_code[: l_code.index("_")]
         mbart_l2l_code[l] = l_code
 
-    def sample2elem_dict(
-        self, sample: GenerationSample, inference_mode: bool
-    ) -> Dict[str, Any]:
-        assert (
-            sample.source_language is not None and sample.target_language is not None
-        ), f"MBARTHFGenerationSampleEncoder requires language specification"
-
-        self.tokenizer: MBartTokenizerFast
-        self.tokenizer.src_lang = self.mbart_l2l_code[sample.source_language]
-        self.tokenizer.tgt_lang = self.mbart_l2l_code[sample.target_language]
+    def dataset_iterator_func(self):
+        for sample in self.samples_iterator():
+            assert (
+                sample.source_language is not None
+                and sample.target_language is not None
+            ), f"MBartHFGenerationDataset requires language specification"
 
-        tokenization_output = self.tokenizer(
-            sample.source_sequence, return_tensors="pt"
-        )
-        elem_dict = {
-            "input_ids": tokenization_output["input_ids"].squeeze(),
-            "attention_mask": tokenization_output["attention_mask"].squeeze(),
-        }
+            self.tokenizer: MBartTokenizerFast
+            self.tokenizer.src_lang = self.mbart_l2l_code.get(
+                sample.source_language, sample.source_language
+            )
+            self.tokenizer.tgt_lang = self.mbart_l2l_code.get(
+                sample.target_language, sample.target_language
+            )
 
-        if not inference_mode:
-            with self.tokenizer.as_target_tokenizer():
-                tokenization_output = self.tokenizer(
-                    sample.reference_annotation, return_tensors="pt"
-                )
-                elem_dict.update(
-                    **{
-                        "labels": tokenization_output["input_ids"].squeeze(),
-                        "decoder_attention_mask": tokenization_output[
-                            "attention_mask"
-                        ].squeeze(),
-                    }
-                )
-        else:
-            elem_dict["decoder_start"] = torch.tensor(
-                [
-                    self.tokenizer.lang_code_to_id[
-                        self.mbart_l2l_code[sample.target_language]
-                    ]
-                ]
+            tokenization_output = self.tokenizer(
+                sample.source_sequence,
+                return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
             )
+            elem_dict = {
+                "input_ids": tokenization_output["input_ids"].squeeze(),
+                "attention_mask": tokenization_output["attention_mask"].squeeze(),
+            }
+
+            if self.teacher_forcing:
+                with self.tokenizer.as_target_tokenizer():
+                    tokenization_output = self.tokenizer(
+                        sample.reference_annotation,
+                        return_tensors="pt",
+                        **(
+                            {"truncation": True, "max_length": self.max_length}
+                            if self.truncation
+                            else {}
+                        ),
+                    )
+                    elem_dict.update(
+                        **{
+                            "labels": tokenization_output["input_ids"].squeeze(),
+                            "decoder_attention_mask": tokenization_output[
+                                "attention_mask"
+                            ].squeeze(),
+                        }
+                    )
+            else:
+                elem_dict["decoder_start_token_id"] = torch.tensor(
+                    [
+                        self.tokenizer.convert_tokens_to_ids(
+                            self.mbart_l2l_code.get(
+                                sample.target_language, sample.target_language
+                            )
+                        )
+                    ]
+                )
 
-        elem_dict["samples"] = sample
-        return elem_dict
+            elem_dict["samples"] = sample
+            yield elem_dict
 
     def group_elements_on_materializations(
-        self, dataset_elements: List[Dict[str, Any]], inference_mode: bool
+        self, dataset_elements: List[Dict[str, Any]], teacher_forcing: bool
     ) -> List[List[Dict[str, Any]]]:
-        if not inference_mode:
+        if teacher_forcing:
             return [dataset_elements]
 
         groups = collections.defaultdict(list)
@@ -293,59 +228,133 @@ def group_elements_on_materializations(
         return [group for group_len, group in groups.items()]
 
 
-class GPT2HFGenerationSampleEcnoder(DecHFGenerationSampleEncoder):
-    def get_fields_batcher(self) -> Dict[str, Callable]:
+class T5HFGenerationDataset(EncDecHFGenerationBaseDataset):
+    def init_fields_batcher(self) -> Dict:
         return {
             "input_ids": lambda lst: batchify(
-                lst, padding_value=0
-            ),  # todo gpt2 does not seem to have a pad token id, double check solution
+                lst, padding_value=self.tokenizer.pad_token_id
+            ),
             "attention_mask": lambda lst: batchify(lst, padding_value=0),
             "samples": None,
             "labels": lambda lst: batchify(
                 lst, padding_value=-100
             ),  # -100 == cross entropy ignore index
+            "decoder_attention_mask": lambda lst: batchify(lst, padding_value=0),
         }
 
-    def sample2elem_dict(
-        self, sample: GenerationSample, inference_mode: bool
-    ) -> Dict[str, Any]:
-        tokenization_output = self.tokenizer(
-            sample.source_sequence, return_tensors="pt"
-        )
-        elem_dict = {
-            "input_ids": tokenization_output["input_ids"].squeeze(),
-            "attention_mask": tokenization_output["attention_mask"].squeeze(),
-            "samples": sample,
+    def dataset_iterator_func(self):
+        for sample in self.samples_iterator():
+
+            assert (
+                sample.source_language is None and sample.target_language is None
+            ), f"T5HFGenerationDataset requires task/language specification to be set in sample.source_sequence"
+
+            tokenization_output = self.tokenizer(
+                sample.source_sequence,
+                return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
+            )
+            elem_dict = {
+                "input_ids": tokenization_output["input_ids"].squeeze(),
+                "attention_mask": tokenization_output["attention_mask"].squeeze(),
+            }
+
+            if self.teacher_forcing:
+                if sample.reference_annotation is not None:
+                    with self.tokenizer.as_target_tokenizer():
+                        tokenization_output = self.tokenizer(
+                            sample.reference_annotation,
+                            return_tensors="pt",
+                            **(
+                                {"truncation": True, "max_length": self.max_length}
+                                if self.truncation
+                                else {}
+                            ),
+                        )
+                        elem_dict.update(
+                            **{
+                                "labels": tokenization_output["input_ids"].squeeze(),
+                                "decoder_attention_mask": tokenization_output[
+                                    "attention_mask"
+                                ].squeeze(),
+                            }
+                        )
+
+            elem_dict["samples"] = sample
+            yield elem_dict
+
+
+class GPT2HFGenerationCataset(DecHFGenerationBaseDataset):
+    def init_fields_batcher(self) -> Dict[str, Callable]:
+        return {
+            "input_ids": lambda lst: batchify(
+                lst, padding_value=0
+            ),  # todo gpt2 does not seem to have a pad token id, double check solution
+            "attention_mask": lambda lst: batchify(lst, padding_value=0),
+            "samples": None,
+            "labels": lambda lst: batchify(
+                lst, padding_value=-100
+            ),  # -100 == cross entropy ignore index
         }
-        if not inference_mode:
-            if sample.reference_annotation is not None:
-                # assume masked clm
-                tokenization_output = self.tokenizer(
-                    sample.reference_annotation, return_tensors="pt"
-                )
-                elem_dict["labels"] = torch.cat(
-                    (
-                        torch.full_like(elem_dict["input_ids"], fill_value=-100),
-                        tokenization_output["input_ids"].squeeze(),
-                    ),
-                    dim=0,
-                )
-                elem_dict["input_ids"] = torch.cat(
-                    (
-                        elem_dict["input_ids"],
-                        tokenization_output["input_ids"].squeeze(),
-                    ),
-                    dim=0,
-                )
-                elem_dict["attention_mask"] = torch.cat(
-                    (
-                        elem_dict["attention_mask"],
-                        tokenization_output["attention_mask"].squeeze(),
-                    ),
-                    dim=0,
-                )
-            else:
-                # assume standard clm
-                elem_dict["labels"] = tokenization_output["input_ids"].squeeze().clone()
 
-        return elem_dict
+    def dataset_iterator_func(self):
+        for sample in self.samples_iterator():
+
+            tokenization_output = self.tokenizer(
+                sample.source_sequence,
+                return_tensors="pt",
+                **(
+                    {"truncation": True, "max_length": self.max_length}
+                    if self.truncation
+                    else {}
+                ),
+            )
+            elem_dict = {
+                "input_ids": tokenization_output["input_ids"].squeeze(),
+                "attention_mask": tokenization_output["attention_mask"].squeeze(),
+                "samples": sample,
+            }
+            if self.teacher_forcing:
+                if sample.reference_annotation is not None:
+                    # assume masked clm
+                    tokenization_output = self.tokenizer(
+                        sample.reference_annotation,
+                        return_tensors="pt",
+                        **(
+                            {"truncation": True, "max_length": self.max_length}
+                            if self.truncation
+                            else {}
+                        ),
+                    )
+                    elem_dict["labels"] = torch.cat(
+                        (
+                            torch.full_like(elem_dict["input_ids"], fill_value=-100),
+                            tokenization_output["input_ids"].squeeze(),
+                        ),
+                        dim=0,
+                    )
+                    elem_dict["input_ids"] = torch.cat(
+                        (
+                            elem_dict["input_ids"],
+                            tokenization_output["input_ids"].squeeze(),
+                        ),
+                        dim=0,
+                    )
+                    elem_dict["attention_mask"] = torch.cat(
+                        (
+                            elem_dict["attention_mask"],
+                            tokenization_output["attention_mask"].squeeze(),
+                        ),
+                        dim=0,
+                    )
+                else:
+                    # assume standard clm
+                    elem_dict["labels"] = (
+                        tokenization_output["input_ids"].squeeze().clone()
+                    )
+
+            yield elem_dict
diff --git a/classy/optional_deps.py b/classy/optional_deps.py
deleted file mode 100644
index 94349f7f..00000000
--- a/classy/optional_deps.py
+++ /dev/null
@@ -1,31 +0,0 @@
-_optional_requirements_str = """
-# classy demo
-streamlit==1.5.0
-st-annotated-text==2.0.0
-# classy serve
-fastapi==0.68.1
-uvicorn[standard]==0.15.0
-# reference api
-pdoc3==0.10.0
-# misc
-plotly==5.5.0
-"""
-_optional_requirements = {}
-
-
-for line in _optional_requirements_str.split("\n")[1:-1]:
-    line = line.strip()
-    if not line.startswith("#"):
-        parts = line.split("==")
-        d = parts[0]
-        if "[" in d:
-            d = d[: d.index("[")]
-        _optional_requirements[d] = line
-
-
-def get_optional_requirement(dependant_package: str) -> str:
-    return _optional_requirements[dependant_package]
-
-
-if __name__ == "__main__":
-    print(_optional_requirements_str[1:])
diff --git a/classy/pl_callbacks/prediction.py b/classy/pl_callbacks/prediction.py
index a4c06088..e3b30b47 100644
--- a/classy/pl_callbacks/prediction.py
+++ b/classy/pl_callbacks/prediction.py
@@ -162,23 +162,31 @@ def on_validation_epoch_start(
             on_result,
         ) in self.settings:
 
-            logger.info(
-                f"Prediction callback processing setting {name} with path={path}"
-            )
-
             if prediction_param_conf_path is not None:
                 model.load_prediction_params(
                     dict(OmegaConf.load(prediction_param_conf_path))
                 )
 
-            # build samples iterator
-            samples_it = iter([])
-            for _path in glob.iglob(path):
-                extension = _path.split(".")[-1]
+            if isinstance(path, str):
+                extension = path.split(".")[-1]
                 data_driver = get_data_driver(model.task, extension)
-                samples_it = itertools.chain(
-                    samples_it, data_driver.read_from_path(_path)
-                )
+                samples_it = data_driver.read_from_path(path)
+            elif isinstance(path, DictConfig):
+                if len(path) > 1:
+                    logger.warning(
+                        "The number of validation datasets is greater than 1."
+                        " At the moment, we automatically set the first validation"
+                        " dataset as the main dataset in the prediction callback."
+                    )
+                path, data_driver = list(path.items())[0]
+                samples_it = data_driver.read_from_path(path)
+                extension = path.split(".")[-1]
+            else:
+                raise NotImplementedError
+
+            logger.info(
+                f"Prediction callback processing setting {name} with path={path}"
+            )
 
             # apply limits (changing path as well)
             if trainer.global_step == 0:
diff --git a/classy/pl_modules/hf/generation.py b/classy/pl_modules/hf/generation.py
index 173eb572..dbab8b16 100644
--- a/classy/pl_modules/hf/generation.py
+++ b/classy/pl_modules/hf/generation.py
@@ -3,7 +3,7 @@
 
 import omegaconf
 import torch
-from torch import nn
+from omegaconf import OmegaConf
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
 
 from classy.data.data_drivers import GenerationSample
@@ -12,29 +12,6 @@
 
 
 class HFGenerationPLModule(GenerationTask, ClassyPLModule):
-    def __init__(
-        self,
-        transformer_model: str,
-        decoding_skip_special_tokens: bool,
-        decoding_clean_up_tokenization_spaces: bool,
-        optim_conf: omegaconf.DictConfig,
-        additional_special_tokens: Optional[List[str]] = None,
-    ):
-        super().__init__(vocabulary=None, optim_conf=optim_conf)
-        self.save_hyperparameters()
-        self.generative_model = HFGenerativeModel.from_transformer_model(
-            transformer_model,
-            decoding_skip_special_tokens=decoding_skip_special_tokens,
-            decoding_clean_up_tokenization_spaces=decoding_clean_up_tokenization_spaces,
-            additional_special_tokens=additional_special_tokens,
-        )
-
-    def load_prediction_params(self, prediction_params: Dict):
-        self.generative_model.load_generation_params(prediction_params)
-
-    def forward(self, *args, **kwargs) -> ClassificationOutput:
-        return self.generative_model.forward(*args, **kwargs)
-
     def training_step(self, batch: dict, batch_idx: int) -> torch.Tensor:
         """ """
         forward_output = self.forward(**batch)
@@ -68,61 +45,23 @@ def test_step(self, batch: dict, batch_idx: int) -> None:
         )
         return forward_output.loss
 
-    def batch_predict(self, *args, **kwargs) -> Iterator[GenerationSample]:
-        return self.generative_model.batch_predict(*args, **kwargs)
-
-
-class HFGenerativeModel(nn.Module):
-    @classmethod
-    def from_transformer_model(cls, transformer_model: str, **kwargs):
-        if re.fullmatch("facebook/bart-(base|large)", transformer_model):
-            return BartGenerativeModule(transformer_model, **kwargs)
-        elif re.fullmatch("facebook/mbart-large-(cc25|50)", transformer_model):
-            return MBartGenerativeModule(transformer_model, **kwargs)
-        elif transformer_model.startswith("gpt2"):
-            return GPT2GenerativeModule(transformer_model, **kwargs)
-        else:
-            raise ValueError
-
-    def __init__(
-        self,
-        transformer_model: str,
-        decoding_skip_special_tokens: bool,
-        decoding_clean_up_tokenization_spaces: bool,
-    ):
-        super().__init__()
-        self.generation_params = {}
-
-    def load_generation_params(self, generation_params: Dict):
-        self.generation_params = generation_params
 
-    def forward(self, *args, **kwargs) -> ClassificationOutput:
-        raise NotImplementedError
-
-    def batch_predict(self, *args, **kwargs) -> Iterator[Tuple[GenerationSample, str]]:
-        raise NotImplementedError
-
-
-class BartGenerativeModule(HFGenerativeModel):
+class BartGenerativeModule(HFGenerationPLModule):
     def __init__(
         self,
         transformer_model: str,
         decoding_skip_special_tokens: bool,
         decoding_clean_up_tokenization_spaces: bool,
+        optim_conf: omegaconf.DictConfig,
         additional_special_tokens: Optional[List[str]] = None,
     ):
-        super().__init__(
-            transformer_model,
-            decoding_skip_special_tokens,
-            decoding_clean_up_tokenization_spaces,
-        )
+        super().__init__(vocabulary=None, optim_conf=optim_conf)
         self.tokenizer = AutoTokenizer.from_pretrained(
             transformer_model,
             additional_special_tokens=list(additional_special_tokens)
             if additional_special_tokens is not None
             else None,
             use_fast=True,
-            add_prefix_space=True,  # todo this should be read from config (like facebook/bart-large-xsum has it False)
         )
         self.model = AutoModelForSeq2SeqLM.from_pretrained(transformer_model)
         if additional_special_tokens is not None and len(additional_special_tokens) > 0:
@@ -132,6 +71,10 @@ def __init__(
             decoding_clean_up_tokenization_spaces
         )
         self.forced_bos_token_id = self.tokenizer.bos_token_id
+        self.generation_params = {}
+
+    def load_prediction_params(self, prediction_params: Dict):
+        self.generation_params = prediction_params
 
     def forward(
         self,
@@ -158,17 +101,15 @@ def batch_predict(
         self,
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
-        decoder_start: torch.Tensor,
-        num_return_sequences: int = 1,  # todo implement
+        decoder_start_token_id: torch.Tensor,
         **kwargs,
     ) -> Iterator[GenerationSample]:
-        assert len(set(decoder_start.squeeze(-1).tolist())) == 1
+        assert len(set(decoder_start_token_id.squeeze(-1).tolist())) == 1
         # generate
         bart_out = self.model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            decoder_start_token_id=decoder_start[0][0],
-            num_return_sequences=num_return_sequences,
+            decoder_start_token_id=decoder_start_token_id[0][0],
             forced_bos_token_id=self.forced_bos_token_id,
             **self.generation_params,
         )
@@ -178,10 +119,17 @@ def batch_predict(
             skip_special_tokens=self.decoding_skip_special_tokens,
             clean_up_tokenization_spaces=self.decoding_clean_up_tokenization_spaces,
         )
+        # handle num sequences
+        num_sequences = int(len(decoded_bart_out) / input_ids.shape[0])
+        grouped_decoded_bart_out = []
+        for i in range(0, len(decoded_bart_out), num_sequences):
+            grouped_decoded_bart_out.append(decoded_bart_out[i : i + num_sequences])
         # postprocess
         samples = kwargs.get("samples")
-        for sample, prediction in zip(samples, decoded_bart_out):
-            sample.predicted_annotation = prediction
+        for sample, prediction in zip(samples, grouped_decoded_bart_out):
+            sample.predicted_annotation = prediction[0]
+            if num_sequences > 1:
+                sample.predicted_annotation_group = prediction
             yield sample
 
 
@@ -191,19 +139,98 @@ def __init__(self, *args, **kwargs):
         self.forced_bos_token_id = None
 
 
-class GPT2GenerativeModule(HFGenerativeModel):
+class T5GenerativeModule(HFGenerationPLModule):
     def __init__(
         self,
         transformer_model: str,
         decoding_skip_special_tokens: bool,
         decoding_clean_up_tokenization_spaces: bool,
+        optim_conf: omegaconf.DictConfig,
         additional_special_tokens: Optional[List[str]] = None,
     ):
-        super().__init__(
+        super().__init__(vocabulary=None, optim_conf=optim_conf)
+        self.tokenizer = AutoTokenizer.from_pretrained(
             transformer_model,
-            decoding_skip_special_tokens,
-            decoding_clean_up_tokenization_spaces,
+            additional_special_tokens=list(additional_special_tokens)
+            if additional_special_tokens is not None
+            else None,
+            use_fast=True,
         )
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(transformer_model)
+        if additional_special_tokens is not None and len(additional_special_tokens) > 0:
+            self.model.resize_token_embeddings(len(self.tokenizer))
+        self.decoding_skip_special_tokens = decoding_skip_special_tokens
+        self.decoding_clean_up_tokenization_spaces = (
+            decoding_clean_up_tokenization_spaces
+        )
+        self.generation_params = {}
+
+    def load_prediction_params(self, prediction_params: Dict):
+        self.generation_params = prediction_params
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        labels: Optional[torch.Tensor],
+        decoder_attention_mask: Optional[torch.Tensor],
+        **kwargs,
+    ) -> ClassificationOutput:
+        t5_out = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            decoder_attention_mask=decoder_attention_mask,
+        )
+        return ClassificationOutput(
+            loss=t5_out.loss,
+            logits=t5_out.logits,
+            probabilities=t5_out.logits.softmax(dim=-1),
+            predictions=t5_out.logits.argmax(dim=-1),
+        )
+
+    def batch_predict(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        **kwargs,
+    ) -> Iterator[GenerationSample]:
+        # generate
+        t5_out = self.model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **self.generation_params,
+        )
+        # decode
+        decoded_t5_out = self.tokenizer.batch_decode(
+            t5_out,
+            skip_special_tokens=self.decoding_skip_special_tokens,
+            clean_up_tokenization_spaces=self.decoding_clean_up_tokenization_spaces,
+        )
+        # handle num sequences
+        num_sequences = int(len(decoded_t5_out) / input_ids.shape[0])
+        grouped_decoded_t5_out = []
+        for i in range(0, len(decoded_t5_out), num_sequences):
+            grouped_decoded_t5_out.append(decoded_t5_out[i : i + num_sequences])
+        # postprocess
+        samples = kwargs.get("samples")
+        for sample, prediction in zip(samples, grouped_decoded_t5_out):
+            sample.predicted_annotation = prediction[0]
+            if num_sequences > 1:
+                sample.predicted_annotation_group = prediction
+            yield sample
+
+
+class GPT2GenerativeModule(HFGenerationPLModule):
+    def __init__(
+        self,
+        transformer_model: str,
+        decoding_skip_special_tokens: bool,
+        decoding_clean_up_tokenization_spaces: bool,
+        optim_conf: omegaconf.DictConfig,
+        additional_special_tokens: Optional[List[str]] = None,
+    ):
+        super().__init__(vocabulary=None, optim_conf=optim_conf)
         self.tokenizer = AutoTokenizer.from_pretrained(
             transformer_model,
             additional_special_tokens=list(additional_special_tokens)
@@ -221,6 +248,10 @@ def __init__(
         self.decoding_clean_up_tokenization_spaces = (
             decoding_clean_up_tokenization_spaces
         )
+        self.generation_params = {}
+
+    def load_prediction_params(self, prediction_params: Dict):
+        self.generation_params = prediction_params
 
     def forward(
         self,
@@ -245,14 +276,12 @@ def batch_predict(
         self,
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
-        num_return_sequences: int = 1,  # todo implement
         **kwargs,
     ) -> Iterator[Tuple[GenerationSample, str]]:
         # generate
         gpt_out = self.model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            num_return_sequences=num_return_sequences,
             **self.generation_params,
         )
         # decode
@@ -261,7 +290,15 @@ def batch_predict(
             skip_special_tokens=self.decoding_skip_special_tokens,
             clean_up_tokenization_spaces=self.decoding_clean_up_tokenization_spaces,
         )
+        # handle num sequences
+        num_sequences = int(len(decoded_gpt_out) / input_ids.shape[0])
+        grouped_decoded_gpt_out = []
+        for i in range(0, len(decoded_gpt_out), num_sequences):
+            grouped_decoded_gpt_out.append(decoded_gpt_out[i : i + num_sequences])
         # postprocess
         samples = kwargs.get("samples")
         for sample, prediction in zip(samples, decoded_gpt_out):
-            yield sample, prediction
+            sample.predicted_annotation = prediction[0]
+            if num_sequences > 1:
+                sample.predicted_annotation_group = prediction
+            yield sample
diff --git a/classy/pl_modules/mixins/prediction.py b/classy/pl_modules/mixins/prediction.py
index 78bf8291..8305f8e9 100644
--- a/classy/pl_modules/mixins/prediction.py
+++ b/classy/pl_modules/mixins/prediction.py
@@ -9,6 +9,9 @@
 from tqdm import tqdm
 
 from classy.data.data_drivers import ClassySample
+from classy.utils.log import get_project_logger
+
+logger = get_project_logger(__name__)
 
 
 class PredictionMixin:
@@ -78,6 +81,13 @@ def samples_it():
                     del position2predicted_sample[next_prediction_position]
                     next_prediction_position += 1
 
+        if len(position2predicted_sample) > 0:
+            logger.warning(
+                "It seems samples have been discarded in your dataset. This means that you WON'T have a prediction for each input sample. Prediction order will also be partially disrupted"
+            )
+            for k, v in sorted(position2predicted_sample.items(), key=lambda x: x[0]):
+                yield v
+
         if progress_bar:
             iterator.close()
 
diff --git a/classy/pl_modules/mixins/saving.py b/classy/pl_modules/mixins/saving.py
index cc90c091..761bd392 100644
--- a/classy/pl_modules/mixins/saving.py
+++ b/classy/pl_modules/mixins/saving.py
@@ -32,6 +32,10 @@ def save_resources_and_update_config(
         # move every paths into "./resources/" and overwrite the config
         Path(experiment_folder / "resources").mkdir()
 
+        # a same resource might be used by multiple components at the same time
+        # avoid copying them multiple times
+        colored_paths = set()
+
         def fix_with_copy_side_effect(path):
             input_path = Path(path)
             assert input_path.exists()
@@ -39,10 +43,12 @@ def fix_with_copy_side_effect(path):
                 experiment_folder / "resources" / input_path.relative_to(working_folder)
             )
             output_path.parent.mkdir(parents=True, exist_ok=True)
-            if os.path.isdir(input_path):
-                shutil.copytree(input_path, output_path)
-            else:
-                shutil.copy(input_path, output_path)
+            if input_path not in colored_paths:
+                if Path(input_path).is_dir():
+                    shutil.copytree(input_path, output_path)
+                else:
+                    shutil.copy(input_path, output_path)
+            colored_paths.add(input_path)
             return str(output_path.relative_to(experiment_folder))
 
         fix_paths(
diff --git a/classy/pl_modules/mixins/task_serve.py b/classy/pl_modules/mixins/task_serve.py
index d05e94b9..333176ea 100644
--- a/classy/pl_modules/mixins/task_serve.py
+++ b/classy/pl_modules/mixins/task_serve.py
@@ -7,6 +7,12 @@
     SequenceSample,
     TokensSample,
 )
+from classy.utils.optional_deps import requires
+
+try:
+    import pydantic
+except ImportError:
+    pydantic = None
 
 
 class MarshalInputSample:
@@ -39,177 +45,206 @@ def serve_output_class(self) -> Type[MarshalOutputSample]:
         raise NotImplementedError
 
 
-from classy.optional_deps import get_optional_requirement
+class SequenceTaskServeMixin(TaskServeMixin):
 
-try:
+    if pydantic is not None:
 
-    from pydantic import BaseModel, Field
+        class MarshalInputSequenceSample(pydantic.BaseModel, MarshalInputSample):
+            sequence: str = pydantic.Field(None, description="Input sequence")
+
+            def unmarshal(self) -> SequenceSample:
+                return SequenceSample(sequence=self.sequence)
+
+        class MarshalOutputSequenceSample(
+            MarshalInputSequenceSample, MarshalOutputSample
+        ):
+            label: str = pydantic.Field(
+                None, description="Label resulting from model classification"
+            )
 
-    class MarshalInputSequenceSample(BaseModel, MarshalInputSample):
-        sequence: str = Field(None, description="Input sequence")
+            @classmethod
+            def marshal(cls, sample: SequenceSample):
+                return cls(sequence=sample.sequence, label=sample.predicted_annotation)
 
-        def unmarshal(self) -> SequenceSample:
-            return SequenceSample(sequence=self.sequence)
+    else:
 
-    class MarshalOutputSequenceSample(MarshalInputSequenceSample, MarshalOutputSample):
-        label: str = Field(
-            None, description="Label resulting from model classification"
-        )
+        MarshalInputSequenceSample, MarshalOutputSequenceSample = None, None
 
-        @classmethod
-        def marshal(cls, sample: SequenceSample):
-            return cls(sequence=sample.sequence, label=sample.predicted_annotation)
+    @property
+    @requires("pydantic")
+    def serve_input_class(self):
+        return SequenceTaskServeMixin.MarshalInputSequenceSample
 
-    class SequenceTaskServeMixin(TaskServeMixin):
-        @property
-        def serve_input_class(self):
-            return MarshalInputSequenceSample
+    @property
+    @requires("pydantic")
+    def serve_output_class(self):
+        return SequenceTaskServeMixin.MarshalOutputSequenceSample
 
-        @property
-        def serve_output_class(self):
-            return MarshalOutputSequenceSample
 
-    class MarshalInputSentencePairSample(BaseModel, MarshalInputSample):
-        sentence1: str = Field(None, description="First input sentence")
-        sentence2: str = Field(None, description="Second input sentence")
+class SentencePairTaskServeMixin(TaskServeMixin):
 
-        def unmarshal(self) -> SentencePairSample:
-            return SentencePairSample(
-                sentence1=self.sentence1, sentence2=self.sentence2
-            )
+    if pydantic is not None:
 
-    class MarshalOutputSentencePairSample(
-        MarshalInputSentencePairSample, MarshalOutputSample
-    ):
-        label: str = Field(
-            None, description="Label resulting from model classification"
-        )
-
-        @classmethod
-        def marshal(cls, sample: SentencePairSample):
-            return cls(
-                sentence1=sample.sentence1,
-                sentence2=sample.sentence2,
-                label=sample.predicted_annotation,
-            )
+        class MarshalInputSentencePairSample(pydantic.BaseModel, MarshalInputSample):
+            sentence1: str = pydantic.Field(None, description="First input sentence")
+            sentence2: str = pydantic.Field(None, description="Second input sentence")
 
-    class SentencePairTaskServeMixin(TaskServeMixin):
-        @property
-        def serve_input_class(self):
-            return MarshalInputSentencePairSample
-
-        @property
-        def serve_output_class(self):
-            return MarshalOutputSentencePairSample
-
-    class MarshalInputTokensSample(BaseModel, MarshalInputSample):
-        tokens: List[str] = Field(None, description="List of input tokens")
-
-        def unmarshal(self) -> TokensSample:
-            return TokensSample(tokens=self.tokens)
-
-    class MarshalOutputTokensSample(MarshalInputTokensSample, MarshalOutputSample):
-        labels: List[str] = Field(
-            None, description="List of labels the model assigned to each input token"
-        )
-
-        @classmethod
-        def marshal(cls, sample: TokensSample):
-            return cls(tokens=sample.tokens, labels=sample.predicted_annotation)
-
-    class TokenTaskServeMixin(TaskServeMixin):
-        @property
-        def serve_input_class(self):
-            return MarshalInputTokensSample
-
-        @property
-        def serve_output_class(self):
-            return MarshalOutputTokensSample
-
-    class MarshalInputQASample(BaseModel, MarshalInputSample):
-        context: str = Field(None, description="Input context")
-        question: str = Field(None, description="Input question")
-
-        def unmarshal(self) -> QASample:
-            return QASample(context=self.context, question=self.question)
-
-    class MarshalOutputQASample(MarshalInputQASample, MarshalOutputSample):
-        answer_char_start: int = Field(None, description="Answer starting char index")
-        answer_char_end: int = Field(None, description="Answer ending char index")
-
-        @classmethod
-        def marshal(cls, sample: QASample):
-            char_start, char_end = sample.predicted_annotation
-            return cls(
-                context=sample.context,
-                question=sample.question,
-                answer_char_start=char_start,
-                answer_char_end=char_end,
-            )
+            def unmarshal(self) -> SentencePairSample:
+                return SentencePairSample(
+                    sentence1=self.sentence1, sentence2=self.sentence2
+                )
 
-    class QATaskServeMixin(TaskServeMixin):
-        @property
-        def serve_input_class(self):
-            return MarshalInputQASample
-
-        @property
-        def serve_output_class(self):
-            return MarshalOutputQASample
-
-    class MarshalInputGenerationSample(BaseModel, MarshalInputSample):
-        source_sequence: str = Field(None, description="Source sequence")
-        source_language: str = Field(None, description="Source language")
-        target_language: str = Field(None, description="Target language")
-
-        def unmarshal(self) -> GenerationSample:
-            return GenerationSample(
-                source_sequence=self.source_sequence,
-                source_language=self.source_language,
-                target_language=self.target_language,
+        class MarshalOutputSentencePairSample(
+            MarshalInputSentencePairSample, MarshalOutputSample
+        ):
+            label: str = pydantic.Field(
+                None, description="Label resulting from model classification"
             )
 
-    class MarshalOutputGenerationSample(
-        MarshalInputGenerationSample, MarshalOutputSample
-    ):
-        target_sequence: str = Field(
-            None, description="Target sequence resulting from model classification"
-        )
-
-        @classmethod
-        def marshal(cls, sample: GenerationSample):
-            return cls(
-                source_sequence=sample.source_sequence,
-                source_language=sample.source_language,
-                target_language=sample.target_language,
-                target_sequence=sample.predicted_annotation,
+            @classmethod
+            def marshal(cls, sample: SentencePairSample):
+                return cls(
+                    sentence1=sample.sentence1,
+                    sentence2=sample.sentence2,
+                    label=sample.predicted_annotation,
+                )
+
+    else:
+
+        MarshalInputSentencePairSample, MarshalOutputSentencePairSample = None, None
+
+    @property
+    @requires("pydantic")
+    def serve_input_class(self):
+        return SentencePairTaskServeMixin.MarshalInputSentencePairSample
+
+    @property
+    @requires("pydantic")
+    def serve_output_class(self):
+        return SentencePairTaskServeMixin.MarshalOutputSentencePairSample
+
+
+class TokenTaskServeMixin(TaskServeMixin):
+
+    if pydantic is not None:
+
+        class MarshalInputTokensSample(pydantic.BaseModel, MarshalInputSample):
+            tokens: List[str] = pydantic.Field(None, description="List of input tokens")
+
+            def unmarshal(self) -> TokensSample:
+                return TokensSample(tokens=self.tokens)
+
+        class MarshalOutputTokensSample(MarshalInputTokensSample, MarshalOutputSample):
+            labels: List[str] = pydantic.Field(
+                None,
+                description="List of labels the model assigned to each input token",
             )
 
-    class GenerationTaskServeMixin(TaskServeMixin):
-        @property
-        def serve_input_class(self):
-            return MarshalInputGenerationSample
+            @classmethod
+            def marshal(cls, sample: TokensSample):
+                return cls(tokens=sample.tokens, labels=sample.predicted_annotation)
 
-        @property
-        def serve_output_class(self):
-            return MarshalOutputGenerationSample
+    else:
 
-except ImportError:
+        MarshalInputTokensSample, MarshalOutputTokensSample = None, None
 
-    class MockTaskServeMixin(TaskServeMixin):
-        @property
-        def serve_input_class(self):
-            raise ImportError(
-                f"Serve mixins could not be properly instantiated as optional requirements ([pydantic]) are missing. To fix this issue, run `pip install {get_optional_requirement('fastapi')} {get_optional_requirement('uvicorn')}`"
+    @property
+    @requires("pydantic")
+    def serve_input_class(self):
+        return TokenTaskServeMixin.MarshalInputTokensSample
+
+    @property
+    @requires("pydantic")
+    def serve_output_class(self):
+        return TokenTaskServeMixin.MarshalOutputTokensSample
+
+
+class QATaskServeMixin(TaskServeMixin):
+
+    if pydantic is not None:
+
+        class MarshalInputQASample(pydantic.BaseModel, MarshalInputSample):
+            context: str = pydantic.Field(None, description="Input context")
+            question: str = pydantic.Field(None, description="Input question")
+
+            def unmarshal(self) -> QASample:
+                return QASample(context=self.context, question=self.question)
+
+        class MarshalOutputQASample(MarshalInputQASample, MarshalOutputSample):
+            answer_char_start: int = pydantic.Field(
+                None, description="Answer starting char index"
+            )
+            answer_char_end: int = pydantic.Field(
+                None, description="Answer ending char index"
             )
 
-        @property
-        def serve_output_class(self):
-            raise ImportError(
-                f"Serve mixins could not be properly instantiated as optional requirements ([pydantic]) are missing. To fix this issue, run `pip install {get_optional_requirement('fastapi')} {get_optional_requirement('uvicorn')}`"
+            @classmethod
+            def marshal(cls, sample: QASample):
+                char_start, char_end = sample.predicted_annotation
+                return cls(
+                    context=sample.context,
+                    question=sample.question,
+                    answer_char_start=char_start,
+                    answer_char_end=char_end,
+                )
+
+    else:
+
+        MarshalInputQASample, MarshalOutputQASample = None, None
+
+    @property
+    @requires("pydantic")
+    def serve_input_class(self):
+        return QATaskServeMixin.MarshalInputQASample
+
+    @property
+    @requires("pydantic")
+    def serve_output_class(self):
+        return QATaskServeMixin.MarshalOutputQASample
+
+
+class GenerationTaskServeMixin(TaskServeMixin):
+
+    if pydantic is not None:
+
+        class MarshalInputGenerationSample(pydantic.BaseModel, MarshalInputSample):
+            source_sequence: str = pydantic.Field(None, description="Source sequence")
+            source_language: str = pydantic.Field(None, description="Source language")
+            target_language: str = pydantic.Field(None, description="Target language")
+
+            def unmarshal(self) -> GenerationSample:
+                return GenerationSample(
+                    source_sequence=self.source_sequence,
+                    source_language=self.source_language,
+                    target_language=self.target_language,
+                )
+
+        class MarshalOutputGenerationSample(
+            MarshalInputGenerationSample, MarshalOutputSample
+        ):
+            target_sequence: str = pydantic.Field(
+                None, description="Target sequence resulting from model classification"
             )
 
-    GenerationTaskServeMixin = MockTaskServeMixin
-    QATaskServeMixin = MockTaskServeMixin
-    SentencePairTaskServeMixin = MockTaskServeMixin
-    SequenceTaskServeMixin = MockTaskServeMixin
-    TokenTaskServeMixin = MockTaskServeMixin
+            @classmethod
+            def marshal(cls, sample: GenerationSample):
+                return cls(
+                    source_sequence=sample.source_sequence,
+                    source_language=sample.source_language,
+                    target_language=sample.target_language,
+                    target_sequence=sample.predicted_annotation,
+                )
+
+    else:
+
+        MarshalInputGenerationSample, MarshalOutputGenerationSample = None, None
+
+    @property
+    @requires("pydantic")
+    def serve_input_class(self):
+        return GenerationTaskServeMixin.MarshalInputGenerationSample
+
+    @property
+    @requires("pydantic")
+    def serve_output_class(self):
+        return GenerationTaskServeMixin.MarshalOutputGenerationSample
diff --git a/classy/pl_modules/mixins/task_ui.py b/classy/pl_modules/mixins/task_ui.py
index a7772330..44c53a14 100644
--- a/classy/pl_modules/mixins/task_ui.py
+++ b/classy/pl_modules/mixins/task_ui.py
@@ -19,9 +19,16 @@
     TokensSample,
     get_data_driver,
 )
-from classy.optional_deps import get_optional_requirement
+from classy.utils.optional_deps import requires
 from classy.utils.streamlit import get_md_200_random_color_generator
 
+try:
+    import annotated_text
+    import streamlit as st
+except ModuleNotFoundError:
+    st = None
+    annotated_text = None
+
 
 class TaskUIMixin:
     def ui_render_task_in_sidebar(self):
@@ -42,338 +49,322 @@ def ui_render(
         raise NotImplementedError
 
 
-try:
-
-    import streamlit as st
-    from annotated_text import annotation
-
-    class SentencePairTaskUIMixin(TaskUIMixin):
+class SentencePairTaskUIMixin(TaskUIMixin):
 
-        __data_driver = get_data_driver(SENTENCE_PAIR, JSONL)
-        truncate_k = 40
+    __data_driver = get_data_driver(SENTENCE_PAIR, JSONL)
+    truncate_k = 40
 
-        def ui_render_task_in_sidebar(self):
-            st.markdown(
-                f"""
-                    * **task**: Sentence-Pair Classification
-                    * **input**: Pair of input sentences
-                """
-            )
+    @requires("streamlit")
+    def ui_render_task_in_sidebar(self):
+        st.markdown(
+            f"""
+                * **task**: Sentence-Pair Classification
+                * **input**: Pair of input sentences
+            """
+        )
 
-        def ui_read_input(
-            self, inference_message: str, inferred_examples: List[SentencePairSample]
-        ) -> Optional[SentencePairSample]:
-            # tuple can't be used for selection boxes, let's use incipts
-            option2example = {}
-            for ie in inferred_examples:
-                option2example[
-                    f"({ie.sentence1[: self.truncate_k]}, {ie.sentence2[: self.truncate_k]})"
-                ] = (
-                    ie.sentence1,
-                    ie.sentence2,
-                )
-            # build selection box
-            selection_message = (
-                inference_message.rstrip(".")
-                + f". Examples are in the format (<sentence1>, <sentence2>); for space constraints, only the first {self.truncate_k} characters of each sentence are shown."
-            )
-            selected_option = st.selectbox(
-                selection_message, options=list(option2example.keys()), index=0
-            )
-            sentence1 = st.text_area(
-                "First input sequence", option2example[selected_option][0]
-            )
-            sentence2 = st.text_area(
-                "Second input sequence", option2example[selected_option][1]
+    @requires("streamlit")
+    def ui_read_input(
+        self, inference_message: str, inferred_examples: List[SentencePairSample]
+    ) -> Optional[SentencePairSample]:
+        # tuple can't be used for selection boxes, let's use incipts
+        option2example = {}
+        for ie in inferred_examples:
+            option2example[
+                f"({ie.sentence1[: self.truncate_k]}, {ie.sentence2[: self.truncate_k]})"
+            ] = (
+                ie.sentence1,
+                ie.sentence2,
             )
-            if st.button("Classify", key="classify"):
-                sample = json.dumps({"sentence1": sentence1, "sentence2": sentence2})
-                return next(self.__data_driver.read([sample]))
-            return None
-
-        def ui_render(self, predicted_sample: SentencePairSample, time: float):
-            st.markdown(
-                f"""
-                <div>
-                    <div class="stAlert">
-                        <p>Model classified input with label: <b>{predicted_sample.predicted_annotation}</b></p>
-                    </div>
-                    <p></p>
-                    <div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>
+        # build selection box
+        selection_message = (
+            inference_message.rstrip(".")
+            + f". Examples are in the format (<sentence1>, <sentence2>); for space constraints, only the first {self.truncate_k} characters of each sentence are shown."
+        )
+        selected_option = st.selectbox(
+            selection_message, options=list(option2example.keys()), index=0
+        )
+        sentence1 = st.text_area(
+            "First input sequence", option2example[selected_option][0]
+        )
+        sentence2 = st.text_area(
+            "Second input sequence", option2example[selected_option][1]
+        )
+        if st.button("Classify", key="classify"):
+            sample = json.dumps({"sentence1": sentence1, "sentence2": sentence2})
+            return next(self.__data_driver.read([sample]))
+        return None
+
+    @requires("streamlit")
+    def ui_render(self, predicted_sample: SentencePairSample, time: float):
+        st.markdown(
+            f"""
+            <div>
+                <div class="stAlert">
+                    <p>Model classified input with label: <b>{predicted_sample.predicted_annotation}</b></p>
                 </div>
-                """,
-                unsafe_allow_html=True,
-            )
+                <p></p>
+                <div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
 
-    class SequenceTaskUIMixin(TaskUIMixin):
-        __data_driver = get_data_driver(SEQUENCE, JSONL)
 
-        def ui_render_task_in_sidebar(self):
-            st.markdown(
-                f"""
-                    * **task**: Sequence Classification
-                    * **input**: String sequence
-                """
-            )
+class SequenceTaskUIMixin(TaskUIMixin):
+    __data_driver = get_data_driver(SEQUENCE, JSONL)
 
-        def ui_read_input(
-            self, inference_message: str, inferred_examples: List[SequenceSample]
-        ) -> Optional[SequenceSample]:
-            placeholder = st.selectbox(
-                inference_message,
-                options=[ie.sequence for ie in inferred_examples],
-                index=0,
-            )
-            input_text = st.text_area("Input sequence to classify", placeholder)
-            if st.button("Classify", key="classify"):
-                sample = json.dumps({"sequence": input_text})
-                return next(self.__data_driver.read([sample]))
-            return None
-
-        def ui_render(self, predicted_sample: SequenceSample, time: float):
-            st.markdown(
-                f"""
-                <div>
-                    <div class="stAlert">
-                        <p>Model classified input with label: <b>{predicted_sample.predicted_annotation}</b></p>
-                    </div>
-                    <p></p>
-                    <div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>
-                </div>
-                """,
-                unsafe_allow_html=True,
-            )
+    @requires("streamlit")
+    def ui_render_task_in_sidebar(self):
+        st.markdown(
+            f"""
+                * **task**: Sequence Classification
+                * **input**: String sequence
+            """
+        )
 
-    class TokenTaskUIMixin(TaskUIMixin):
-        __data_driver = get_data_driver(TOKEN, JSONL)
-        color_generator = get_md_200_random_color_generator()
-        color_mapping = collections.defaultdict(
-            lambda: TokenTaskUIMixin.color_generator()
+    @requires("streamlit")
+    def ui_read_input(
+        self, inference_message: str, inferred_examples: List[SequenceSample]
+    ) -> Optional[SequenceSample]:
+        placeholder = st.selectbox(
+            inference_message,
+            options=[ie.sequence for ie in inferred_examples],
+            index=0,
+        )
+        input_text = st.text_area("Input sequence to classify", placeholder)
+        if st.button("Classify", key="classify"):
+            sample = json.dumps({"sequence": input_text})
+            return next(self.__data_driver.read([sample]))
+        return None
+
+    @requires("streamlit")
+    def ui_render(self, predicted_sample: SequenceSample, time: float):
+        st.markdown(
+            f"""
+            <div>
+                <div class="stAlert">
+                    <p>Model classified input with label: <b>{predicted_sample.predicted_annotation}</b></p>
+                </div>
+                <p></p>
+                <div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>
+            </div>
+            """,
+            unsafe_allow_html=True,
         )
 
-        def ui_render_task_in_sidebar(self):
-            st.markdown(
-                f"""
-                    * **task**: Token Classification
-                    * **input**: Space-separeted list of tokens
-                """
-            )
 
-        def ui_read_input(
-            self, inference_message: str, inferred_examples: List[TokensSample]
-        ) -> Optional[TokensSample]:
-            placeholder = st.selectbox(
-                inference_message,
-                options=[" ".join(ie.tokens) for ie in inferred_examples],
-                index=0,
-            )
-            input_text = st.text_area(
-                "Space-separeted list of tokens to classify", placeholder
-            )
-            if st.button("Classify", key="classify"):
-                sample = json.dumps({"tokens": input_text.split(" ")})
-                return next(self.__data_driver.read([sample]))
-            return None
+class TokenTaskUIMixin(TaskUIMixin):
+    __data_driver = get_data_driver(TOKEN, JSONL)
+    color_generator = get_md_200_random_color_generator()
+    color_mapping = collections.defaultdict(lambda: TokenTaskUIMixin.color_generator())
 
-        def ui_render(self, predicted_sample: TokensSample, time: float):
+    @requires("streamlit")
+    def ui_render_task_in_sidebar(self):
+        st.markdown(
+            f"""
+                * **task**: Token Classification
+                * **input**: Space-separeted list of tokens
+            """
+        )
 
-            tokens, labels = (
-                predicted_sample.tokens,
-                predicted_sample.predicted_annotation,
-            )
+    @requires("streamlit")
+    def ui_read_input(
+        self, inference_message: str, inferred_examples: List[TokensSample]
+    ) -> Optional[TokensSample]:
+        placeholder = st.selectbox(
+            inference_message,
+            options=[" ".join(ie.tokens) for ie in inferred_examples],
+            index=0,
+        )
+        input_text = st.text_area(
+            "Space-separeted list of tokens to classify", placeholder
+        )
+        if st.button("Classify", key="classify"):
+            sample = json.dumps({"tokens": input_text.split(" ")})
+            return next(self.__data_driver.read([sample]))
+        return None
+
+    @requires("streamlit")
+    def ui_render(self, predicted_sample: TokensSample, time: float):
 
-            # check if any token encodings (e.g. bio) are used
-            if all(
-                l == "O" or re.fullmatch("^[BI]-.*$", l)
-                for l in predicted_sample.predicted_annotation
-            ):
-                _tokens, _labels = [], []
-                for t, l in zip(tokens, labels):
-                    if l.startswith("I"):
-                        _tokens[-1] += f" {t}"
-                    else:
-                        _tokens.append(t)
-                        _labels.append(None if l == "O" else l[2:])
-                tokens, labels = _tokens, _labels
-
-            assert len(tokens) == len(labels)
-
-            annotated_html_components = []
+        tokens, labels = (
+            predicted_sample.tokens,
+            predicted_sample.predicted_annotation,
+        )
+
+        # check if any token encodings (e.g. bio) are used
+        if all(
+            l == "O" or re.fullmatch("^[BI]-.*$", l)
+            for l in predicted_sample.predicted_annotation
+        ):
+            _tokens, _labels = [], []
             for t, l in zip(tokens, labels):
-                if l is None:
-                    annotated_html_components.append(str(html.escape(f" {t} ")))
+                if l.startswith("I"):
+                    _tokens[-1] += f" {t}"
                 else:
-                    annotated_html_components.append(
-                        str(annotation(*(t, l, self.color_mapping[l])))
-                    )
-
-            st.markdown(
-                "\n".join(
-                    [
-                        "<div>",
-                        *annotated_html_components,
-                        "<p></p>"
-                        f'<div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>'
-                        "</div>",
-                    ]
-                ),
-                unsafe_allow_html=True,
-            )
-
-    class QATaskUIMixin(TaskUIMixin):
-        __data_driver = get_data_driver(QA, JSONL)
-        truncate_k = 40
+                    _tokens.append(t)
+                    _labels.append(None if l == "O" else l[2:])
+            tokens, labels = _tokens, _labels
+
+        assert len(tokens) == len(labels)
+
+        annotated_html_components = []
+        for t, l in zip(tokens, labels):
+            if l is None:
+                annotated_html_components.append(str(html.escape(f" {t} ")))
+            else:
+                annotated_html_components.append(
+                    str(annotated_text.annotation(*(t, l, self.color_mapping[l])))
+                )
 
-        def ui_render_task_in_sidebar(self):
-            st.markdown(
-                f"""
-                    * **task**: QA
-                    * **input**: context and question
-                """
-            )
+        st.markdown(
+            "\n".join(
+                [
+                    "<div>",
+                    *annotated_html_components,
+                    "<p></p>"
+                    f'<div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>'
+                    "</div>",
+                ]
+            ),
+            unsafe_allow_html=True,
+        )
 
-        def ui_read_input(
-            self, inference_message: str, inferred_examples: List[QASample]
-        ) -> Optional[QASample]:
-            # tuple can't be used for selection boxes, let's use incipts
-            option2example = {}
-            for ie in inferred_examples:
-                option2example[
-                    f"({ie.question[: self.truncate_k]}, {ie.context[: self.truncate_k]})"
-                ] = (
-                    ie.question,
-                    ie.context,
-                )
-            # build selection box
-            selection_message = (
-                inference_message.rstrip(".")
-                + f". Examples are in the format (<question>, <context>); for space constraints, only the first {self.truncate_k} characters of both are shown."
-            )
-            selected_option = st.selectbox(
-                selection_message, options=list(option2example.keys()), index=0
-            )
-            # build input area
-            question = st.text_area("Question", option2example[selected_option][0])
-            context = st.text_area("Context", option2example[selected_option][1])
-            if st.button("Classify", key="classify"):
-                sample = json.dumps({"question": question, "context": context})
-                return next(self.__data_driver.read([sample]))
-            return None
-
-        def ui_render(self, predicted_sample: QASample, time: float):
-            char_start, char_end = predicted_sample.predicted_annotation
-            annotated_html_components = [
-                str(html.escape(f"{predicted_sample.context[: char_start]} ")),
-                str(
-                    annotation(
-                        f"{predicted_sample.context[char_start: char_end]}",
-                        background="#f1e740",
-                    )
-                ),
-                str(html.escape(f"{predicted_sample.context[char_end:]} ")),
-            ]
-            print(annotated_html_components)
-            st.markdown(
-                "\n".join(
-                    [
-                        "<div>",
-                        *annotated_html_components,
-                        "<p></p>"
-                        f'<div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>'
-                        "</div>",
-                    ]
-                ),
-                unsafe_allow_html=True,
-            )
 
-    class GenerationTaskUIMixin(TaskUIMixin):
-        __data_driver = get_data_driver(GENERATION, JSONL)
+class QATaskUIMixin(TaskUIMixin):
+    __data_driver = get_data_driver(QA, JSONL)
+    truncate_k = 40
 
-        def ui_render_task_in_sidebar(self):
-            st.markdown(
-                f"""
-                    * **task**: QA
-                    * **input**: source sequence and, optionally (depending on the model), source and target language
-                """
-            )
+    @requires("streamlit")
+    def ui_render_task_in_sidebar(self):
+        st.markdown(
+            f"""
+                * **task**: QA
+                * **input**: context and question
+            """
+        )
 
-        def ui_read_input(
-            self,
-            inference_message: str,
-            inferred_examples: List[GenerationSample],
-        ) -> Optional[GenerationSample]:
-            # tuple can't be used for selection boxes, let's use incipts
-            option2example = {}
-            for ie in inferred_examples:
-                option2example[f"{ie.source_sequence}"] = (
-                    ie.source_sequence,
-                    ie.source_language if ie.source_language is not None else "",
-                    ie.target_language if ie.target_language is not None else "",
-                )
-            # actual reading
-            selected_option = st.selectbox(
-                inference_message, options=list(option2example.keys()), index=0
-            )
-            source_sequence = st.text_area(
-                "Source sequence", option2example[selected_option][0]
+    @requires("streamlit")
+    def ui_read_input(
+        self, inference_message: str, inferred_examples: List[QASample]
+    ) -> Optional[QASample]:
+        # tuple can't be used for selection boxes, let's use incipts
+        option2example = {}
+        for ie in inferred_examples:
+            option2example[
+                f"({ie.question[: self.truncate_k]}, {ie.context[: self.truncate_k]})"
+            ] = (
+                ie.question,
+                ie.context,
             )
-            source_language = st.text_input(
-                "Source language (empty to set it to None)",
-                option2example[selected_option][1],
-            ).strip()
-            target_language = st.text_input(
-                "Target language (empty to set it to None)",
-                option2example[selected_option][2],
-            ).strip()
-            if st.button("Classify", key="classify"):
-                sample = json.dumps(
-                    dict(
-                        source_sequence=source_sequence,
-                        source_language=source_language or None,
-                        target_language=target_language or None,
-                    )
+        # build selection box
+        selection_message = (
+            inference_message.rstrip(".")
+            + f". Examples are in the format (<question>, <context>); for space constraints, only the first {self.truncate_k} characters of both are shown."
+        )
+        selected_option = st.selectbox(
+            selection_message, options=list(option2example.keys()), index=0
+        )
+        # build input area
+        question = st.text_area("Question", option2example[selected_option][0])
+        context = st.text_area("Context", option2example[selected_option][1])
+        if st.button("Classify", key="classify"):
+            sample = json.dumps({"question": question, "context": context})
+            return next(self.__data_driver.read([sample]))
+        return None
+
+    @requires("streamlit")
+    def ui_render(self, predicted_sample: QASample, time: float):
+        char_start, char_end = predicted_sample.predicted_annotation
+        annotated_html_components = [
+            str(html.escape(f"{predicted_sample.context[: char_start]} ")),
+            str(
+                annotated_text.annotation(
+                    f"{predicted_sample.context[char_start: char_end]}",
+                    background="#f1e740",
                 )
-                return next(self.__data_driver.read([sample]))
-            return None
-
-        def ui_render(self, predicted_sample: GenerationSample, time: float):
-            st.markdown(
-                f"""
-                        <div>
-                            <div class="stAlert">
-                                <p>Model generated: <br>{predicted_sample.predicted_annotation}</p>
-                            </div>
-                            <p></p>
-                            <div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>
-                        </div>
-                        """,
-                unsafe_allow_html=True,
-            )
+            ),
+            str(html.escape(f"{predicted_sample.context[char_end:]} ")),
+        ]
+        print(annotated_html_components)
+        st.markdown(
+            "\n".join(
+                [
+                    "<div>",
+                    *annotated_html_components,
+                    "<p></p>"
+                    f'<div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>'
+                    "</div>",
+                ]
+            ),
+            unsafe_allow_html=True,
+        )
 
-except ImportError:
 
-    class MockTaskServeMixin(TaskUIMixin):
-        def ui_render_task_in_sidebar(self):
-            raise ImportError(
-                f"UI mixins could not be properly instantiated as one of its requirements ([streamlit, st-annotated-text]) are not isntalled. To fix this issue, run `pip install {get_optional_requirement('streamlit')} {get_optional_requirement('st-annotated-text')}`"
-            )
+class GenerationTaskUIMixin(TaskUIMixin):
+    __data_driver = get_data_driver(GENERATION, JSONL)
 
-        def ui_read_input(
-            self, inference_message: str, inferred_examples: List[ClassySample]
-        ) -> Optional[ClassySample]:
+    @requires("streamlit")
+    def ui_render_task_in_sidebar(self):
+        st.markdown(
+            f"""
+                * **task**: QA
+                * **input**: source sequence and, optionally (depending on the model), source and target language
+            """
+        )
 
-            raise ImportError(
-                f"UI mixins could not be properly instantiated as one of its requirements ([streamlit, st-annotated-text]) are not isntalled. To fix this issue, run `pip install {get_optional_requirement('streamlit')} {get_optional_requirement('st-annotated-text')}`"
+    @requires("streamlit")
+    def ui_read_input(
+        self,
+        inference_message: str,
+        inferred_examples: List[GenerationSample],
+    ) -> Optional[GenerationSample]:
+        # tuple can't be used for selection boxes, let's use incipts
+        option2example = {}
+        for ie in inferred_examples:
+            option2example[f"{ie.source_sequence}"] = (
+                ie.source_sequence,
+                ie.source_language if ie.source_language is not None else "",
+                ie.target_language if ie.target_language is not None else "",
             )
-
-        def ui_render(self, predicted_sample: ClassySample, time: float):
-
-            raise ImportError(
-                f"UI mixins could not be properly instantiated as one of its requirements ([streamlit, st-annotated-text]) are not isntalled. To fix this issue, run `pip install {get_optional_requirement('streamlit')} {get_optional_requirement('st-annotated-text')}`"
+        # actual reading
+        selected_option = st.selectbox(
+            inference_message, options=list(option2example.keys()), index=0
+        )
+        source_sequence = st.text_area(
+            "Source sequence", option2example[selected_option][0]
+        )
+        source_language = st.text_input(
+            "Source language (empty to set it to None)",
+            option2example[selected_option][1],
+        ).strip()
+        target_language = st.text_input(
+            "Target language (empty to set it to None)",
+            option2example[selected_option][2],
+        ).strip()
+        if st.button("Classify", key="classify"):
+            sample = json.dumps(
+                dict(
+                    source_sequence=source_sequence,
+                    source_language=source_language or None,
+                    target_language=target_language or None,
+                )
             )
-
-    GenerationTaskUIMixin = MockTaskServeMixin
-    QATaskUIMixin = MockTaskServeMixin
-    SentencePairTaskUIMixin = MockTaskServeMixin
-    SequenceTaskUIMixin = MockTaskServeMixin
-    TokenTaskUIMixin = MockTaskServeMixin
+            return next(self.__data_driver.read([sample]))
+        return None
+
+    @requires("streamlit")
+    def ui_render(self, predicted_sample: GenerationSample, time: float):
+        st.markdown(
+            f"""
+                    <div>
+                        <div class="stAlert">
+                            <p>Model generated: <br>{predicted_sample.predicted_annotation}</p>
+                        </div>
+                        <p></p>
+                        <div style="text-align: right"><p style="color: gray">Time: {time:.2f}s</p></div>
+                    </div>
+                    """,
+            unsafe_allow_html=True,
+        )
diff --git a/classy/scripts/cli/__init__.py b/classy/scripts/cli/__init__.py
index d80a4199..d93b666f 100644
--- a/classy/scripts/cli/__init__.py
+++ b/classy/scripts/cli/__init__.py
@@ -15,6 +15,10 @@ def get_commands():
     from classy.scripts.cli.download import main as download_main
     from classy.scripts.cli.evaluate import get_parser as evaluate_parser
     from classy.scripts.cli.evaluate import main as evaluate_main
+    from classy.scripts.cli.export import get_parser as export_parser
+    from classy.scripts.cli.export import main as export_main
+    from classy.scripts.cli.import_ import get_parser as import_parser
+    from classy.scripts.cli.import_ import main as import_main
     from classy.scripts.cli.predict import get_parser as predict_parser
     from classy.scripts.cli.predict import main as predict_main
     from classy.scripts.cli.serve import get_parser as serve_parser
@@ -24,7 +28,7 @@ def get_commands():
     from classy.scripts.cli.upload import get_parser as upload_parser
     from classy.scripts.cli.upload import main as upload_main
 
-    return dict(
+    commands = dict(
         train=dict(
             parser=train_parser,
             main=train_main,
@@ -57,8 +61,13 @@ def get_commands():
             parser=upload_parser,
             main=upload_main,
         ),
+        export=dict(parser=export_parser, main=export_main),
     )
 
+    commands["import"] = dict(parser=import_parser, main=import_main)
+
+    return commands
+
 
 def parse_args(commands: dict):
 
@@ -135,7 +144,7 @@ def main():
         )
         exit(1)
 
-    to_import = args.package_dir or maybe_find_directory(("src", "source"))
+    to_import = args.package_dir or maybe_find_directory(("src", "source", "classy"))
 
     if to_import is not None:
         import_module_and_submodules(to_import)
diff --git a/classy/scripts/cli/demo.py b/classy/scripts/cli/demo.py
index 424417e4..d6218a48 100644
--- a/classy/scripts/cli/demo.py
+++ b/classy/scripts/cli/demo.py
@@ -1,12 +1,12 @@
 from argparse import ArgumentParser
 
-from classy.optional_deps import get_optional_requirement
 from classy.scripts.cli.utils import (
     autocomplete_model_path,
     checkpoint_path_from_user_input,
     get_device,
 )
 from classy.utils.help_cli import HELP_MODEL_PATH, HELP_PREDICTION_PARAMS
+from classy.utils.optional_deps import requires
 
 
 def populate_parser(parser: ArgumentParser):
@@ -25,8 +25,8 @@ def populate_parser(parser: ArgumentParser):
     parser.add_argument(
         "-d",
         "--device",
-        default="gpu",
-        help="On which device the model for the demo will be loaded.",
+        default=None,
+        help="On which device the model for the demo will be loaded. If not provided, classy will try to infer the desired behavior from the available environment.",
     )
     parser.add_argument(
         "--prediction-params", type=str, default=None, help=HELP_PREDICTION_PARAMS
@@ -54,24 +54,19 @@ def parse_args():
     return get_parser().parse_args()
 
 
+@requires("streamlit", "demo")
 def main(args):
     # import here to avoid importing before needed
     import sys
 
-    try:
-        from streamlit.cli import main as st_main
-    except ImportError:
-        print(
-            f"classy demo [...] requires `pip install {get_optional_requirement('streamlit')}`"
-        )
-        exit(1)
+    from streamlit.cli import main as st_main
 
     device = get_device(args.device)
 
     # script params
     script_params = [args.model_path]
-    if device != -1:
-        # todo ugly workaround for straemlit which interprets -1 as a streamlit param)
+    if device is not None and device != -1:
+        # todo ugly workaround for streamlit which interprets -1 as a streamlit param)
         script_params += ["cuda_device", str(device)]
     if args.prediction_params is not None:
         script_params += ["prediction_params", args.prediction_params]
diff --git a/classy/scripts/cli/describe.py b/classy/scripts/cli/describe.py
index 9dd0029f..95626b4e 100644
--- a/classy/scripts/cli/describe.py
+++ b/classy/scripts/cli/describe.py
@@ -1,8 +1,8 @@
 from argparse import ArgumentParser
 
 from classy.data.data_drivers import GENERATION, QA, SENTENCE_PAIR, SEQUENCE, TOKEN
-from classy.optional_deps import get_optional_requirement
 from classy.utils.help_cli import HELP_TASKS
+from classy.utils.optional_deps import requires
 
 
 def populate_parser(parser: ArgumentParser):
@@ -51,17 +51,12 @@ def parse_args():
     return get_parser().parse_args()
 
 
+@requires("streamlit", "describe")
 def main(args):
     # import here to avoid importing before needed
     import sys
 
-    try:
-        from streamlit.cli import main as st_main
-    except ImportError:
-        print(
-            f"classy demo [...] requires `pip install {get_optional_requirement('streamlit')}`"
-        )
-        exit(1)
+    from streamlit.cli import main as st_main
 
     # script params
     script_params = [args.task, args.dataset]
diff --git a/classy/scripts/cli/evaluate.py b/classy/scripts/cli/evaluate.py
index 5a29d05a..352aefd6 100644
--- a/classy/scripts/cli/evaluate.py
+++ b/classy/scripts/cli/evaluate.py
@@ -1,8 +1,11 @@
 from argparse import ArgumentParser
 from pathlib import Path
+from typing import Dict, Union
 
 from argcomplete import FilesCompleter
+from omegaconf import OmegaConf
 
+from classy.data.data_drivers import DataDriver
 from classy.scripts.cli.utils import (
     autocomplete_model_path,
     checkpoint_path_from_user_input,
@@ -15,6 +18,7 @@
     HELP_PREDICTION_PARAMS,
     HELP_TOKEN_BATCH_SIZE,
 )
+from classy.utils.train_coordinates import load_bundle
 
 
 def populate_parser(parser: ArgumentParser):
@@ -32,8 +36,8 @@ def populate_parser(parser: ArgumentParser):
     parser.add_argument(
         "-d",
         "--device",
-        default="gpu",
-        help="The device you will use for the evaluation.",
+        default=None,
+        help="The device you will use for the evaluation. If not provided, classy will try to infer the desired behavior from the available environment.",
     )
     parser.add_argument(
         "-o",
@@ -54,6 +58,13 @@ def populate_parser(parser: ArgumentParser):
     parser.add_argument(
         "--prediction-params", type=str, default=None, help=HELP_PREDICTION_PARAMS
     )
+    parser.add_argument(
+        "-t",
+        "--output-type",
+        type=str,
+        default="tree",
+        choices=("tree", "json", "list"),
+    )
 
 
 def get_parser(subparser=None) -> ArgumentParser:
@@ -77,7 +88,7 @@ def parse_args():
     return get_parser().parse_args()
 
 
-def automatically_infer_test_path(model_path: str) -> str:
+def automatically_infer_test_path(model_path: str) -> Union[str, Dict[str, DataDriver]]:
     from classy.utils.lightning import load_training_conf_from_checkpoint
 
     checkpoint_path = Path(model_path)
@@ -94,16 +105,35 @@ def automatically_infer_test_path(model_path: str) -> str:
     # check if dataset_path provided at training time was a folder that contained a test set
     training_conf = load_training_conf_from_checkpoint(model_path)
     dataset_path = Path(training_conf.data.datamodule.dataset_path)
-    if dataset_path.exists() and dataset_path.is_dir():
-        possible_test_files = [fp for fp in dataset_path.iterdir() if fp.stem == "test"]
-        if len(possible_test_files) == 1:
-            return str(possible_test_files[0])
+    if dataset_path.exists():
+        if dataset_path.is_file():
+            coordinates_dict = OmegaConf.load(dataset_path)
+            if "test_dataset" in coordinates_dict:
+                test_bundle = load_bundle(
+                    coordinates_dict["test_dataset"], training_conf.data.datamodule.task
+                )
+                return test_bundle
+            elif "validation_dataset" in coordinates_dict:
+                validation_bundle = load_bundle(
+                    coordinates_dict["validation_dataset"],
+                    training_conf.data.datamodule.task,
+                )
+                return validation_bundle
+
+        if dataset_path.is_dir():
+            possible_test_files = [
+                fp for fp in dataset_path.iterdir() if fp.stem == "test"
+            ]
+            if len(possible_test_files) == 1:
+                return str(possible_test_files[0])
 
     raise ValueError
 
 
 def main(args):
     # import here to avoid importing torch before it's actually needed
+    import torch
+
     from classy.scripts.model.evaluate import evaluate
 
     # input_path: if provided, use default one
@@ -119,13 +149,19 @@ def main(args):
             print("Failed to automatically infer test path")
             input_path = input("Please, explicitly enter test path: ").strip()
 
-    device = get_device(args.device)
+    # read device
+    device = args.device
+    if device is None and torch.cuda.is_available():
+        device = 0
+    device = get_device(device)
+
     evaluate(
         args.model_path,
         device,
         args.token_batch_size,
         input_path,
-        args.output_path,
+        output_type=args.output_type,
+        output_path=args.output_path,
         evaluate_config_path=args.evaluation_config,
         prediction_params=args.prediction_params,
         metrics_fn=None,
diff --git a/classy/scripts/cli/export.py b/classy/scripts/cli/export.py
new file mode 100644
index 00000000..d8627d37
--- /dev/null
+++ b/classy/scripts/cli/export.py
@@ -0,0 +1,53 @@
+from argparse import ArgumentParser
+
+
+def populate_parser(parser: ArgumentParser):
+    parser.add_argument("model_name", help="The model you want to export")
+    parser.add_argument(
+        "--zip-name",
+        help="Name of the output file. Defaults to classy-export-{model_name}.zip",
+    )
+    parser.add_argument(
+        "-ns",
+        "--no-strip",
+        action="store_true",
+        default=False,
+        help="Whether to strip the checkpoint of optimizer states, schedulers and callbacks. "
+        "Should only do this if you're not planning on resuming training (i.e., for inference).",
+    )
+    parser.add_argument(
+        "-a",
+        "--all-ckpts",
+        action="store_true",
+        default=False,
+        help="Whether to include every checkpoint under the <RUN>/checkpoints/ folder or just the `best.ckpt`.",
+    )
+
+
+def get_parser(subparser=None) -> ArgumentParser:
+    parser_kwargs = dict(
+        name="export",
+        description="export a trained model as a zip file",
+        help="Export a trained model as a zip file",
+    )
+    parser = (subparser.add_parser if subparser is not None else ArgumentParser)(
+        **parser_kwargs
+    )
+
+    populate_parser(parser)
+
+    return parser
+
+
+def parse_args():
+    return get_parser().parse_args()
+
+
+def main(args):
+    from classy.scripts.model.export import export
+
+    export(args.model_name, args.no_strip, args.all_ckpts, args.zip_name)
+
+
+if __name__ == "__main__":
+    main(parse_args())
diff --git a/classy/scripts/cli/import_.py b/classy/scripts/cli/import_.py
new file mode 100644
index 00000000..2758fe9b
--- /dev/null
+++ b/classy/scripts/cli/import_.py
@@ -0,0 +1,41 @@
+from argparse import ArgumentParser
+
+
+def populate_parser(parser: ArgumentParser):
+    parser.add_argument(
+        "path", help="Path to the zip file with the model run you want to import"
+    )
+    parser.add_argument(
+        "--exp-dir",
+        help="Path to the experiments folder where the exported model should be added. "
+        "Optional, automatically inferred if running from a classy project root dir.",
+    )
+
+
+def get_parser(subparser=None) -> ArgumentParser:
+    parser_kwargs = dict(
+        name="import",
+        description="import a previously exported trained model from a zip file",
+        help="import a previously exported trained model from a zip file",
+    )
+    parser = (subparser.add_parser if subparser is not None else ArgumentParser)(
+        **parser_kwargs
+    )
+
+    populate_parser(parser)
+
+    return parser
+
+
+def parse_args():
+    return get_parser().parse_args()
+
+
+def main(args):
+    from classy.scripts.model.import_ import import_zip
+
+    import_zip(args.path, args.exp_dir)
+
+
+if __name__ == "__main__":
+    main(parse_args())
diff --git a/classy/scripts/cli/predict.py b/classy/scripts/cli/predict.py
index ab852e94..18f7edfb 100644
--- a/classy/scripts/cli/predict.py
+++ b/classy/scripts/cli/predict.py
@@ -30,8 +30,8 @@ def populate_parser(parser: ArgumentParser):
     interactive_parser.add_argument(
         "-d",
         "--device",
-        default="gpu",
-        help="The device where the dataset prediction will be run.",
+        default=None,
+        help="The device where the dataset prediction will be run. If not provided, classy will try to infer the desired behavior from the available environment.",
     )
     interactive_parser.add_argument(
         "--prediction-params", type=str, default=None, help="Path to prediction params."
@@ -87,11 +87,17 @@ def parse_args():
 
 def main(args):
     # import here to avoid importing torch before it's actually needed
+    import torch
+
     from classy.scripts.model.predict import file_main, interactive_main
 
     subcmd = args.subcmd
 
-    device = get_device(args.device)
+    # read device
+    device = args.device
+    if device is None and torch.cuda.is_available():
+        device = 0
+    device = get_device(device)
 
     if subcmd == "file":
         file_main(
diff --git a/classy/scripts/cli/serve.py b/classy/scripts/cli/serve.py
index 45c43083..72e895db 100644
--- a/classy/scripts/cli/serve.py
+++ b/classy/scripts/cli/serve.py
@@ -26,8 +26,8 @@ def populate_parser(parser: ArgumentParser):
     parser.add_argument(
         "-d",
         "--device",
-        default="gpu",
-        help="On which device the model for the REST api will be loaded.",
+        default=None,
+        help="On which device the model for the REST api will be loaded. If not provided, classy will try to infer the desired behavior from the available environment.",
     )
     parser.add_argument(
         "--token-batch-size", type=int, default=1024, help=HELP_TOKEN_BATCH_SIZE
@@ -60,9 +60,16 @@ def parse_args():
 
 def main(args):
     # import here to avoid importing torch before it's actually needed
+    import torch
+
     from classy.scripts.model.serve import serve
 
-    device = get_device(args.device)
+    # read device
+    device = args.device
+    if device is None and torch.cuda.is_available():
+        device = 0
+    device = get_device(device)
+
     serve(
         args.model_path,
         args.port,
diff --git a/classy/scripts/cli/train.py b/classy/scripts/cli/train.py
index a7c2b887..40e5d690 100644
--- a/classy/scripts/cli/train.py
+++ b/classy/scripts/cli/train.py
@@ -1,13 +1,15 @@
+import copy
+import json
 import os
 import shutil
 import tempfile
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import List, Optional
+from typing import Dict, List, Optional, Tuple
 
+import hydra
 from omegaconf import DictConfig, OmegaConf
 
-import classy
 from classy.data.data_drivers import GENERATION, QA, SENTENCE_PAIR, SEQUENCE, TOKEN
 from classy.scripts.cli.utils import get_device, maybe_find_directory
 from classy.utils.help_cli import HELP_TASKS
@@ -67,8 +69,8 @@ def populate_parser(parser: ArgumentParser):
     parser.add_argument(
         "-d",
         "--device",
-        default="gpu",
-        help="The device you will use for the training of your model.",
+        default=None,
+        help="The device you will use for the training of your model. If not provided, classy will try to infer the desired behavior from the available environment.",
     )  # TODO: add validator?
     parser.add_argument(
         "-cn",
@@ -251,175 +253,223 @@ def _main_resume(model_dir: str):
     hydra.main(config_path=None)(_main_mock)()
 
 
-def apply_profile_on_dir(
-    profile: DictConfig, profile_name: str, config_name: str, config_dir: str
-):
+def classy_merge(
+    base_cfg: DictConfig,
+    updating_cfg: DictConfig,
+    reference_folder: Optional[str] = None,
+) -> List[str]:
 
-    blames = []
-
-    def recurse_and_fix(
-        prefix,
-        profile_node,
-        cfg,
-        blame_prefix,
-        path_to_target_config,
-        defaults,
-        potential_defaults,
-    ):
-
-        if OmegaConf.is_dict(profile_node):
-            # if profile overrides a dict, the original dict should be:
-            target_node = OmegaConf.select(cfg, prefix)
-            if target_node is None:
-                # inserted if the dict was not present
-                OmegaConf.update(cfg, prefix, profile_node, force_add=True)
-                blames.append(
-                    (
-                        [(blame_prefix + "." + prefix).lstrip(".")],
-                        ClassyBlame(f"--profile {profile_name}"),
-                    )
-                )
-            else:
-                if "_target_" in profile_node:
-                    # discarded if _target_ is changed
-                    if prefix == "":
-                        cfg = profile_node
-                        assert potential_defaults is not None
-                        for k in potential_defaults:
-                            if k in cfg:
-                                assert type(cfg[k]) == str
-                                defaults[k] = cfg[k]
-                                cfg.pop(k)
-                    else:
-                        OmegaConf.update(
-                            cfg, prefix, profile_node, merge=False, force_add=True
-                        )
-                    blames.append(
-                        (
-                            [(blame_prefix + "." + prefix).lstrip(".")],
-                            ClassyBlame(f"--profile {profile_name}"),
-                        )
-                    )
-                else:
-                    # merged and updated recursively if it was present
-                    for k, v in profile_node.items():
-                        if potential_defaults is not None and k in potential_defaults:
-                            # potential defaults is not None only for direct children of a root (where the defaults logic should be applied)
-                            # if (k, v) refers to a new default
-                            if type(v) == str:
-                                # this should be added
-                                defaults[k] = v
-                            else:
-                                # launch fixing logic on child file
-                                child_file = (
-                                    path_to_target_config.parent
-                                    / k
-                                    / (defaults[k] + ".yaml")
-                                )
-                                assert (
-                                    child_file.exists()
-                                ), f"{child_file} not found in config dir"
-                                apply_recursively(
-                                    v, child_file, (prefix + "." + k).lstrip(".")
-                                )
-                        else:
-                            # otherwise, standard recursion
-                            recurse_and_fix(
-                                k,
-                                v,
-                                target_node,
-                                (blame_prefix + "." + prefix).lstrip("."),
-                                path_to_target_config=None,
-                                defaults=None,
-                                potential_defaults=None,
-                            )
-        elif OmegaConf.is_list(profile_node):
-            # if profile overrides a list, the original list should be completely overwritten
-            if prefix == "":
-                cfg = profile_node
-            else:
-                OmegaConf.update(cfg, prefix, profile_node, merge=False, force_add=True)
-            blames.append(
-                (
-                    [(blame_prefix + "." + prefix).lstrip(".")],
-                    ClassyBlame(f"--profile {profile_name}"),
-                )
+    changes = []
+
+    def is_interpolation(key):
+        # OmegaConf.is_interpolation(base_cfg, key) requires key to be a direct children of base_cfg
+        # this function patch this behavior so as to support any successor
+        parent = base_cfg
+        if "." in key:
+            parent = OmegaConf.select(base_cfg, key[: key.rindex(".")])
+            key = key[key.rindex(".") + 1 :]
+        return OmegaConf.is_interpolation(parent, key)
+
+    def rec(node, key):
+
+        original_node = OmegaConf.select(base_cfg, key)
+
+        # check if node is interpolation and, if it is, duplicate it
+        if original_node is not None and is_interpolation(key):
+            OmegaConf.update(
+                base_cfg, key, copy.deepcopy(original_node), force_add=True
             )
-        elif type(profile_node) in [str, float, int, bool] or profile_node is None:
-            OmegaConf.update(cfg, prefix, profile_node, force_add=True)
-            blames.append(
-                (
-                    [(blame_prefix + "." + prefix).lstrip(".")],
-                    ClassyBlame(f"--profile {profile_name}"),
+
+        # check if node refers to a config group and, if it does, read it
+        # and replace node
+        if type(node) == str and reference_folder is not None:
+            # check whether key is a primitive or if it refers to a config group
+            for extension in ["yaml", "yml"]:
+                config_group_path = (
+                    Path(reference_folder)
+                    / key.replace(".", "/")
+                    / f"{node}.{extension}"
                 )
+                if config_group_path.exists():
+                    # update node to be the config in the file
+                    node = OmegaConf.load(config_group_path)
+                    break
+
+        # handle node
+        if type(original_node) != type(node):
+            # if types differ, overwrite
+            # note that this is also handle the case when one is None and the other is not
+            OmegaConf.update(base_cfg, key, node, merge=False, force_add=True)
+            changes.append(key)
+        elif OmegaConf.is_dict(node):
+            if "_target_" in node:
+                # overwrite dictionary
+                OmegaConf.update(base_cfg, key, node, merge=False, force_add=True)
+                changes.append(key)
+            else:
+                # recurse
+                for k, v in node.items():
+                    rec(v, f"{key}.{k}")
+        elif OmegaConf.is_list(node):
+            # append
+            OmegaConf.update(
+                base_cfg,
+                key,
+                OmegaConf.select(base_cfg, key) + node,
+                merge=False,
+                force_add=True,
             )
+            changes.append(key)
+        elif type(node) in [float, int, bool, str]:
+            # overwrite
+            OmegaConf.update(base_cfg, key, node, force_add=True)
+            changes.append(key)
         else:
-            raise ValueError(f"Unexpected type {type(profile_node)}: {profile_node}")
-
-        return cfg
+            raise ValueError(f"Unexpected type {type(node)}: {node}")
+
+    for k, v in updating_cfg.items():
+        rec(v, k)
+
+    return changes
+
+
+def apply_profiles_and_cli(
+    config_name: str,
+    config_dir: str,
+    profile_path: Optional[str],
+    cli: Dict[ClassyBlame, List[str]],
+    output_config_name: str,
+    is_profile_path: bool = False,
+) -> List[Tuple[List[str], ClassyBlame]]:
+    def parse_primitive(s: str):
+        o = s
+        if "." in s:
+            try:
+                o = float(s)
+            except ValueError:
+                o = s
+        elif s.isdigit():
+            try:
+                o = int(s)
+            except ValueError:
+                o = s
+        elif s.lower() in ["true", "false"]:
+            return s.lower() == "true"
+        else:
+            try:
+                return json.loads(s)
+            except json.JSONDecodeError:
+                return s
+        return o
+
+    # load initial configuration from folder
+    with hydra.initialize_config_dir(config_dir=config_dir, job_name="train"):
+        base_cfg = hydra.compose(config_name=config_name, return_hydra_config=True)
+        blames = base_cfg.__dict__["_blame"]
+
+    # load profile
+    profile_cfg = (
+        OmegaConf.load(profile_path)
+        if profile_path is not None
+        else OmegaConf.create({})
+    )
 
-    def apply_recursively(profile_node, path_to_target_config: Path, prefix: str):
+    # apply cli edits on profile
+    cli_changes = set()
+    for blame, changes in cli.items():
+        _changes = []
+        for change in changes:
+            _k, _v = change.split("=")
+            _changes.extend(
+                classy_merge(profile_cfg, OmegaConf.create({_k: parse_primitive(_v)}))
+            )
+        cli_changes.update(_changes)
+        blames.append((_changes, blame))
+
+    # apply edits over base_cfg
+    profile_changes = classy_merge(base_cfg, profile_cfg, reference_folder=config_dir)
+
+    # add profile blames (if profile was passed)
+    if profile_path is not None:
+        profile_blame_name = str(profile_path)
+        if not is_profile_path:
+            profile_blame_name = profile_blame_name.split("/")[-1]
+            profile_blame_name = profile_blame_name[: profile_blame_name.rindex(".")]
+        blames.append(
+            (
+                [change for change in profile_changes if change not in cli_changes],
+                ClassyBlame(f"--profile {profile_blame_name}"),
+            )
+        )
 
-        # load conf
-        cfg = OmegaConf.load(path_to_target_config)
+    # save and save
+    OmegaConf.save(base_cfg, f"{config_dir}/{output_config_name}.yaml")
+    return blames
 
-        # compute potential defaults dict (folders present)
-        potential_defaults = set(
-            [
-                d.name
-                for d in path_to_target_config.parent.iterdir()
-                if d.is_dir() and d.name != "__pycache__"
-            ]
-        )
 
-        # extract defaults dict
-        is_self_first = None
-        defaults = {}
-
-        if "defaults" in cfg:
-            for i, d in enumerate(cfg.defaults):
-                if d == "_self_":
-                    is_self_first = i == 0
-                    continue
-                for k, v in d.items():
-                    assert (
-                        k not in defaults
-                    ), f"Key {k} already present in defaults list. Check your defaults list"
-                    defaults[k] = v
-
-        # check all defaults are in potential defaults
-        assert all(d in potential_defaults for d in defaults)
-
-        # apply profile
-        cfg = recurse_and_fix(
-            "",
-            profile_node,
-            cfg,
-            blame_prefix=prefix,
-            path_to_target_config=path_to_target_config,
-            defaults=defaults,
-            potential_defaults=potential_defaults,
-        )
+def handle_device(
+    args, profile_path: Optional[Path], cli_overrides: Dict[ClassyBlame, List[str]]
+):
 
-        # update defaults
-        if len(defaults) > 0:
-            cfg.defaults = [{k: v} for k, v in defaults.items()]
-            if is_self_first is not None:
-                if is_self_first:
-                    cfg.defaults = ["_self_"] + cfg.defaults
-                else:
-                    cfg.defaults = cfg.defaults + ["_self_"]
+    import torch.cuda
 
-        # update config
-        OmegaConf.save(cfg, path_to_target_config)
+    # read profile
+    profile_cfg = (
+        OmegaConf.load(profile_path)
+        if profile_path is not None
+        else OmegaConf.create({})
+    )
 
-    apply_recursively(profile, Path(config_dir) / (config_name + ".yaml"), prefix="")
+    # check that either user used cli params (-d 0, --fp16, ...) or profile or neither of the two
+    profile_overrides_device = "device" in profile_cfg
+    cli_specifies_device = args.device is not None
+    assert int(profile_overrides_device) + int(cli_specifies_device) in [
+        0,
+        1,
+    ], f"You are specifying your device in both profile ({profile_cfg.device}) and cli (-d {args.device}. This is not supported: either specify everything in profile or use only CLI"
 
-    return blames
+    if profile_overrides_device:
+        # here we can do nothing and return as profile logic will handle it
+        return
+    elif cli_specifies_device:
+        # here we need to side-effect on cli_overrides
+        device = get_device(args.device)
+        if device != -1:
+            if args.fp16:
+                cli_overrides[ClassyBlame(f"-d {args.device} --fp16")] = [
+                    f"training.pl_trainer.gpus=[{device}]",
+                    "training.pl_trainer.precision=16",
+                ]
+            else:
+                cli_overrides[ClassyBlame(f"-d {args.device}")] = [
+                    f"training.pl_trainer.gpus=[{device}]",
+                    "training.pl_trainer.precision=32",
+                ]
+        else:
+            # user requested to train on cpu
+            # we just need to check that fp16 has not been requested
+            if args.fp16:
+                logger.error("fp16 is only available when training on a GPU")
+                exit(1)
+    else:
+        # auto-mode: infer what we should be doing from the environment
+        # todo should this code be improved?
+        if torch.cuda.is_available():
+            # use first gpu
+            cli_overrides[ClassyBlame(f"[-d {args.device}]")] = [
+                "training.pl_trainer.accelerator=gpu",
+                "training.pl_trainer.devices=1",
+                "training.pl_trainer.auto_select_gpus=True",
+            ]
+        else:
+            # use cpu
+            pass
 
 
 def main(args):
 
+    import classy
+
     if args.resume_from is not None:
         _main_resume(args.resume_from)
         return
@@ -439,9 +489,6 @@ def main(args):
             ]
         )
 
-        # set blames list
-        blames = []
-
         # copy config dir and installed classy configurations into tmp_dir
         classy_dir = str(Path(classy.__file__).parent.parent / "configurations")
         shutil.copytree(classy_dir, tmp_dir, dirs_exist_ok=True)
@@ -451,70 +498,40 @@ def main(args):
             Path(tmp_dir) / (config_name + ".yaml")
         ).exists(), f"No config name file {config_name} found in temporary config dir"
 
-        # apply profile on config dir
-        if args.profile is not None:
-            if args.profile.endswith(".yaml") or args.profile.endswith(".yml"):
-                logger.info(f"Passed profile {args.profile} was detected to be a path")
-                profile_path = Path(args.profile)
-            else:
-                profile_path = Path(tmp_dir) / "profiles" / (args.profile + ".yaml")
-            assert profile_path.exists(), f"No profile found at {profile_path}"
-            blames += apply_profile_on_dir(
-                OmegaConf.load(profile_path), args.profile, config_name, tmp_dir
-            )
-
-        cmd = ["classy-train", "-cn", args.config_name or args.task, "-cd", tmp_dir]
+        # read user-provided configuration overrides from cli
+        cli_overrides = {}
 
-        # choose device
-        device = get_device(args.device)
-        if device != -1:
-            if args.fp16:
-                cmd.append("device=cuda_amp")
-            else:
-                cmd.append(f"device=cuda")
-            if not isinstance(device, int) or device > 0:
-                cmd.append(f"device.gpus=[{device}]")
-        else:
-            if args.fp16:
-                logger.error("fp16 is only available when training on a GPU")
-                return
-            cmd.append(f"device=cpu")
-        blames.append((["device"], ClassyBlame(f"-d {args.device}")))
-
-        cmd.append(f"exp_name={args.exp_name}")
-        blames.append((["exp_name"], ClassyBlame(f"-n {args.exp_name}")))
+        # read args.config
+        for override in args.config:
+            if "~" in override:
+                raise NotImplementedError('Usage of "~" is currently not supported')
+            cli_overrides[ClassyBlame(f"-c {override}")] = [override]
 
         # add dataset path
-        cmd.append(f"data.datamodule.dataset_path={args.dataset}")
-        blames.append(
-            (["data.datamodule.dataset_path"], ClassyBlame(f"{args.dataset}"))
-        )
+        cli_overrides[ClassyBlame(f"{args.dataset}")] = [
+            f"data.datamodule.dataset_path={args.dataset}"
+        ]
 
         # turn off shuffling if requested
         if args.no_shuffle:
-            cmd.append("data.datamodule.shuffle_dataset=False")
-            blames.append(
-                (["data.datamodule.shuffle_dataset"], ClassyBlame("--no-shuffle"))
-            )
+            cli_overrides[ClassyBlame("--no-shuffle")] = [
+                "data.datamodule.shuffle_dataset=False"
+            ]
 
+        # set number of epochs
         if args.epochs:
-            cmd.append(f"++training.pl_trainer.max_epochs={args.epochs}")
-            blames.append(
-                (
-                    ["training.pl_trainer.max_epochs"],
-                    ClassyBlame(f"--epochs {args.epochs}"),
-                )
-            )
+            cli_overrides[ClassyBlame(f"--epochs {args.epochs}")] = [
+                f"training.pl_trainer.max_epochs={args.epochs}"
+            ]
 
-        # wandb logging
+        # setup wandb logging
         if args.wandb is not None:
-            cmd.append(f"logging.wandb.use_wandb=True")
-            configs = ["logging.wandb.use_wandb"]
+            wandb_overrides = []
+            wandb_overrides.append(f"logging.wandb.use_wandb=True")
 
             if args.wandb == "anonymous":
-                cmd.append(f"logging.wandb.anonymous=allow")
-                configs.append("logging.wandb.anonymous")
-                to_blame = ClassyBlame("--wandb anonymous")
+                wandb_overrides.append(f"logging.wandb.anonymous=allow")
+                wandb_blame = ClassyBlame("--wandb anonymous")
             else:
                 if "@" not in args.wandb:
                     print(
@@ -525,69 +542,83 @@ def main(args):
                     exit(1)
 
                 project, experiment = args.wandb.split("@")
-                cmd.append(f"logging.wandb.project_name={project}")
-                cmd.append(f"logging.wandb.experiment_name={experiment}")
-                configs.extend(
-                    ("logging.wandb.project_name", "logging.wandb.experiment_name")
-                )
-                to_blame = ClassyBlame(f"--wandb {args.wandb}")
+                wandb_overrides.append(f"logging.wandb.project_name={project}")
+                wandb_overrides.append(f"logging.wandb.experiment_name={experiment}")
+                wandb_blame = ClassyBlame(f"--wandb {args.wandb}")
 
-            blames.append((configs, to_blame))
+            cli_overrides[wandb_blame] = wandb_overrides
 
         # change the underlying transformer model
         if args.transformer_model is not None:
-            cmd.append(f"transformer_model={args.transformer_model}")
-            blames.append(
-                (
-                    ["transformer_model"],
-                    ClassyBlame(f"--transformer-model {args.transformer_model}"),
-                )
-            )
+            cli_overrides[
+                ClassyBlame(f"--transformer-model {args.transformer_model}")
+            ] = [f"transformer_model={args.transformer_model}"]
 
-        # precomputed vocabulary from the user
+        # check if the user provided a pre-computed vocabulary
         if args.vocabulary_dir is not None:
-            cmd.append(f"data.vocabulary_dir={args.vocabulary_dir}")
-            blames.append(
-                (
-                    ["data.vocabulary_dir"],
-                    ClassyBlame(f"--vocabulary-dir {args.vocabulary_dir}"),
-                )
-            )
+            cli_overrides[ClassyBlame(f"--vocabulary-dir {args.vocabulary_dir}")] = [
+                f"data.vocabulary_dir={args.vocabulary_dir}"
+            ]
 
         # bid-dataset option
         if args.big_dataset:
             logger.info(
-                "The user selected the --big-dataset option. "
+                "The --big-dataset option has been selected. "
                 "Hence we will: 1) assume the training dataset is ALREADY SHUFFLED "
                 "2) evaluate the model performance every 2 thousand steps"
                 "3) If the dataset provided is a file path when splitting the whole dataset in train, validation and test"
                 "we will partition with the following ratio: 0.90 / 0.05 / 0.05"
             )
-            cmd.append("data.datamodule.shuffle_dataset=False")
-            cmd.append(
-                "training.pl_trainer.val_check_interval=2000"
-            )  # TODO: 2K steps seems quite arbitrary
-            cmd.append("data.datamodule.validation_split_size=0.05")
-            cmd.append("data.datamodule.test_split_size=0.05")
-            blames.append(
-                (
-                    [
-                        "data.datamodule.shuffle_dataset",
-                        "training.pl_trainer.val_check_interval",
-                        "data.datamodule.validation_split_size",
-                        "data.datamodule.test_split_size",
-                    ],
-                    ClassyBlame("--big-dataset"),
-                )
-            )
+            cli_overrides[ClassyBlame("--big-dataset")] = [
+                "data.datamodule.shuffle_dataset=False",
+                "training.pl_trainer.val_check_interval=2000",  # TODO: 2K steps seems quite arbitrary
+                "data.datamodule.validation_split_size=0.05",
+                "data.datamodule.test_split_size=0.05",
+            ]
 
-        # append all user-provided configuration overrides
-        cmd += args.config
-        for override in args.config:
-            key, value = override.split("=")
-            key = key.lstrip("+~")
-            blames.append(([key], ClassyBlame(f"-c {override}")))
+        # read profile
+        profile, profile_path, is_profile_path = args.profile, None, False
+        if profile is not None:
+            if profile.endswith(".yaml") or profile.endswith(".yml"):
+                logger.info(f"Passed profile {profile} was detected to be a path")
+                profile_path = Path(profile)
+                is_profile_path = True
+                assert (
+                    profile_path.exists()
+                ), f"Passed profile {profile} does not seem to exist"
+            else:
+                for extension in ["yaml", "yml"]:
+                    profile_path = (
+                        Path(tmp_dir) / "profiles" / (profile + f".{extension}")
+                    )
+                    if profile_path.exists():
+                        break
+                assert (
+                    profile_path.exists()
+                ), f"Passed profile {profile} does not seem to exist"
+
+        # handle device
+        handle_device(args, profile_path=profile_path, cli_overrides=cli_overrides)
+
+        # set experiment name
+        cli_overrides[ClassyBlame(f"-n {args.exp_name}")] = [
+            f"exp_name={args.exp_name}"
+        ]
+
+        # apply profile and cli overrides
+        blames = apply_profiles_and_cli(
+            config_name=config_name,
+            config_dir=tmp_dir,
+            profile_path=profile_path,
+            is_profile_path=is_profile_path,
+            cli=cli_overrides,
+            output_config_name="run",
+        )
+
+        # setup cmd
+        cmd = ["classy-train", "-cn", "run", "-cd", tmp_dir]
 
+        # we wrap this under a try-catch block because streamlit is an optional dependency
         try:
 
             # we import streamlit so that the stderr handler is added to the root logger here and we can remove it
diff --git a/classy/scripts/model/demo.py b/classy/scripts/model/demo.py
index b0102c7c..836bc6e2 100644
--- a/classy/scripts/model/demo.py
+++ b/classy/scripts/model/demo.py
@@ -4,15 +4,12 @@
 from pathlib import Path
 from typing import Callable, Dict, List, Optional, Tuple
 
-from classy.optional_deps import get_optional_requirement
+from classy.utils.optional_deps import requires
 
 try:
     import streamlit as st
 except ImportError:
-    print(
-        f"classy demo [...] requires `pip install {get_optional_requirement('streamlit')}`"
-    )
-    exit(1)
+    st = None
 
 import torch
 from omegaconf import OmegaConf
@@ -25,6 +22,7 @@
 )
 
 
+@requires("streamlit", "demo")
 def auto_infer_examples(
     task: str, model_checkpoint_path: str
 ) -> Tuple[str, List[ClassySample]]:
@@ -50,6 +48,7 @@ def auto_infer_examples(
         )
 
 
+@requires("streamlit", "demo")
 def tabbed_navigation(
     tabs: Dict[str, Tuple[str, Callable[[], None]]], default_tab: Optional[str] = None
 ):
@@ -100,6 +99,7 @@ def tabbed_navigation(
     tabs.get(active_tab)[1]()
 
 
+@requires("streamlit", "demo")
 def demo(
     model_checkpoint_path: str,
     cuda_device: int,
@@ -195,6 +195,7 @@ def render_model_demo():
     tabbed_navigation(tabs, "model")
 
 
+@requires("streamlit", "demo")
 def main():
     args = parse_args()
     # todo proper solution with argparse named arguments
diff --git a/classy/scripts/model/describe.py b/classy/scripts/model/describe.py
index ab261ec4..403ddc80 100644
--- a/classy/scripts/model/describe.py
+++ b/classy/scripts/model/describe.py
@@ -4,23 +4,20 @@
 
 import numpy as np
 
-from classy.optional_deps import get_optional_requirement
+from classy.utils.optional_deps import requires
 
 try:
     import plotly.express as px
+
+    from classy.utils.plotly import boxplot
 except ImportError:
-    print(
-        f"classy describe [...] requires `pip install {get_optional_requirement('plotly')}`"
-    )
-    exit(1)
+    px = None
+    boxplot = None
 
 try:
     import streamlit as st
 except ImportError:
-    print(
-        f"classy demo [...] requires `pip install {get_optional_requirement('streamlit')}`"
-    )
-    exit(1)
+    st = None
 
 from sacremoses import MosesTokenizer
 
@@ -37,7 +34,6 @@
     TokensSample,
     get_data_driver,
 )
-from classy.utils.plotly import boxplot
 
 # colors
 from classy.utils.streamlit import get_md_400_random_color_generator
@@ -732,6 +728,8 @@ def init_layout(task: str, dataset_path: str, metrics: Iterable[UIMetric]):
     st.title("Metrics Visualization")
 
 
+@requires("streamlit", "describe")
+@requires("plotly.express", "describe")
 def describe(task: str, dataset_path: str, tokenize: Optional[str]) -> None:
     """
     Main method - orchestrator
diff --git a/classy/scripts/model/evaluate.py b/classy/scripts/model/evaluate.py
index 72e85ec7..ce044339 100644
--- a/classy/scripts/model/evaluate.py
+++ b/classy/scripts/model/evaluate.py
@@ -1,10 +1,12 @@
-from typing import Callable, Dict, List, Optional, Tuple
+import itertools
+import logging
+from typing import Callable, Dict, List, Optional, Union
 
 import hydra
 import torch
 from omegaconf import OmegaConf
 
-from classy.data.data_drivers import ClassySample, get_data_driver
+from classy.data.data_drivers import ClassySample, DataDriver, get_data_driver
 from classy.utils.lightning import (
     load_classy_module_from_checkpoint,
     load_prediction_dataset_conf_from_checkpoint,
@@ -16,7 +18,8 @@ def evaluate(
     model_checkpoint_path: str,
     cuda_device: int,
     token_batch_size: int,
-    input_path: str,
+    input_path: Union[str, Dict[str, DataDriver]],
+    output_type: Optional[str] = None,
     output_path: Optional[str] = None,
     evaluate_config_path: Optional[str] = None,
     prediction_params: Optional[str] = None,
@@ -38,8 +41,16 @@ def evaluate(
 
     # load dataset conf and driver
     dataset_conf = load_prediction_dataset_conf_from_checkpoint(model_checkpoint_path)
-    input_extension = input_path.split(".")[-1]
-    data_driver = get_data_driver(model.task, input_extension)
+
+    if isinstance(input_path, str):
+        input_extension = input_path.split(".")[-1]
+        data_driver = get_data_driver(model.task, input_extension)
+        dataset_bundle = {input_path: data_driver}
+    elif isinstance(input_path, dict):
+        dataset_bundle = input_path
+    else:
+        logging.error("input_path must be a str or a DictConfig")
+        raise ValueError
 
     # load evaluation metric
     if metrics_fn is not None:
@@ -58,7 +69,9 @@ def evaluate(
     predicted_samples = list(
         model.predict(
             model=model,
-            samples=data_driver.read_from_path(input_path),
+            samples=itertools.chain(
+                *[dd.read_from_path(p) for p, dd in dataset_bundle.items()]
+            ),
             dataset_conf=dataset_conf,
             token_batch_size=token_batch_size,
             progress_bar=True,
@@ -67,11 +80,52 @@ def evaluate(
 
     # dump predictions if requested
     if output_path is not None:
-        with open(output_path, "w") as f:
-            for sample in predicted_samples:
-                f.write(sample.pretty_print() + "\n")
+        output_data_driver = get_data_driver(model.task, output_path.split(".")[-1])
+        output_data_driver.save(
+            predicted_samples, output_path, use_predicted_annotation=True
+        )
 
     # run evaluation and print metrics
     result = metrics_fn(input_path, predicted_samples)
-    for metric_name, metric_f in result.items():
-        print(f"* {metric_name}: {metric_f}")
+
+    def to_primitives(dictionary):
+        def to_primitive(v):
+            if hasattr(v, "item"):
+                return v.item()
+            return v
+
+        return {k: to_primitive(v) for k, v in dictionary.items()}
+
+    result = to_primitives(result)
+
+    output_type = output_type or "tree"
+
+    if output_type == "json":
+        import json
+
+        print(json.dumps(result))
+
+    elif output_type == "list":
+        for metric_name, metric_f in result.items():
+            print(f"* {metric_name}: {metric_f}")
+
+    elif output_type == "tree":
+        from classy.utils.rich_config import print_config
+
+        # taken from https://stackoverflow.com/a/35508197/1908499
+        def nest_dict(flat_dict, sep="_"):
+            """Return nested dict by splitting the keys on a delimiter."""
+            tree = {}
+            for key, val in flat_dict.items():
+                t = tree
+                prev = None
+                for part in key.split(sep):
+                    if prev is not None:
+                        t = t.setdefault(prev, {})
+                    prev = part
+                else:
+                    t.setdefault(prev, val)
+            return tree
+
+        c = OmegaConf.create(nest_dict(result))
+        print_config(c, tree_label="<scores>")
diff --git a/classy/scripts/model/export.py b/classy/scripts/model/export.py
new file mode 100644
index 00000000..2ccc5b31
--- /dev/null
+++ b/classy/scripts/model/export.py
@@ -0,0 +1,114 @@
+import logging
+import zipfile
+from pathlib import Path
+from typing import Optional
+
+from classy.utils.experiment import Experiment, Run
+
+logger = logging.getLogger(__name__)
+
+
+def strip_checkpoint(
+    checkpoint_path: Path,
+    destination: Path,
+    keys_to_remove=("callbacks", "optimizer_states", "lr_schedulers"),
+):
+    import torch
+
+    logger.debug(f"loading checkpoint {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location="cpu")
+
+    for key in keys_to_remove:
+        if ckpt.pop(key, None) is None:
+            logger.debug(f"key {key} did not exist in checkpoint {checkpoint_path}")
+
+    logger.debug(f"saving stripped checkpoint to {destination}")
+    torch.save(ckpt, destination)
+
+
+def zip_run(
+    run: Run,
+    tmpdir: Path,
+    zip_name: str = "model.zip",
+    strip_ckpt: bool = True,
+    is_export: bool = False,
+    best_only: bool = True,
+) -> Path:
+
+    logger.debug(f"zipping run {run} to {tmpdir}")
+    # creates a zip version of the provided Run (with a single stripped checkpoint) in a model.zip file under `tmpdir`
+    run_dir = run.directory
+    ckpt_path = tmpdir / "best.ckpt"
+    zip_path = tmpdir / zip_name
+
+    relative_directory = run.experiment.directory.parent if is_export else run_dir
+
+    with zipfile.ZipFile(zip_path, "w") as zip_file:
+
+        # fully zip the run directory maintaining its structure
+        for file in run_dir.rglob("*.*"):
+            relative_name = file.relative_to(relative_directory)
+
+            if file.is_dir():
+                continue
+
+            # skip checkpoints as we add a single checkpoint later
+            if "checkpoints/" in str(relative_name):
+                if best_only:
+                    continue
+
+                if strip_ckpt:
+                    strip_checkpoint(file, ckpt_path)
+                    zip_file.write(
+                        ckpt_path, arcname=file.relative_to(relative_directory)
+                    )
+                    continue
+
+            zip_file.write(file, arcname=file.relative_to(relative_directory))
+
+        if best_only:
+            ckpt_name = (
+                run_dir.relative_to(relative_directory) / "checkpoints/best.ckpt"
+            )
+
+            if strip_ckpt:
+                logger.debug("Stripping checkpoint before writing to zip file")
+                strip_checkpoint(run.best_checkpoint, ckpt_path)
+                logger.debug("Writing stripped checkpoint to zip file")
+                zip_file.write(ckpt_path, arcname=ckpt_name)
+            else:
+                zip_file.write(run.best_checkpoint, arcname=ckpt_name)
+
+        # remove stripped checkpoint file as it's inside the zip
+        ckpt_path.unlink()
+
+    return zip_path
+
+
+def export(
+    model_name: str,
+    no_strip: bool,
+    all_ckpts: bool,
+    zip_name: Optional[str] = None,
+):
+    exp = Experiment.from_name(model_name)
+    if exp is None:
+        print(f"No experiment named {model_name} found. Exiting...")
+        return
+
+    run = exp.last_valid_run
+    if run is None:
+        print(f"No valid run found for experiment {model_name}. Exiting...")
+        return
+
+    zip_name = zip_name or f"classy-export-{model_name}.zip"
+
+    zip_file = zip_run(
+        run,
+        Path.cwd(),
+        zip_name=zip_name,
+        strip_ckpt=not no_strip,
+        is_export=True,
+        best_only=not all_ckpts,
+    )
+    print(f"Model exported at {zip_file}")
diff --git a/classy/scripts/model/import_.py b/classy/scripts/model/import_.py
new file mode 100644
index 00000000..e54e8794
--- /dev/null
+++ b/classy/scripts/model/import_.py
@@ -0,0 +1,26 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from classy.utils.experiment import Experiment
+
+logger = logging.getLogger(__name__)
+
+
+def unzip(zip_path: str, target: Path):
+    import zipfile
+
+    """
+    Unzip the contents of `zip_path` into `target`.
+    """
+    logger.debug(f"Unzipping {zip_path} to {target}")
+    with zipfile.ZipFile(zip_path) as f:
+        f.extractall(target)
+
+
+def import_zip(zip_path: str, target_path: Optional[str] = None):
+    if target_path is None:
+        target = Experiment.try_get_experiment_dir()
+    else:
+        target = Path(target_path)
+    unzip(zip_path, target)
diff --git a/classy/scripts/model/serve.py b/classy/scripts/model/serve.py
index a95c92a4..c15715af 100644
--- a/classy/scripts/model/serve.py
+++ b/classy/scripts/model/serve.py
@@ -4,16 +4,14 @@
 import torch
 from omegaconf import OmegaConf
 
-from classy.optional_deps import get_optional_requirement
+from classy.utils.optional_deps import requires
 
 try:
     import uvicorn
     from fastapi import FastAPI
 except ImportError:
-    print(
-        f"classy serve [...] requires `pip install {get_optional_requirement('fastapi')} {get_optional_requirement('uvicorn')}`"
-    )
-    exit(1)
+    uvicorn = None
+    FastAPI = None
 
 from classy.utils.commons import get_local_ip_address
 from classy.utils.lightning import (
@@ -25,6 +23,8 @@
 logger = get_project_logger(__name__)
 
 
+@requires("uvicorn", "serve")
+@requires("fastapi", "serve")
 def serve(
     model_checkpoint_path: str,
     port: int,
diff --git a/classy/scripts/model/train.py b/classy/scripts/model/train.py
index 40b8b102..70bd985b 100644
--- a/classy/scripts/model/train.py
+++ b/classy/scripts/model/train.py
@@ -4,8 +4,12 @@
 import hydra
 import pytorch_lightning as pl
 from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning.callbacks import RichProgressBar
 
 from classy.data.data_modules import ClassyDataModule
+from classy.utils.log import get_project_logger
+
+python_logger = get_project_logger(__name__)
 
 
 def train(conf: DictConfig) -> None:
@@ -28,7 +32,7 @@ def train(conf: DictConfig) -> None:
     pl_module = hydra.utils.instantiate(conf.model, **pl_module_init)
 
     # callbacks declaration
-    callbacks_store = []
+    callbacks_store = [RichProgressBar()]
 
     # lightning callbacks
     if conf.training.early_stopping_callback is not None:
@@ -49,8 +53,20 @@ def train(conf: DictConfig) -> None:
             == "classy.pl_callbacks.prediction.PredictionPLCallback"
             and callback.get("path", None) is None
         ):
-            callback["path"] = pl_data_module.validation_path
-        callbacks_store.append(hydra.utils.instantiate(callback, _recursive_=False))
+            validation_bundle = pl_data_module.train_coordinates.validation_bundle
+            dataset_path = list(validation_bundle.items())[0][0]
+
+            python_logger.info(
+                f"Callback dataset path automatically set to: {dataset_path}"
+            )
+
+            callbacks_store.append(
+                hydra.utils.instantiate(
+                    callback, path=validation_bundle, _recursive_=False
+                )
+            )
+        else:
+            callbacks_store.append(hydra.utils.instantiate(callback, _recursive_=False))
 
     # logging
     logger = None
@@ -95,14 +111,12 @@ def train(conf: DictConfig) -> None:
             resume_from_checkpoint=conf.training.resume_from,
             callbacks=callbacks_store,
             logger=logger,
-            **conf.device,
         )
     else:
         trainer: pl.trainer.Trainer = hydra.utils.instantiate(
             conf.training.pl_trainer,
             callbacks=callbacks_store,
             logger=logger,
-            **conf.device,
         )
 
     # save resources
@@ -113,14 +127,5 @@ def train(conf: DictConfig) -> None:
         data_module=pl_data_module,
     )
 
-    # save updated config (and bk old config)
-    Path(".hydra/config.yaml").rename(".hydra/config.bk.yaml")
-    with open(".hydra/config.yaml", "w") as f:
-        f.write(OmegaConf.to_yaml(conf))
-
-    # saving post trainer-init conf
-    with open(".hydra/config_post_trainer_init.yaml", "w") as f:
-        f.write(OmegaConf.to_yaml(conf))
-
     # module fit
     trainer.fit(pl_module, datamodule=pl_data_module)
diff --git a/classy/scripts/model/upload.py b/classy/scripts/model/upload.py
index a6b35138..285aaaca 100644
--- a/classy/scripts/model/upload.py
+++ b/classy/scripts/model/upload.py
@@ -2,7 +2,6 @@
 import json
 import logging
 import tempfile
-import zipfile
 from datetime import datetime
 from pathlib import Path
 from typing import Optional
@@ -10,59 +9,12 @@
 import huggingface_hub
 
 from classy.scripts.model.download import CLASSY_DATE_FORMAT, get_md5
-from classy.utils.experiment import Experiment, Run
+from classy.scripts.model.export import zip_run
+from classy.utils.experiment import Experiment
 
 logger = logging.getLogger(__name__)
 
 
-def strip_checkpoint(
-    checkpoint_path: Path,
-    destination: Path,
-    keys_to_remove=("callbacks", "optimizer_states", "lr_schedulers"),
-):
-    import torch
-
-    logger.debug(f"loading checkpoint {checkpoint_path}")
-    ckpt = torch.load(checkpoint_path, map_location="cpu")
-
-    for key in keys_to_remove:
-        ckpt.pop(key, None)
-        logger.debug(f"key {key} did not exist in checkpoint {checkpoint_path}")
-
-    logger.debug(f"saving stripped checkpoint to {destination}")
-    torch.save(ckpt, destination)
-
-
-def zip_run(run: Run, tmpdir: Path) -> Path:
-    logger.debug(f"zipping run {run} to {tmpdir}")
-    # creates a zip version of the provided Run (with a single stripped checkpoint) in a model.zip file under `tmpdir`
-    run_dir = run.directory
-    ckpt_path = tmpdir / "best.ckpt"
-    zip_path = tmpdir / "model.zip"
-
-    with zipfile.ZipFile(zip_path, "w") as zip_file:
-
-        # fully zip the run directory maintaining its structure
-        for file in run_dir.rglob("*.*"):
-            relative_name = file.relative_to(run_dir)
-
-            # skip checkpoints as we add a single checkpoint later
-            if file.is_dir() or str(relative_name).startswith("checkpoints"):
-                continue
-
-            zip_file.write(file, arcname=file.relative_to(run_dir))
-
-        logger.debug("Stripping checkpoint before writing to zip file")
-        strip_checkpoint(run.best_checkpoint, ckpt_path)
-        logger.debug("Writing stripped checkpoint to zip file")
-        zip_file.write(ckpt_path, arcname="checkpoints/best.ckpt")
-
-    # remove stripped checkpoint as we don't need to upload it to the hub
-    ckpt_path.unlink()
-
-    return zip_path
-
-
 def create_info_file(tmpdir: Path):
     logger.debug("Computing md5 of model.zip")
     md5 = get_md5(tmpdir / "model.zip")
@@ -91,7 +43,7 @@ def upload(
         print(f"No experiment named {model_name} found. Exiting...")
         return
 
-    run = exp.last_run
+    run = exp.last_valid_run
     if run is None:
         print(f"No valid run found for experiment {model_name}. Exiting...")
         return
diff --git a/classy/utils/config.py b/classy/utils/config.py
index 9031c18d..ddc4bbc6 100644
--- a/classy/utils/config.py
+++ b/classy/utils/config.py
@@ -19,7 +19,7 @@ def __init__(self, blame):
 
         it = (k for k in self.iter_all() if k not in self.key2blame)
         for key in sorted(it, key=lambda e: e.count("."), reverse=True):
-            self.solve_unseen(key)
+            self.key2blame[key] = self.get_most_common_blame(key)
 
     @staticmethod
     def _expand(key):
@@ -35,17 +35,6 @@ def iter_all(self):
                 yield k
                 seen.add(k)
 
-    def children(self, key):
-        return [
-            item
-            for item in self.key2blame
-            if item.startswith(key) and item.count(".") - 1 == key.count(".")
-        ]
-
-    def solve_unseen(self, key):
-        children = self.children(key)
-        self.key2blame[key] = self.get_most_common_blame(children[0])
-
     def blame(self, key):
         if key not in self.key2blame:
             self.key2blame[key] = self.compute_blame(key)
diff --git a/classy/utils/data.py b/classy/utils/data.py
index de204039..37ca4b0a 100644
--- a/classy/utils/data.py
+++ b/classy/utils/data.py
@@ -1,5 +1,6 @@
+import itertools
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Dict, Optional, Tuple
 
 import numpy as np
 
@@ -16,45 +17,50 @@ def create_data_dir():
 
 
 def load_dataset(
-    dataset_path: str,
-    data_driver: DataDriver,
+    path2data_driver: Dict[str, DataDriver],
 ) -> list:
-    return list(data_driver.read_from_path(dataset_path))
+    dataset_iterator = itertools.chain(
+        *[dd.read_from_path(p) for p, dd in path2data_driver.items()]
+    )
+    return list(dataset_iterator)
 
 
 def shuffle_dataset(
-    dataset_path: str,
-    data_driver: DataDriver,
+    path2data_driver: Dict[str, DataDriver],
 ) -> list:
-    samples = load_dataset(dataset_path, data_driver)
+    samples = load_dataset(path2data_driver)
     np.random.shuffle(samples)
     return samples
 
 
 def shuffle_and_store_dataset(
-    dataset_path: str,
-    data_driver: DataDriver,
+    path2data_driver: Dict[str, DataDriver],
+    main_data_driver: DataDriver,
     output_path: str,
 ) -> None:
-    samples = shuffle_dataset(dataset_path, data_driver)
-    data_driver.save(samples, output_path)
+    samples = shuffle_dataset(path2data_driver)
+    main_data_driver.save(samples, output_path)
 
 
 # TODO: we have to modify this script in order to support the split without loading the whole dataset in memory
 def split_dataset(
-    dataset_path: str,
-    data_driver: DataDriver,
+    path2data_driver: Dict[str, DataDriver],
+    main_data_driver: DataDriver,
+    main_extension: str,
     output_folder: str,
     validation_split_size: Optional[float] = None,
     test_split_size: Optional[float] = None,
     data_max_split: Optional[int] = None,
     shuffle: bool = True,
-) -> Tuple[str, Optional[str], Optional[str]]:
+) -> Tuple[
+    Dict[str, DataDriver],
+    Optional[Dict[str, DataDriver]],
+    Optional[Dict[str, DataDriver]],
+]:
 
     assert (
         sum([validation_split_size or 0.0, test_split_size or 0.0]) > 0.0
     ), "At least one between validation_split_size and test_split_size must be provided with a value > 0"
-    output_extension = dataset_path.split(".")[-1]
 
     # create output folder
     create_data_dir()
@@ -62,10 +68,10 @@ def split_dataset(
     # read samples and shuffle
     if shuffle:
         logger.info("Materializing and shuffling dataset before splitting it")
-        samples = shuffle_dataset(dataset_path, data_driver)
+        samples = shuffle_dataset(path2data_driver)
     else:
         logger.info("Materializing dataset before splitting it")
-        samples = load_dataset(dataset_path, data_driver)
+        samples = load_dataset(path2data_driver)
 
     # splitting
     training_samples = samples
@@ -81,8 +87,8 @@ def split_dataset(
             training_samples[:n_validation_samples],
             training_samples[n_validation_samples:],
         )
-        validation_path = str(output_folder.joinpath(f"validation.{output_extension}"))
-        data_driver.save(validation_samples, validation_path)
+        validation_path = str(output_folder.joinpath(f"validation.{main_extension}"))
+        main_data_driver.save(validation_samples, validation_path)
 
     if test_split_size is not None:
         n_test_samples = min(
@@ -92,10 +98,16 @@ def split_dataset(
             training_samples[:n_test_samples],
             training_samples[n_test_samples:],
         )
-        test_path = str(output_folder.joinpath(f"test.{output_extension}"))
-        data_driver.save(test_samples, test_path)
+        test_path = str(output_folder.joinpath(f"test.{main_extension}"))
+        main_data_driver.save(test_samples, test_path)
 
-    train_path = str(output_folder.joinpath(f"train.{output_extension}"))
-    data_driver.save(training_samples, train_path)
+    train_path = str(output_folder.joinpath(f"train.{main_extension}"))
+    main_data_driver.save(training_samples, train_path)
 
-    return train_path, validation_path, test_path
+    train_bundle = {train_path: main_data_driver} if train_path is not None else None
+    validation_bundle = (
+        {validation_path: main_data_driver} if validation_path is not None else None
+    )
+    test_bundle = {test_path: main_data_driver} if test_path is not None else None
+
+    return train_bundle, validation_bundle, test_bundle
diff --git a/classy/utils/experiment.py b/classy/utils/experiment.py
index 25525190..32de5beb 100644
--- a/classy/utils/experiment.py
+++ b/classy/utils/experiment.py
@@ -22,6 +22,10 @@ def name(self):
     def last_run(self) -> "Run":
         return max(self.iter_runs(), key=lambda r: r.date, default=None)
 
+    @property
+    def last_valid_run(self) -> Optional["Run"]:
+        return max(self.iter_candidate_runs(), key=lambda r: r.date, default=None)
+
     @property
     def has_checkpoint(self):
         return any(run.has_checkpoints for run in self.iter_runs())
@@ -83,6 +87,19 @@ def list_available_experiments(
     def list_downloaded_experiments(cls) -> List[str]:
         return cls.list_available_experiments(CLASSY_MODELS_CACHE_PATH)
 
+    @classmethod
+    def try_get_experiment_dir(cls) -> Optional[Path]:
+        cwd = Path.cwd()
+
+        if cwd.name == "experiments":
+            return cwd
+
+        exp_tentative = cwd / "experiments"
+        if exp_tentative.exists() and exp_tentative.is_dir():
+            return exp_tentative
+
+        return None
+
     def iter_runs(self) -> Iterator["Run"]:
         for config_file in self.directory.rglob(".hydra/config.yaml"):
             yield Run.from_hydra_config(hydra_config_path=config_file, experiment=self)
diff --git a/classy/utils/lightning.py b/classy/utils/lightning.py
index 3e0d2ab6..2ab55be7 100644
--- a/classy/utils/lightning.py
+++ b/classy/utils/lightning.py
@@ -18,9 +18,8 @@ def load_training_conf_from_checkpoint(
 ) -> DictConfig:
     # find hydra config path
     experiment_folder = Path(checkpoint_path).parent.parent
-    # load hydra config
-    conf_file = "config_post_trainer_init.yaml" if post_trainer_init else "config.yaml"
-    conf = OmegaConf.load(f"{experiment_folder}/.hydra/{conf_file}")
+    # load hydra configs
+    conf = OmegaConf.load(f"{experiment_folder}/.hydra/config.yaml")
 
     # fix paths
     def check_fn(path):
@@ -57,7 +56,7 @@ def load_classy_module_from_checkpoint(checkpoint_path: str) -> ClassyPLModule:
     conf = load_training_conf_from_checkpoint(checkpoint_path)
 
     # check if the model requires a vocab
-    train_dataset_class = conf["data"]["datamodule"]["train_dataset"]["_target_"]
+    train_dataset_class = conf["data"]["datamodule"]["dataset"]["_target_"]
     if not train_dataset_class.split(".")[-1][
         0
     ].isupper():  # if it is not upper then it is a class method
diff --git a/classy/utils/omegaconf.py b/classy/utils/omegaconf.py
new file mode 100644
index 00000000..f0a219c3
--- /dev/null
+++ b/classy/utils/omegaconf.py
@@ -0,0 +1,82 @@
+import copy
+import re
+
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
+
+def adapt_dataset_from(training_dataset: DictConfig, setting: str):
+    # duplicate configuration
+    training_dataset = copy.deepcopy(training_dataset)
+    # resolve otherwise hydra.utils.instantiate will crash
+    OmegaConf.resolve(training_dataset)
+    # identify adaptation code
+    train_dataset_class = training_dataset["_target_"]
+    if not train_dataset_class.split(".")[-1][
+        0
+    ].isupper():  # if it is not upper then it is a class method
+        train_dataset_class = ".".join(train_dataset_class.split(".")[:-1])
+    # invoke it and return
+    return hydra.utils.instantiate(
+        {"_target_": f"{train_dataset_class}.adapt_dataset_from"},
+        training_dataset=training_dataset,
+        setting=setting,
+        _recursive_=False,
+    )
+
+
+OmegaConf.register_new_resolver("adapt_dataset_from", adapt_dataset_from)
+
+
+def resolve_hf_generation_base_dataset_on_transformer_model(
+    transformer_model: str,
+) -> str:
+    if re.fullmatch("facebook/bart-(base|large)", transformer_model):
+        return "classy.data.dataset.hf.generation.BartHFGenerationDataset.from_file"
+    elif re.fullmatch("facebook/mbart-large-(cc25|50)", transformer_model):
+        return "classy.data.dataset.hf.generation.MBartHFGenerationDataset.from_file"
+    elif transformer_model.startswith("gpt2"):
+        return "classy.data.dataset.hf.generation.GPT2HFGenerationCataset.from_file"
+    elif (
+        transformer_model.startswith("t5-")
+        or transformer_model.startswith("google/t5-")
+        or transformer_model.startswith("google/mt5-")
+    ):
+        return "classy.data.dataset.hf.generation.T5HFGenerationDataset.from_file"
+    else:
+        raise ValueError(
+            f"{transformer_model} not currently supported in automatic resolution. But you can still write your own dataset (write _target_ and its parameters)."
+        )
+
+
+OmegaConf.register_new_resolver(
+    "resolve_hf_generation_base_dataset_on_transformer_model",
+    resolve_hf_generation_base_dataset_on_transformer_model,
+)
+
+
+def resolve_hf_generation_module_on_transformer_model(
+    transformer_model: str,
+) -> str:
+    if re.fullmatch("facebook/bart-(base|large)", transformer_model):
+        return "classy.pl_modules.hf.generation.BartGenerativeModule"
+    elif re.fullmatch("facebook/mbart-large-(cc25|50)", transformer_model):
+        return "classy.pl_modules.hf.generation.MBartGenerativeModule"
+    elif transformer_model.startswith("gpt2"):
+        return "classy.pl_modules.hf.generation.GPT2GenerativeModule"
+    elif (
+        transformer_model.startswith("t5-")
+        or transformer_model.startswith("google/t5-")
+        or transformer_model.startswith("google/mt5-")
+    ):
+        return "classy.pl_modules.hf.generation.T5GenerativeModule"
+    else:
+        raise ValueError(
+            f"{transformer_model} not currently supported in automatic resolution. But you can still write your own dataset (write _target_ and its parameters)."
+        )
+
+
+OmegaConf.register_new_resolver(
+    "resolve_hf_generation_module_on_transformer_model",
+    resolve_hf_generation_module_on_transformer_model,
+)
diff --git a/classy/utils/optional_deps.py b/classy/utils/optional_deps.py
new file mode 100644
index 00000000..dd0376d5
--- /dev/null
+++ b/classy/utils/optional_deps.py
@@ -0,0 +1,19 @@
+import importlib
+from typing import Optional
+
+
+def requires(library, extra_required: Optional[str] = None):
+    def closure(decorated_arg):
+        def inner(*args, **kwargs):
+            try:
+                importlib.import_module(library)
+            except ModuleNotFoundError as e:
+                error_message = f"ModuleNotFoundError: {library} not found."
+                if extra_required is not None:
+                    error_message += f" It seems you haven't installed classy[{extra_required}], try doing `pip install classy[{extra_required}]`"
+                raise ModuleNotFoundError(error_message)
+            return decorated_arg(*args, **kwargs)
+
+        return inner
+
+    return closure
diff --git a/classy/utils/plotly.py b/classy/utils/plotly.py
index aad79cc8..4e51f60b 100644
--- a/classy/utils/plotly.py
+++ b/classy/utils/plotly.py
@@ -1,7 +1,14 @@
 import numpy as np
-import plotly.graph_objects as go
 
+from classy.utils.optional_deps import requires
 
+try:
+    import plotly.graph_objects as go
+except ImportError:
+    go = None
+
+
+@requires("plotly.graph_objects")
 def boxplot(y: np.ndarray, x_name: str, y_name: str, color: str):
     fig = go.Figure()
     fig.add_trace(
diff --git a/classy/utils/rich_config.py b/classy/utils/rich_config.py
index 3d6407bc..f8867114 100644
--- a/classy/utils/rich_config.py
+++ b/classy/utils/rich_config.py
@@ -102,8 +102,7 @@ def __rich__(self):
                 ), f"Unknown blame: {blame}"
                 blame_val = blame[len("[source: ") : -1]
                 provider, config = blame_val.split("/", 1)
-                if provider == "classy":
-                    # if rich.console
+                if provider == "main":
                     config_url = (
                         f"{RichNodeInfo._CLASSY_GITHUB_CONFIG_URL}/{config}.yaml"
                     )
@@ -138,14 +137,16 @@ def __init__(
         ),
         skip_remaining: bool = False,
         additional_blames: Optional[List[Tuple[List[str], "ConfigBlame"]]] = None,
+        label: str = "<root>",
     ):
         self.expl = ExplainableConfig(cfg, additional_blames)
         self.fields_order = fields_order
         self.skip_remaining = skip_remaining
+        self.label = label
 
     def get_rich_tree(self) -> Tree:
         style = "dim"
-        tree = Tree("<root>", guide_style=style)
+        tree = Tree(self.label, guide_style=style)
 
         ordered_keys = list(self.fields_order)
 
@@ -155,6 +156,9 @@ def get_rich_tree(self) -> Tree:
             )
 
         for key in ordered_keys:
+            if key not in self.expl.cfg:
+                continue
+
             for branch in self.walk_config(key, sort=False):
                 tree.add(branch)
 
@@ -193,15 +197,21 @@ def walk_config(self, key, sort: bool = True) -> List[Tree]:
 
 
 def get_rich_tree_config(
-    cfg: DictConfig, blames: Optional[List[Tuple[List[str], ConfigBlame]]] = None
+    cfg: DictConfig,
+    blames: Optional[List[Tuple[List[str], ConfigBlame]]] = None,
+    tree_label: str = "<root>",
 ):
-    return ConfigPrinter(cfg, additional_blames=blames).get_rich_tree()
+    return ConfigPrinter(
+        cfg, additional_blames=blames, label=tree_label
+    ).get_rich_tree()
 
 
 def print_config(
-    cfg: DictConfig, blames: Optional[List[Tuple[List[str], ConfigBlame]]] = None
+    cfg: DictConfig,
+    blames: Optional[List[Tuple[List[str], ConfigBlame]]] = None,
+    tree_label: str = "<root>",
 ):
-    rich.print(get_rich_tree_config(cfg, blames))
+    rich.print(get_rich_tree_config(cfg, blames, tree_label))
 
 
 RICH_ST_CODE_FORMAT = (
diff --git a/classy/utils/train_coordinates.py b/classy/utils/train_coordinates.py
new file mode 100644
index 00000000..5ef45a98
--- /dev/null
+++ b/classy/utils/train_coordinates.py
@@ -0,0 +1,59 @@
+import collections
+from typing import Dict, Optional, Tuple, Union
+
+import hydra
+from omegaconf import DictConfig, ListConfig
+
+from classy.data.data_drivers import DataDriver, get_data_driver
+from classy.utils.log import get_project_logger
+
+logger = get_project_logger(__name__)
+
+
+def load_bundle(
+    bundle_conf: Optional[Union[str, Dict[str, str]]],
+    task: str,
+    compute_main_extension: bool = False,
+) -> Optional[Union[Dict[str, DataDriver], Tuple[Dict[str, DataDriver], str]]]:
+    if bundle_conf is None:
+        return None
+
+    main_extension = None
+    if type(bundle_conf) == str:
+        file_extension = bundle_conf.split(".")[-1]
+        bundle_store = {
+            hydra.utils.to_absolute_path(bundle_conf): get_data_driver(
+                task, file_extension
+            )
+        }
+        if compute_main_extension:
+            main_extension = file_extension
+    elif type(bundle_conf) == ListConfig:
+        file_extensions = [path.split(".")[-1] for path in bundle_conf]
+        bundle_store = {
+            hydra.utils.to_absolute_path(path): get_data_driver(task, file_extension)
+            for path, file_extension in zip(bundle_conf, file_extensions)
+        }
+        if compute_main_extension:
+            main_extension = collections.Counter(file_extensions).most_common(1)[0][0]
+    elif type(bundle_conf) == DictConfig:
+        bundle_store = {
+            hydra.utils.to_absolute_path(path): get_data_driver(task, file_extension)
+            for path, file_extension in bundle_conf.items()
+        }
+        if compute_main_extension:
+            main_extension = collections.Counter(bundle_conf.values()).most_common(1)[
+                0
+            ][0]
+    else:
+        logger.error(
+            "The value of the dataset in the coordinates file "
+            "must be either a string indicating the dataset, a "
+            "list of string or  a dict path -> file_extension"
+        )
+        raise NotImplementedError
+
+    if main_extension is not None:
+        return bundle_store, main_extension
+    else:
+        return bundle_store
diff --git a/classy/utils/vocabulary.py b/classy/utils/vocabulary.py
index 60f24c2a..3d28d3a7 100644
--- a/classy/utils/vocabulary.py
+++ b/classy/utils/vocabulary.py
@@ -1,4 +1,5 @@
 import collections
+import logging
 from pathlib import Path
 from typing import Dict, Iterable
 
@@ -12,10 +13,13 @@ class Vocabulary:
     UNK = "<unk>"
 
     @classmethod
-    def from_samples(cls, samples: Iterable[Dict[str, str]]):
-        backend_vocab = collections.defaultdict(
-            lambda: {Vocabulary.PAD: 0, Vocabulary.UNK: 1}
-        )
+    def from_samples(cls, samples: Iterable[Dict[str, str]], add_pad_unk: bool = True):
+        if add_pad_unk:
+            backend_vocab = collections.defaultdict(
+                lambda: {Vocabulary.PAD: 0, Vocabulary.UNK: 1}
+            )
+        else:
+            backend_vocab = collections.defaultdict(dict)
         for sample in samples:
             for k, v in sample.items():
                 elem2idx = backend_vocab[k]
@@ -47,7 +51,16 @@ def get_size(self, k: str) -> int:
         return len(self.backend_vocab[k])
 
     def get_idx(self, k: str, elem: str) -> int:
-        return self.backend_vocab[k].get(elem, self.backend_vocab[k][Vocabulary.UNK])
+        idx = self.backend_vocab[k].get(elem)
+        if idx is None:
+            if Vocabulary.UNK in self.backend_vocab[k]:
+                idx = self.backend_vocab[k][Vocabulary.UNK]
+            else:
+                logging.error(
+                    f"Unknown element found {elem} but no {Vocabulary.UNK} in vocabulary"
+                )
+                raise KeyError
+        return idx
 
     def get_elem(self, k: str, idx: int) -> str:
         return self.reverse_backend_vocab[k][idx]
diff --git a/classy/version.py b/classy/version.py
index ad8fc0c3..529fec91 100644
--- a/classy/version.py
+++ b/classy/version.py
@@ -1,10 +1,10 @@
 import os
 
 _MAJOR = "0"
-_MINOR = "2"
+_MINOR = "3"
 # On main and in a nightly release the patch should be one ahead of the last
 # released build.
-_PATCH = "1"
+_PATCH = "0"
 # This is mainly for nightly builds which have the suffix ".dev$DATE". See
 # https://semver.org/#is-v123-a-semantic-version for the semantics.
 _SUFFIX = os.environ.get("CLASSY_VERSION_SUFFIX", "")
diff --git a/configurations/callbacks/evaluation.yaml b/configurations/callbacks/evaluation.yaml
index 8c9c288b..a740df7d 100644
--- a/configurations/callbacks/evaluation.yaml
+++ b/configurations/callbacks/evaluation.yaml
@@ -8,7 +8,7 @@
       _target_: "classy.pl_callbacks.prediction.EvaluationPredictionCallback"
       evaluation: ${evaluation}
   settings:
-    - name: "validation"
+    - name: "val"
       path: null  # leave it to null to set it to PredictionPLCallback.path
       token_batch_size: 800
       prediction_param_conf_path: null
diff --git a/configurations/data/generation.yaml b/configurations/data/generation.yaml
index a3cd9c91..388d8c29 100644
--- a/configurations/data/generation.yaml
+++ b/configurations/data/generation.yaml
@@ -2,30 +2,21 @@ datamodule:
   _target_: 'classy.data.data_modules.ClassyDataModule'
   task: ${task}
   dataset_path: null  # set via kwargs
-  train_dataset:
-    _target_: 'classy.data.dataset.hf.generation.HFGenerationDataset.from_file'
+  dataset:
+    _target_: "${resolve_hf_generation_base_dataset_on_transformer_model:${transformer_model}}"
     transformer_model: ${transformer_model}
+    truncation: true
     additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: -1
-    max_length: 500
+    teacher_forcing: True
+    min_length: 5
+    max_length: 512
     tokens_per_batch: 1000
     max_batch_size: 10
     section_size: 10000
     prebatch: True
     materialize: False
     for_inference: False
-  validation_dataset:
-    _target_: 'classy.data.dataset.hf.generation.HFGenerationDataset.from_file'
-    transformer_model: ${transformer_model}
-    additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: -1
-    max_length: 500
-    tokens_per_batch: 1000
-    max_batch_size: 10
-    section_size: 10000
-    prebatch: True
-    materialize: True
-    for_inference: False
+  validation_dataset: "${adapt_dataset_from:${data.datamodule.dataset},validation}"
   validation_split_size: 0.1
   test_split_size: 0.1
   max_nontrain_split_size: 500
diff --git a/configurations/data/qa.yaml b/configurations/data/qa.yaml
index 84e7fc74..7b6af6b2 100644
--- a/configurations/data/qa.yaml
+++ b/configurations/data/qa.yaml
@@ -2,30 +2,20 @@ datamodule:
   _target_: 'classy.data.data_modules.ClassyDataModule'
   task: ${task}
   dataset_path: null  # set via kwargs
-  train_dataset:
+  dataset:
     _target_: 'classy.data.dataset.hf.classification.HFQADataset.from_file'
     transformer_model: ${transformer_model}
+    truncation: true
     additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
     min_length: 5
-    max_length: 500
+    max_length: 512
     tokens_per_batch: 2000
     max_batch_size: 10
     section_size: 10000
     prebatch: True
     materialize: False
     for_inference: False
-  validation_dataset:
-    _target_: 'classy.data.dataset.hf.classification.HFQADataset.from_file'
-    transformer_model: ${transformer_model}
-    additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: 5
-    max_length: 500
-    tokens_per_batch: 2000
-    max_batch_size: 10
-    section_size: 10000
-    prebatch: True
-    materialize: True
-    for_inference: True
+  validation_dataset: "${adapt_dataset_from:${data.datamodule.dataset},validation}"
   validation_split_size: 0.1
   test_split_size: 0.1
   max_nontrain_split_size: 10000
diff --git a/configurations/data/sentence-pair.yaml b/configurations/data/sentence-pair.yaml
index 29c6bb70..2bace72e 100644
--- a/configurations/data/sentence-pair.yaml
+++ b/configurations/data/sentence-pair.yaml
@@ -2,30 +2,20 @@ datamodule:
   _target_: 'classy.data.data_modules.ClassyDataModule'
   task: ${task}
   dataset_path: null  # set via kwargs
-  train_dataset:
+  dataset:
     _target_: 'classy.data.dataset.hf.classification.HFSentencePairDataset.from_file'
     transformer_model: ${transformer_model}
+    truncation: true
     additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
     min_length: 5
-    max_length: 500
+    max_length: 512
     tokens_per_batch: 800
     max_batch_size: 10
     section_size: 10000
     prebatch: True
     materialize: False
     for_inference: False
-  validation_dataset:
-    _target_: 'classy.data.dataset.hf.classification.HFSentencePairDataset.from_file'
-    transformer_model: ${transformer_model}
-    additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: 5
-    max_length: 500
-    tokens_per_batch: 800
-    max_batch_size: 10
-    section_size: 10000
-    prebatch: True
-    materialize: True
-    for_inference: True
+  validation_dataset: "${adapt_dataset_from:${data.datamodule.dataset},validation}"
   validation_split_size: 0.1
   test_split_size: 0.1
   max_nontrain_split_size: 10000
diff --git a/configurations/data/sequence.yaml b/configurations/data/sequence.yaml
index 6d44e0c2..f04aded0 100644
--- a/configurations/data/sequence.yaml
+++ b/configurations/data/sequence.yaml
@@ -2,30 +2,20 @@ datamodule:
   _target_: 'classy.data.data_modules.ClassyDataModule'
   task: ${task}
   dataset_path: null  # set via kwargs
-  train_dataset:
+  dataset:
     _target_: 'classy.data.dataset.hf.classification.HFSequenceDataset.from_file'
     transformer_model: ${transformer_model}
+    truncation: true
     additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
     min_length: 5
-    max_length: 500
+    max_length: 512
     tokens_per_batch: 2000
     max_batch_size: 10
     section_size: 10000
     prebatch: True
     materialize: False
     for_inference: False
-  validation_dataset:
-    _target_: 'classy.data.dataset.hf.classification.HFSequenceDataset.from_file'
-    transformer_model: ${transformer_model}
-    additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: 5
-    max_length: 500
-    tokens_per_batch: 2000
-    max_batch_size: 10
-    section_size: 10000
-    prebatch: True
-    materialize: True
-    for_inference: True
+  validation_dataset: "${adapt_dataset_from:${data.datamodule.dataset},validation}"
   validation_split_size: 0.1
   test_split_size: 0.1
   max_nontrain_split_size: 10000
diff --git a/configurations/data/token.yaml b/configurations/data/token.yaml
index 2ff9fd4e..0eddaab0 100644
--- a/configurations/data/token.yaml
+++ b/configurations/data/token.yaml
@@ -2,30 +2,20 @@ datamodule:
   _target_: 'classy.data.data_modules.ClassyDataModule'
   task: ${task}
   dataset_path: null  # set via kwargs
-  train_dataset:
+  dataset:
     _target_: 'classy.data.dataset.hf.classification.HFTokenDataset.from_file'
     transformer_model: ${transformer_model}
+    truncation: true
     additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: -1
-    max_length: 500
+    min_length: 5
+    max_length: 512
     tokens_per_batch: 1000
     max_batch_size: 16
     section_size: 10000
     prebatch: True
     materialize: False
     for_inference: False
-  validation_dataset:
-    _target_: 'classy.data.dataset.hf.classification.HFTokenDataset.from_file'
-    transformer_model: ${transformer_model}
-    additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-    min_length: -1
-    max_length: 500
-    tokens_per_batch: 1000
-    max_batch_size: 16
-    section_size: 10000
-    prebatch: True
-    materialize: True
-    for_inference: True
+  validation_dataset: "${adapt_dataset_from:${data.datamodule.dataset},validation}"
   validation_split_size: 0.1
   test_split_size: 0.1
   max_nontrain_split_size: 10000
diff --git a/configurations/device/cpu.yaml b/configurations/device/cpu.yaml
deleted file mode 100644
index 0967ef42..00000000
--- a/configurations/device/cpu.yaml
+++ /dev/null
@@ -1,1 +0,0 @@
-{}
diff --git a/configurations/device/cuda.yaml b/configurations/device/cuda.yaml
deleted file mode 100644
index a233b900..00000000
--- a/configurations/device/cuda.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-gpus:
-  - 0
-precision: 32
diff --git a/configurations/device/cuda_amp.yaml b/configurations/device/cuda_amp.yaml
deleted file mode 100644
index b6e3751c..00000000
--- a/configurations/device/cuda_amp.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-gpus:
-  - 0
-precision: 16
diff --git a/configurations/generation.yaml b/configurations/generation.yaml
index 36c8183d..226cff5a 100644
--- a/configurations/generation.yaml
+++ b/configurations/generation.yaml
@@ -12,8 +12,6 @@ hydra:
   # customize working dir
   run:
     dir: ./experiments/${exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
-  # customize logging
-  verbose: [ classy ]
   job_logging:
     formatters:
       simple:
@@ -25,10 +23,9 @@ hydra:
 defaults:
   - callbacks: empty
   - data: generation
-  - device: cuda
   - evaluation: generation
   - model: generation
-  - prediction: generation
+  - prediction: default
   - training: default
   - logging: default
   - _self_
diff --git a/configurations/model/generation.yaml b/configurations/model/generation.yaml
index 28617905..c9a9b2c3 100644
--- a/configurations/model/generation.yaml
+++ b/configurations/model/generation.yaml
@@ -1,4 +1,4 @@
-_target_: 'classy.pl_modules.hf.generation.HFGenerationPLModule'
+_target_: "${resolve_hf_generation_module_on_transformer_model:${transformer_model}}"
 transformer_model: ${transformer_model}
 additional_special_tokens: []
 decoding_skip_special_tokens: True
diff --git a/configurations/prediction-params/generation-beam.yaml b/configurations/prediction-params/generation-beam.yaml
index 03ba7030..7e654af8 100644
--- a/configurations/prediction-params/generation-beam.yaml
+++ b/configurations/prediction-params/generation-beam.yaml
@@ -1,3 +1,4 @@
+num_return_sequences: 1
 num_beams: 5
 min_length: 5
 max_length: 100
diff --git a/configurations/prediction-params/generation-sample.yaml b/configurations/prediction-params/generation-sample.yaml
index 30ccfe65..03168594 100644
--- a/configurations/prediction-params/generation-sample.yaml
+++ b/configurations/prediction-params/generation-sample.yaml
@@ -1,3 +1,4 @@
+num_return_sequences: 1
 num_beams: 1
 do_sample: True
 top_p: 0.8
diff --git a/configurations/prediction/default.yaml b/configurations/prediction/default.yaml
new file mode 100644
index 00000000..34947007
--- /dev/null
+++ b/configurations/prediction/default.yaml
@@ -0,0 +1,1 @@
+dataset: "${adapt_dataset_from:${data.datamodule.dataset},prediction}"
diff --git a/configurations/prediction/generation.yaml b/configurations/prediction/generation.yaml
deleted file mode 100644
index b51ded97..00000000
--- a/configurations/prediction/generation.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-dataset:
-  _target_: 'classy.data.dataset.hf.generation.HFGenerationDataset.from_samples'
-  transformer_model: ${transformer_model}
-  additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-  min_length: -1
-  max_length: -1
-  tokens_per_batch: 400
-  max_batch_size: -1
-  section_size: 10000
-  prebatch: True
-  materialize: False
-  for_inference: True
diff --git a/configurations/prediction/qa.yaml b/configurations/prediction/qa.yaml
deleted file mode 100644
index e2792422..00000000
--- a/configurations/prediction/qa.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-dataset:
-  _target_: 'classy.data.dataset.hf.classification.HFQADataset.from_samples'
-  transformer_model: ${transformer_model}
-  additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-  min_length: -1
-  max_length: -1
-  tokens_per_batch: 800
-  max_batch_size: -1
-  section_size: 10000
-  prebatch: True
-  materialize: False
-  for_inference: True
diff --git a/configurations/prediction/sentence-pair.yaml b/configurations/prediction/sentence-pair.yaml
deleted file mode 100644
index d1f85d5e..00000000
--- a/configurations/prediction/sentence-pair.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-dataset:
-  _target_: 'classy.data.dataset.hf.classification.HFSentencePairDataset.from_samples'
-  transformer_model: ${transformer_model}
-  additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-  min_length: -1
-  max_length: -1
-  tokens_per_batch: 800
-  max_batch_size: -1
-  section_size: 10000
-  prebatch: True
-  materialize: False
-  for_inference: True
diff --git a/configurations/prediction/sequence.yaml b/configurations/prediction/sequence.yaml
deleted file mode 100644
index e1307ee8..00000000
--- a/configurations/prediction/sequence.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-dataset:
-  _target_: 'classy.data.dataset.hf.classification.HFSequenceDataset.from_samples'
-  transformer_model: ${transformer_model}
-  additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-  min_length: -1
-  max_length: -1
-  tokens_per_batch: 800
-  max_batch_size: -1
-  section_size: 10000
-  prebatch: True
-  materialize: False
-  for_inference: True
diff --git a/configurations/prediction/token.yaml b/configurations/prediction/token.yaml
deleted file mode 100644
index cb686796..00000000
--- a/configurations/prediction/token.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-dataset:
-  _target_: 'classy.data.dataset.hf.classification.HFTokenDataset.from_samples'
-  transformer_model: ${transformer_model}
-  additional_special_tokens: "${oc.select:'model.additional_special_tokens',${oc.decode:'[]'}}"
-  min_length: -1
-  max_length: -1
-  tokens_per_batch: 800
-  max_batch_size: -1
-  section_size: 10000
-  prebatch: True
-  materialize: False
-  for_inference: True
diff --git a/configurations/qa.yaml b/configurations/qa.yaml
index feb445f5..61def1d0 100644
--- a/configurations/qa.yaml
+++ b/configurations/qa.yaml
@@ -12,8 +12,6 @@ hydra:
   # customize working dir
   run:
     dir: ./experiments/${exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
-  # customize logging
-  verbose: [ classy ]
   job_logging:
     formatters:
       simple:
@@ -25,10 +23,9 @@ hydra:
 defaults:
   - callbacks: empty
   - data: qa
-  - device: cuda
   - evaluation: qa
   - model: qa
-  - prediction: qa
+  - prediction: default
   - training: default
   - logging: default
   - _self_
diff --git a/configurations/sentence-pair.yaml b/configurations/sentence-pair.yaml
index 2d8d5391..1ad3b51e 100644
--- a/configurations/sentence-pair.yaml
+++ b/configurations/sentence-pair.yaml
@@ -12,8 +12,6 @@ hydra:
   # customize working dir
   run:
     dir: ./experiments/${exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
-  # customize logging
-  verbose: [ classy ]
   job_logging:
     formatters:
       simple:
@@ -25,10 +23,9 @@ hydra:
 defaults:
   - callbacks: empty
   - data: sentence-pair
-  - device: cuda
   - evaluation: sentence-pair
   - model: sentence-pair
-  - prediction: sentence-pair
+  - prediction: default
   - training: default
   - logging: default
   - _self_
diff --git a/configurations/sequence.yaml b/configurations/sequence.yaml
index af32dba4..25c31a19 100644
--- a/configurations/sequence.yaml
+++ b/configurations/sequence.yaml
@@ -12,8 +12,6 @@ hydra:
   # customize working dir
   run:
     dir: ./experiments/${exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
-  # customize logging
-  verbose: [ classy ]
   job_logging:
     formatters:
       simple:
@@ -25,10 +23,9 @@ hydra:
 defaults:
   - callbacks: empty
   - data: sequence
-  - device: cuda
   - evaluation: sequence
   - model: sequence
-  - prediction: sequence
+  - prediction: default
   - training: default
   - logging: default
   - _self_
diff --git a/configurations/token.yaml b/configurations/token.yaml
index 07c8df3d..a0524492 100644
--- a/configurations/token.yaml
+++ b/configurations/token.yaml
@@ -12,8 +12,6 @@ hydra:
   # customize working dir
   run:
     dir: ./experiments/${exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
-  # customize logging
-  verbose: [ classy ]
   job_logging:
     formatters:
       simple:
@@ -25,10 +23,9 @@ hydra:
 defaults:
   - callbacks: empty
   - data: token
-  - device: cuda
   - evaluation: token
   - model: token
-  - prediction: token
+  - prediction: default
   - training: default
   - logging: default
   - _self_
diff --git a/docs/docs/getting-started/basic/data-formatting.md b/docs/docs/getting-started/basic/data-formatting.md
index a4f9bb25..386d411d 100644
--- a/docs/docs/getting-started/basic/data-formatting.md
+++ b/docs/docs/getting-started/basic/data-formatting.md
@@ -1,11 +1,13 @@
 ---
 sidebar_position: 2
-title: Formatting your data
+title: Organizing your data
 ---
 
+## Data Formatting
+
 `classy` requires data to be formatted in a specific way according to the task you're tackling (check out [Tasks and Input Formats](/docs/reference-manual/tasks-and-formats) in the documentation).
 
-In our case of **Named Entity Recognition** (i.e., *Token Classification*), we need the data to be formatted in a way that each line represents a single sample.
+In our case of **Named Entity Recognition** (i.e., *Token Classification*), we need the data to be formatted such that each line represents a single sample.
 For instance, taking again our running example of *Barack Obama visited Google in California*, we can format it as follows:
 
 ```text
@@ -23,3 +25,22 @@ but you can [add custom formats](/docs/getting-started/customizing-things/custom
 If your dataset is already formatted like this, great! Otherwise, this is the only bit where coding is required.
 You can either convert it yourself (via a python or bash script, whatever you're comfortable with), or you can register
 a [custom data reader](/docs/getting-started/customizing-things/custom-data-format/) to support your dataset format.
+
+
+## Organizing Datasets
+In `classy`, as in standard machine learning projects, the most simple way to organize your datasets is to create
+a directory containing the train, validation and test datasets.
+```
+data/ner-data
+├── train.tsv
+├── validation.tsv
+└── test.tsv
+```
+
+In this way, `classy` will automatically infer the splits of your dataset from the directory structure.
+
+:::tip
+If you have multiple training files, or you want to specify the splits using a different directory structure, you can
+use a _training coordinates_ file. You can find a complete guide on how to do it in the
+[Reference Manual](/docs/reference-manual/cli/train/).
+:::
diff --git a/docs/docs/getting-started/basic/intro.md b/docs/docs/getting-started/basic/intro.md
index 639b5054..69cc2e2e 100644
--- a/docs/docs/getting-started/basic/intro.md
+++ b/docs/docs/getting-started/basic/intro.md
@@ -8,7 +8,8 @@ In the following tutorial, we are going to go over a standard `classy` example,
 
 ## The Task
 
-*Named Entity Recognition* is the task of identifying Named Entities (from a predefined set, e.g., **ORG**, **LOC** and **PER**) in a text.
+*Named Entity Recognition* is the task of identifying Named Entities in a text linking them to a predefined
+set of classes, e.g., **ORG** (organization), **LOC** (location), **PER** (person) and **O** (other) for non-entities.
 For instance, consider what will be our running example for this whole tutorial:
 
 
diff --git a/docs/docs/getting-started/basic/train.md b/docs/docs/getting-started/basic/train.md
index f406e5eb..3ab225b5 100644
--- a/docs/docs/getting-started/basic/train.md
+++ b/docs/docs/getting-started/basic/train.md
@@ -10,7 +10,7 @@ Since we are doing *Named Entity Recognition*, and we want to go with a fast mod
 let's name this experiment *fast-ner*:
 
 <ReactTermynal>
-  <span data-ty="input">classy train token data/train.tsv -n fast-ner --profile distilbert</span>
+  <span data-ty="input">classy train token data/ner-data/ --exp-name fast-ner --profile distilbert</span>
   <span data-ty="progress"></span>
   <span data-ty>Training completed</span>
 </ReactTermynal>
@@ -22,7 +22,7 @@ let's name this experiment *fast-ner*:
 
 Your model and experiment data will be saved in `experiments/<exp-name>/YYYY-MM-DD/HH-mm-ss/`.
 
-`classy` automatically saves *best* and *last* checkpoints, as well as pre- and post-trainer initialization configurations.
+`classy` automatically saves *best* and *last* checkpoints, as well as the training configuration.
 
 :::
 
diff --git a/docs/docs/getting-started/customizing-things/config.md b/docs/docs/getting-started/customizing-things/config.md
index 8540a335..eb3dd03c 100644
--- a/docs/docs/getting-started/customizing-things/config.md
+++ b/docs/docs/getting-started/customizing-things/config.md
@@ -1,11 +1,11 @@
 ---
 sidebar_position: 0
-title: classy Configuration
+title: Classy Configuration
 ---
 
 Behind the scenes, ` classy` works with `.yaml` configuration files to configure your experiment. This allows to modify
 its behavior without having to change a single line of code! You can use `classy train [...] --print` to visualize the
-underlying configuration structure (there's actually more to it, but we'll come back later).
+underlying configuration structure (we'll come back to it later with more details).
 
 ![Classy Train Print - Token](/img/intro/classy-train-print-tok.png)
 
@@ -20,7 +20,7 @@ You can change the config from the command line by using the `-c` parameter of `
 :::tip
 
 For example, say you want to specify the min and max number of steps your train can run, you can do that by adding
-`-c trainer.pl_trainer.max_steps=X trainer.pl_trainer.min_steps=Y` to your `classy train` command.
+`-c trainer.pl_trainer.min_steps=X trainer.pl_trainer.max_steps=Y` to your `classy train` command.
 
 :::
 
diff --git a/docs/docs/getting-started/customizing-things/custom-data-format.md b/docs/docs/getting-started/customizing-things/custom-data-format.md
index a05a41c7..aa3f969d 100644
--- a/docs/docs/getting-started/customizing-things/custom-data-format.md
+++ b/docs/docs/getting-started/customizing-things/custom-data-format.md
@@ -8,6 +8,7 @@ You just need to implement your own data driver and register it:
 
 ```python
 # implement your data driver
+@DataDriver.register(YOUR_TASK, YOUR_FILE_EXTENSION)
 class CustomDataDriver(DataDriver):
     def read(self, lines: Iterator[str]) -> Iterator[ClassySample]:
         raise NotImplementedError
@@ -16,12 +17,9 @@ class CustomDataDriver(DataDriver):
         self,
         samples: Iterator[ClassySample],
         path: str,
+        use_predicted_annotation: bool = False,
     ):
         raise NotImplementedError
-
-
-# register it
-READERS_DICT[(YOUR_TASK, YOUR_FILE_EXTENSION)] = CustomDataDriver
 ```
 
 :::caution
@@ -36,6 +34,7 @@ postpending file extensions is mandatory, even on Unix systems.
 For instance, imagine you were to reimplement the `.jsonl` data driver for Sequence Classification:
 
 ```python
+@DataDriver.register(SEQUENCE, "jsonl")
 class JSONLSequenceDataDriver(SequenceDataDriver):
     pass
 ```
diff --git a/docs/docs/getting-started/customizing-things/custom-metric.md b/docs/docs/getting-started/customizing-things/custom-metric.md
index d771e175..21ab9815 100644
--- a/docs/docs/getting-started/customizing-things/custom-metric.md
+++ b/docs/docs/getting-started/customizing-things/custom-metric.md
@@ -14,6 +14,7 @@ Adding a custom metric for evaluation is easy in `classy`, and you can use it fo
 class Evaluation:
     def __call__(
         self,
+        path: str,
         predicted_samples: List[ClassySample],
     ) -> Dict:
         raise NotImplementedError
diff --git a/docs/docs/getting-started/customizing-things/custom-model.md b/docs/docs/getting-started/customizing-things/custom-model.md
index f4f73856..8d4bbda9 100644
--- a/docs/docs/getting-started/customizing-things/custom-model.md
+++ b/docs/docs/getting-started/customizing-things/custom-model.md
@@ -12,7 +12,7 @@ before proceeding.
 
 :::
 
-Implementing your own model within classy is easy. You just need to:
+Implementing your own model within `classy` is easy. You just need to:
 * subclass `ClassyPLModule` and your task mixin (*SequenceTask*, *SentencePairTask*, *TokensTask*, *QATask*)
 * implement abstract methods
 * (optional) override any other method
diff --git a/docs/docs/getting-started/customizing-things/custom-optimizer.md b/docs/docs/getting-started/customizing-things/custom-optimizer.md
index b31dafc8..161c40cb 100644
--- a/docs/docs/getting-started/customizing-things/custom-optimizer.md
+++ b/docs/docs/getting-started/customizing-things/custom-optimizer.md
@@ -64,7 +64,7 @@ To use it, put the following yaml lines in your own profile or config.
 ```yaml
 model:
   optim_conf:
-    _target_: classy.optim.factories.AdamWWithWarmupFactory
+    _target_: classy.optim.factories.AdafactorWithWarmupFactory
     lr: 2e-5
     warmup_steps: 5000
     total_steps: ${training.pl_trainer.max_steps}
diff --git a/docs/docs/getting-started/customizing-things/template.md b/docs/docs/getting-started/customizing-things/template.md
index fe5c205e..3f97c91a 100644
--- a/docs/docs/getting-started/customizing-things/template.md
+++ b/docs/docs/getting-started/customizing-things/template.md
@@ -3,7 +3,7 @@ sidebar_position: 1
 title: Classy Template
 ---
 
-Welcome! If you're reading this, it's probably because you want to do something that `classy` does not support natively.
+If you're reading this, it's probably because you want to do something that `classy` does not support natively.
 
 We have prepared a [GitHub template](https://github.com/sunglasses-ai/classy-template) that you can use to start things off without having to worry about folder structure and so on.
 
diff --git a/docs/docs/installation.md b/docs/docs/installation.md
index a8ccf11a..76106696 100644
--- a/docs/docs/installation.md
+++ b/docs/docs/installation.md
@@ -7,7 +7,7 @@ title: Installation
 We strongly recommend using [Conda](https://conda.io/) as the environment manager when dealing with deep learning / data science / machine learning.
 :::
 
-`classy` requires Python 3.7 or later, and is built on [PyTorch Lightning](https://pytorchlightning.ai/).
+`classy` requires Python 3.8 or later, and is built on [PyTorch Lightning](https://pytorchlightning.ai/).
 It's recommended that you install the PyTorch ecosystem **before** installing `classy` by following the instructions on [pytorch.org](https://pytorch.org/).
 
 Or, simply put:
@@ -30,11 +30,12 @@ The preferred way to install `classy` is via `pip`. Just run `pip install classy
 
 [Conda](https://conda.io/) can be used set up a virtual environment with the
 version of Python required for `classy`.  If you already have a Python 3
-environment you want to use, you can skip to the 'installing via pip' section.
+environment you want to use, you can skip to the
+[Installing the library and dependencies](/docs/installation/#installing-the-library-and-dependencies) section.
 
 1.  [Download and install Conda](https://conda.io/projects/conda/en/latest/user-guide/install/index.html).
 
-2.  Create a Conda environment with Python 3.7-3.9:
+2.  Create a Conda environment with Python 3.8+:
 
     ```yaml
     conda create -n classy python=3.8
diff --git a/docs/docs/intro.md b/docs/docs/intro.md
index 0b8739b5..2b1e2716 100644
--- a/docs/docs/intro.md
+++ b/docs/docs/intro.md
@@ -8,7 +8,7 @@ import ReactTermynal from '/src/components/termynal';
 
 `classy` is a simple-to-use library for building high-performance Machine Learning models in NLP.
 It wraps the best libraries around ([PyTorch](https://pytorch.org/), [PyTorch Lightning](https://www.pytorchlightning.ai/),
-[Transformers](https://huggingface.co/transformers/), [Streamlit](https://streamlit.io/), ...)
+[Transformers](https://huggingface.co/transformers/), [Hydra](https://hydra.cc), [Streamlit](https://streamlit.io/), ...)
 and offers them to users with a simple CLI interface.
 
 ## Installation
@@ -81,7 +81,7 @@ organizing data, that classy expects you to do. We'll go back to this later on.
 
 Now that we have trained our model, `classy` offers a number of commands that let you use it:
 
-- `classy demo` spawns a [Streamlit](https://streamlit.io/)-based interface that lets you query your model visually;
-- `classy predict` lets you choose whether to use your model for file prediction or in bash-interactive mode (similarly to `demo`, but in the terminal);
-- `classy serve` spawns a REST API (through [FastAPI](https://fastapi.tiangolo.com/)) that you can query using any REST-compatible client (`requests`, `curl`, you name it);
-- `classy evaluate` runs your model on a given test set and reports metrics on the task.
+- [`classy demo`](/docs/reference-manual/cli/inference/#demo) spawns a [Streamlit](https://streamlit.io/)-based interface that lets you query your model visually;
+- [`classy predict`](/docs/reference-manual/cli/predict) lets you choose whether to use your model for file prediction or in bash-interactive mode (similarly to `demo`, but in the terminal);
+- [`classy serve`](/docs/reference-manual/cli/inference/#serve) spawns a REST API (through [FastAPI](https://fastapi.tiangolo.com/)) that you can query using any REST-compatible client (`requests`, `curl`, you name it);
+- [`classy evaluate`](/docs/reference-manual/cli/inference/#evaluate) runs your model on a given test set and reports metrics on the task.
diff --git a/docs/docs/reference-manual/cli/describe.md b/docs/docs/reference-manual/cli/describe.md
index 83e643ef..026e75a8 100644
--- a/docs/docs/reference-manual/cli/describe.md
+++ b/docs/docs/reference-manual/cli/describe.md
@@ -3,6 +3,12 @@ sidebar_position: 5
 title: describe
 ---
 
+:::info
+
+You need to have installed `classy` with describe support. You can do this with `pip install classy-core[describe]`
+
+:::
+
 `classy describe` is a data analysis tool with general and task-specific stats over your dataset.
 
 Say that you have a Sequence Classification dataset, and you want some preliminary information about the data itself
diff --git a/docs/docs/reference-manual/cli/export.md b/docs/docs/reference-manual/cli/export.md
new file mode 100644
index 00000000..beba7203
--- /dev/null
+++ b/docs/docs/reference-manual/cli/export.md
@@ -0,0 +1,37 @@
+---
+sidebar_position: 5
+title: import/export
+---
+
+import ReactTermynal from '/src/components/termynal';
+
+### Export
+
+You can use `classy export` to save a zip file of a trained model to be shared with others.
+
+
+<ReactTermynal>
+  <span data-ty="input">classy train [...] -n modelname</span>
+  <span data-ty="input">classy export modelname</span>
+  <span data-ty="progress"></span>
+  <span data-ty>Model saved to classy-export-modelname.zip</span>
+</ReactTermynal>
+
+### Import
+
+You can use `classy import` to import a model previously exported by someone else with `classy export`. By default,
+everything will be imported under your `experiments/` folder (if you run from a classy project),
+but you can specify a target path via the `--exp-dir` argument.
+
+
+<ReactTermynal>
+  <span data-ty="input">classy import classy-export-modelname.zip</span>
+  <span data-ty="input">classy serve modelname</span>
+  <span data-ty data-ty-start-delay="2000">Demo up and running at http://0.0.0.0:8000</span>
+</ReactTermynal>
+
+<br />
+
+:::tip
+Once a model has been imported, it is accessible from anywhere on the machine using any of the model-based `classy` commands (e.g. `serve`, `predict`, etc).
+:::
diff --git a/docs/docs/reference-manual/cli/inference.md b/docs/docs/reference-manual/cli/inference.md
index 6a2ccfe6..3577588f 100644
--- a/docs/docs/reference-manual/cli/inference.md
+++ b/docs/docs/reference-manual/cli/inference.md
@@ -32,6 +32,12 @@ path explicitly should solve the issue.
 
 ## Serve
 
+:::info
+
+You need to have installed `classy` with serve support. You can do this with `pip install classy-core[serve]`
+
+:::
+
 You can use `classy serve` to expose your model via a REST API with [FastAPI](https://fastapi.tiangolo.com/).
 
 <ReactTermynal>
@@ -56,6 +62,12 @@ By default, `classy serve` uses port 8000. Use the *-p* parameter to specify a d
 
 ## Demo
 
+:::info
+
+You need to have installed `classy` with demo support. You can do this with `pip install classy-core[demo]`
+
+:::
+
 You can use `classy demo` to spawn a [Streamlit](https://streamlit.io/) demo of your model.
 
 <ReactTermynal>
diff --git a/docs/docs/reference-manual/cli/train.md b/docs/docs/reference-manual/cli/train.md
index e56f379f..d0a9f862 100644
--- a/docs/docs/reference-manual/cli/train.md
+++ b/docs/docs/reference-manual/cli/train.md
@@ -32,22 +32,45 @@ consider moving only one of them, for instance the best one, *checkpoints/best.c
 
 :::
 
-Note that *&lt;path-to-dataset&gt;* is a bit of a special parameter and can be either:
+**Note** that *&lt;path-to-dataset&gt;* is a bit of a special parameter and can be either:
 * a **folder**, or, actually, a *ML-ready* folder: that is, it must contain a training file, named *train.#*
   (# denotes a classy supported extension), and, optionally, a validation file and a test file, *validation.#* and *test.#*
 * a **file**: classy uses the file provided to automatically generate an ML-ready folder for you, storing it in the
   *data/* folder inside the current experiment
+* a **yaml file**: classy reads the yaml file and infer the datasets involved in the training process from it. It should be organized as follows:
+  ```yaml
+  # Optional parameter that tells classy which is your favourite format.
+  # If classy has to save something, it will save it in the format you specify,
+  #  using the corresponding DataDriver.
+  main_data_driver: "json"
+
+  # You can specify a split using a dictionary path -> extension
+  #  that tells classy the dataset involved and which DataDriver
+  #  to use for each dataset.
+  train_dataset:
+    "first-train-dataset-path.tsv": "tsv"
+    "second-train-dataset-path.json": "my-special-format"
+
+  # You can also specify a split as a list of paths, and the DataDriver
+  #  to use for each dataset will be inferred from the extension.
+  validation_dataset:
+    - "fist-validation-dataset.tsv"
+    - "second-validation-dataset.tsv"
+
+  # And finally, you can specify a split as single path
+  test_dataset: "test-dataset.json"
+  ```
 
 :::info
 
-In the folder case, if *validation.#* is not present, classy automatically generates it by reserving some samples
+In both _folder_ and _yaml file_ cases, if the validation dataset is not specified, classy automatically generates it by reserving some samples
 from the training file: the new training and validation files are saved in the *data/* folder inside the current experiment.
-This does not hold for *test.#* and, if not present, classy **will not** create it.
+This does not hold for the test set and, if not present, classy **will not** create it.
 
 :::
 
 <ReactTermynal>
-  <span data-ty="input">classy train sequence data/output.tsv -n sequence-example</span>
+  <span data-ty="input">classy train sequence data/sequence-dataset -n sequence-example</span>
   <span data-ty="progress"></span>
   <span data-ty>Training completed</span>
   <span data-ty="input">ls experiments</span>
@@ -55,14 +78,16 @@ This does not hold for *test.#* and, if not present, classy **will not** create
   <span data-ty="input">tree -L 2 sequence-example</span>
   <span className="data-ty-treefix" data-ty>experiments/sequence-example/&lt;day&gt;/&lt;time&gt;/
     <div>
-        {`├── checkpoints
-            │   ├── best.ckpt
-            │   ├── epoch=00-val_loss=0.51.ckpt
-            │   └── last.ckpt
-            ├── data
-            │   ├── train.tsv
-            │   └── validation.tsv
-            └── ...
+        {`├── .hydra
+          │   └── config.yaml
+          ├── checkpoints
+          │   ├── best.ckpt
+          │   ├── epoch=00-val_loss=0.51.ckpt
+          │   └── last.ckpt
+          ├── data
+          │   ├── train.shuffled.tsv
+          │   └── validation.tsv
+          └── ...
         `.split('\n').map( (it, i) => <p key={i} style={{lineHeight: "1.0"}}>{it}</p>)}
     </div>
   </span>
diff --git a/docs/docs/reference-manual/structured-configs/overall-structure.md b/docs/docs/reference-manual/structured-configs/overall-structure.md
index 835cbe74..446638f2 100644
--- a/docs/docs/reference-manual/structured-configs/overall-structure.md
+++ b/docs/docs/reference-manual/structured-configs/overall-structure.md
@@ -106,7 +106,7 @@ Besides, there are 5 root `.yaml` files already defined, one for each task:
 
 :::info
 
-If you manually inspect the `configurations/` folder, you'll see there are 3 additional folders: `device/`, `logging/` and
+If you manually inspect the `configurations/` folder, you'll see there are 3 additional folders: `logging/` and
 `profiles/`. You can safely ignore the first two, and we cover `profiles/` in depth [here](/docs/getting-started/customizing-things/changing-profile).
 
 :::
diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js
index 5d1d7c24..b2ad207a 100644
--- a/docs/docusaurus.config.js
+++ b/docs/docusaurus.config.js
@@ -22,7 +22,7 @@ const darkCodeTheme = require('prism-react-renderer/themes/dracula');
         docs: {
           sidebarPath: require.resolve('./sidebars.js'),
           // Please change this to your repo.
-          editUrl: ({docPath}) => docPath.startsWith('api/') ? undefined : `https://github.com/sunglasses-ai/classy/edit/main/docs/${docPath}`,
+          editUrl: ({docPath}) => docPath.startsWith('api/') ? undefined : `https://github.com/sunglasses-ai/classy/edit/main/docs/docs/${docPath}`,
           // editUrl: 'https://github.com/sunglasses-ai/classy/edit/main/docs/',
           include: [
             "/glossary/*.md",
diff --git a/docs/local-dev.sh b/docs/local-dev.sh
index 094c0779..edb3e566 100644
--- a/docs/local-dev.sh
+++ b/docs/local-dev.sh
@@ -28,8 +28,7 @@ function ctrl_c() {
 
 docker exec $container_id bash -c "
   cd /classy && \
-  pip install -r requirements.txt && \
-  pip install -r <(python3 classy/optional_deps.py) && \
+  pip install -e .[all] && \
   pdoc -f --template-dir docs/pdoc/templates -o docs/docs classy && \
   rm -rf docs/docs/api && \
   mv docs/docs/classy docs/docs/api && \
diff --git a/extra-requirements.txt b/extra-requirements.txt
new file mode 100644
index 00000000..5b6ca46b
--- /dev/null
+++ b/extra-requirements.txt
@@ -0,0 +1,8 @@
+# FORMAT: put your extra requirements here in the following format
+# package[version_required]: tag1, tag2, ...
+streamlit>=1.4.0,<1.5.0	demo,describe
+st-annotated-text==2.0.0	demo,describe
+fastapi==0.68.1	serve
+uvicorn[standard]==0.15.0	serve
+plotly==5.5.0	describe
+pdoc3==0.10.0	docs
diff --git a/requirements.txt b/requirements.txt
index 116e4fdc..7e280363 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,19 +1,18 @@
-torch==1.9.0  # todo change to >=
-pytorch-lightning==1.5.9
-torchmetrics==0.5.1
-transformers==4.11.3
+torch>=1.9
+pytorch-lightning>=1.5.0,<1.6.0
+torchmetrics>=0.5.0,<0.8.0
+transformers>=4.1,<4.18
 hydra-core==1.1.1
 # argcomplete
 argcomplete==1.12.3
-# datasets and metrics
-datasets==1.12.1
-seqeval==1.2.2
-scikit-learn==0.24.2
-matplotlib==3.4.3
-# optional(?)
-rich==11.0.0
-sentencepiece==0.1.96  # required by unigram tokenizers (mbart)
-wandb==0.12.2  # required for logging
-nltk==3.6.5
+# logging
+wandb>=0.12.0
+# data and metrics
+datasets>=1.12,<1.19
+nltk>=3.4.5,<4.0.0
+seqeval==1.2.*  # required by span callback
 rouge-score==0.0.4  # required by rouge callback
 sacrebleu==2.0.0  # required by sacrebleu callback
+# extra
+rich>=11.0.0
+sentencepiece==0.1.96  # required by unigram tokenizers (mbart)
diff --git a/setup.py b/setup.py
index 1a3a1535..6b530b0c 100644
--- a/setup.py
+++ b/setup.py
@@ -1,3 +1,6 @@
+import collections
+import re
+
 from setuptools import setup, find_packages, find_namespace_packages
 
 # PEP0440 compatible formatted version, see:
@@ -16,13 +19,35 @@
 
 # version.py defines the VERSION and VERSION_SHORT variables.
 # We use exec here so we don't import classy whilst setting up.
+
 VERSION = {}  # type: ignore
 with open("classy/version.py", "r") as version_file:
     exec(version_file.read(), VERSION)
 
 
+# read requirements
+requirements = []
 with open("requirements.txt") as f:
-    requirements = f.readlines()
+    for line in f:
+        line = line.strip()
+        if line and not line.startswith("#"):
+            requirements.append(line)
+
+
+# read extra requirements
+extra_requirements = collections.defaultdict(set)
+with open("extra-requirements.txt") as f:
+    for line in f:
+        line = line.strip()
+        if line and not line.startswith("#"):
+            k, tags = line, set()
+            if '\t' in k:
+                k, _tags = k.split('\t')
+                tags.update(_tag.strip() for _tag in _tags.split(','))
+            tags.add(re.split("[<=>]", k)[0])
+            for tag in tags:
+                extra_requirements[tag].add(k)
+extra_requirements["all"] = set([req for reqs in extra_requirements.values() for req in reqs])
 
 setup(
     name="classy-core",
@@ -49,6 +74,7 @@
     packages=find_packages(),
     package_data={"configurations": ["*.yaml", "*/*.yaml"]},
     install_requires=requirements,
+    extras_require=extra_requirements,
     entry_points={"console_scripts": ["classy=classy.scripts.cli.__init__:main"]},
     python_requires=">=3.8.0",
     zip_safe=False,
diff --git a/setup.sh b/setup.sh
index 481024a5..1478c9a0 100644
--- a/setup.sh
+++ b/setup.sh
@@ -7,14 +7,13 @@ source ~/miniconda3/etc/profile.d/conda.sh
 
 ## create conda env
 read -rp "Enter environment name: " env_name
-read -rp "Enter python version (e.g. >=3.8): " python_version
+read -rp "Enter python version (must be >=3.8): " python_version
 conda create -yn "$env_name" python="$python_version"
 conda activate "$env_name"
 
 # install torch
-read -rp "Enter torch version (recommended 1.9.0): " torch_version
-echo "Not sure which CUDA version you have? Check out https://stackoverflow.com/a/68499241/1908499"
-read -rp "Enter cuda version (10.2, 11.3 or none to avoid installing cuda support): " cuda_version
+read -rp "Enter torch version (check it is compatible with all your requirements): " torch_version
+read -rp "Enter cuda version (10.2, 11.3 or none to avoid installing cuda support. Not sure? Check out https://stackoverflow.com/a/68499241/1908499): " cuda_version
 if [ "$cuda_version" == "none" ]; then
     conda install -y pytorch=$torch_version torchvision cpuonly -c pytorch
 else
@@ -34,6 +33,8 @@ if [[ $REPLY =~ ^[Yy]$ ]]
 then
   # install black
   pip install black
+  # install pytest
+  pip install pytest
   # install and hook pre-commit
   pip install pre-commit
   pre-commit install


Test patch:
diff --git a/tests/test_profile_application.py b/tests/test_profile_application.py
deleted file mode 100644
index 60763f7d..00000000
--- a/tests/test_profile_application.py
+++ /dev/null
@@ -1,257 +0,0 @@
-import shutil
-from pathlib import Path
-
-import hydra
-import pytest
-from omegaconf import OmegaConf
-
-import classy
-from classy.scripts.cli.train import apply_profile_on_dir
-
-
-@pytest.fixture
-def hydra_test_env(tmpdir):
-
-    tmpdir = Path(tmpdir.strpath)
-
-    OmegaConf.save(
-        OmegaConf.create(
-            {
-                "i": 0,
-                "l": [0, 0],
-                "d": {"a": "a", "b": "b", "c": "c"},
-                "o": {"_target_": "my.custom.module", "p1": "${interpolation}"},
-                "interpolation": "interpolation",
-                "defaults": [{"x": "a"}, {"zl": "a"}, {"zo": "a"}, "_self_"],
-            }
-        ),
-        tmpdir / "root.yaml",
-    )
-
-    (tmpdir / "x").mkdir()
-    OmegaConf.save(OmegaConf.create({"a": "0"}), tmpdir / "x" / "a.yaml")
-    OmegaConf.save(OmegaConf.create({"b": "0"}), tmpdir / "x" / "b.yaml")
-    OmegaConf.save(OmegaConf.create({"c": "0"}), tmpdir / "x" / "c.yaml")
-
-    (tmpdir / "y").mkdir()
-    OmegaConf.save(OmegaConf.create({"a": "0"}), tmpdir / "y" / "a.yaml")
-    OmegaConf.save(OmegaConf.create({"b": "0"}), tmpdir / "y" / "b.yaml")
-
-    (tmpdir / "zl").mkdir()
-    OmegaConf.save(OmegaConf.create([0, 1]), tmpdir / "zl" / "a.yaml")
-
-    (tmpdir / "zo").mkdir()
-    OmegaConf.save(OmegaConf.create({"_target_": "my.custom.module", "p1": "p1", "defaults": [{"x": "a"}, "_self_"]}), tmpdir / "zo" / "a.yaml")
-
-    (tmpdir / "zo" / "x").mkdir()
-    OmegaConf.save(OmegaConf.create({"a": "0"}), tmpdir / "zo" / "x" / "a.yaml")
-    OmegaConf.save(OmegaConf.create({"b": "0"}), tmpdir / "zo" / "x" / "b.yaml")
-
-    return str(tmpdir), "root"
-
-
-class TestProfileApplication:
-    def _load_hydra_conf(self, hydra_config_dir, hydra_config_name, cli):
-        with hydra.initialize_config_dir(config_dir=hydra_config_dir, job_name="test"):
-            return hydra.compose(config_name=hydra_config_name, overrides=cli)
-
-    def _run_skeleton(self, hydra_config_dir, hydra_config_name, profile, cli):
-
-        before_cfg = self._load_hydra_conf(hydra_config_dir, hydra_config_name, cli)
-
-        apply_profile_on_dir(profile, profile_name="", config_name=hydra_config_name, config_dir=hydra_config_dir)
-
-        after_cfg = self._load_hydra_conf(hydra_config_dir, hydra_config_name, cli)
-
-        return before_cfg, after_cfg
-
-    def test_empty_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({}), cli=[])
-        assert before_cfg == after_cfg
-
-    def test_root_primitive_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"i": 1}), cli=[])
-
-        assert before_cfg.i != after_cfg.i
-        assert after_cfg.i == 1
-
-        # check that the above change was the only one applied
-        after_cfg.i = before_cfg.i
-        assert before_cfg == after_cfg
-
-    def test_none_on_root_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"o": None}), cli=[])
-
-        assert before_cfg.o != after_cfg.o
-        assert after_cfg.o is None
-
-        # check that the above change was the only one applied
-        after_cfg.o = before_cfg.o
-        assert before_cfg == after_cfg
-
-    def test_root_list_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"l": [1]}), cli=[])
-
-        assert before_cfg.l != after_cfg.l
-        assert after_cfg.l == [1]
-
-        # check that the above change was the only one applied
-        after_cfg.l = before_cfg.l
-        assert before_cfg == after_cfg
-
-    def test_root_variable_in_list_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"l.1": [1]}), cli=[])
-
-        assert before_cfg.l != after_cfg.l
-        assert after_cfg.l[1] == [1]
-
-        # check that the above change was the only one applied
-        after_cfg.l = before_cfg.l
-        assert before_cfg == after_cfg
-
-    def test_root_dict_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(
-            *hydra_test_env, profile=OmegaConf.create({"d": {"a": "_a", "d": "_d"}}), cli=[]
-        )
-
-        assert before_cfg.d != after_cfg
-
-        # check replacement
-        assert before_cfg.d.a != after_cfg.d.a
-        assert after_cfg.d.a == "_a"
-
-        # check addittivity (b has remained identical)
-        assert before_cfg.d.b == after_cfg.d.b
-
-        # check new insertion
-        assert "d" not in before_cfg.d
-        assert "d" in after_cfg.d
-
-        # check that the above change was the only one applied
-        after_cfg.d = before_cfg.d
-        assert before_cfg == after_cfg
-
-    def test_root_obj_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(
-            *hydra_test_env, profile=OmegaConf.create({"o": {"_target_": "my.other.custom.module", "p2": "p2"}}), cli=[]
-        )
-
-        assert before_cfg.d != after_cfg
-
-        # check _target has changed correctly
-        assert before_cfg.o._target_ != after_cfg.o._target_
-        assert after_cfg.o._target_ == "my.other.custom.module"
-
-        # check p1 has been removed
-        assert "p1" in before_cfg.o
-        assert "p1" not in after_cfg.o
-
-        # check p2 has been added
-        assert "p2" not in before_cfg.o
-        assert "p2" in after_cfg.o
-
-    def test_root_additional_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"s": "s"}), cli=[])
-
-        assert "s" not in before_cfg
-        assert after_cfg.s == "s"
-
-        # check that the above change was the only one applied
-        del after_cfg.s
-        assert before_cfg == after_cfg
-
-    def test_root_interpolation_variable_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(
-            *hydra_test_env, profile=OmegaConf.create({"interpolation": "another-interpolation"}), cli=[]
-        )
-
-        assert before_cfg.interpolation != after_cfg.interpolation
-        assert after_cfg.interpolation == "another-interpolation"
-        assert after_cfg.o.p1 == "another-interpolation"
-
-        # check that the above change was the only one applied
-        # this works also here, with an interpolation, as no materialization took place
-        after_cfg.interpolation = before_cfg.interpolation
-        assert before_cfg == after_cfg
-
-    def test_dict_override_in_config_group_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"x": {"a": "1"}}), cli=[])
-
-        assert before_cfg.x.a == "0"
-        assert after_cfg.x.a == "1"
-
-        # check that the above change was the only one applied
-        # this works also here, with an interpolation, as no materialization took place
-        after_cfg.x.a = before_cfg.x.a
-        assert before_cfg == after_cfg
-
-    def test_list_override_in_config_group_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"zl": [1]}), cli=[])
-
-        assert before_cfg.zl == [0, 1]
-        assert after_cfg.zl == [1]
-
-        # check that the above change was the only one applied
-        # this works also here, with an interpolation, as no materialization took place
-        after_cfg.zl = before_cfg.zl
-        assert before_cfg == after_cfg
-
-    def test_object_override_in_config_group_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(
-            *hydra_test_env,
-            profile=OmegaConf.create({"zo": {"_target_": "my.other.custom.module", "p2": "p2", "x": "b"}}),
-            cli=[]
-        )
-
-        assert before_cfg.zo != after_cfg.zo
-        assert after_cfg.zo == {"_target_": "my.other.custom.module", "p2": "p2", "x": {"b": "0"}}
-
-        # check that the above change was the only one applied
-        # this works also here, with an interpolation, as no materialization took place
-        after_cfg.zo = before_cfg.zo
-        assert before_cfg == after_cfg
-
-    def test_config_group_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"x": "b"}), cli=[])
-
-        assert before_cfg.x != after_cfg.x
-
-        # check a was deleted
-        assert "a" in before_cfg.x
-        assert "a" not in after_cfg.x
-
-        # check b was introduced
-        assert "b" not in before_cfg.x
-        assert "b" in after_cfg.x
-
-        # check that the above change was the only one applied
-        after_cfg.x = before_cfg.x
-        assert before_cfg == after_cfg
-
-    def test_new_config_group_set_from_profile(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"y": "a"}), cli=[])
-
-        assert before_cfg != after_cfg
-
-        assert "y" not in before_cfg
-        assert "y" in after_cfg
-        assert after_cfg.y.a == "0"
-
-        # check that the above change was the only one applied
-        del after_cfg.y
-        assert before_cfg == after_cfg
-
-    def test_variable_intersection_between_profile_and_cli(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"i": "1"}), cli=["i=2"])
-
-        assert after_cfg.i == 2
-
-    def test_config_group_intersection_between_profile_and_cli(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({"x": "b"}), cli=["x=c"])
-
-        assert after_cfg.x == OmegaConf.create({"c": "0"})
-
-    def test_config_group_set_from_cli(self, hydra_test_env):
-        before_cfg, after_cfg = self._run_skeleton(*hydra_test_env, profile=OmegaConf.create({}), cli=["+y=a"])
-
-        assert after_cfg.y == OmegaConf.create({"a": "0"})
diff --git a/tests/test_profiles.py b/tests/test_profiles.py
new file mode 100644
index 00000000..b92fb853
--- /dev/null
+++ b/tests/test_profiles.py
@@ -0,0 +1,318 @@
+from pathlib import Path
+
+import hydra
+import pytest
+
+from classy.scripts.cli.train import apply_profiles_and_cli, ClassyBlame
+
+
+def parse_test_spec(spec: str, hydra_test_env):
+    path, content = None, []
+    cli = []
+    for line in spec.split("\n"):
+        line = line.strip()
+        if line == "":
+            if path is not None:
+                assert len(content) > 0
+                if path == "cli":
+                    for l in content:
+                        assert not l.startswith(" ") or l.startswith("\t")
+                        cli.append(l.replace(": ", "="))
+                else:
+                    output_path = f"{hydra_test_env}/{path}"
+                    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+                    with open(output_path, "w") as f:
+                        f.write("\n".join(content) + "\n")
+                path, content = None, []
+        elif line.startswith("|"):
+            content.append(line[2:])
+        else:
+            path, _content = line.split("|")
+            path = path.rstrip()
+            content.append(_content[1:])
+    return cli
+
+
+@pytest.fixture
+def hydra_test_env(tmpdir):
+    tmpdir = Path(tmpdir.strpath)
+    (tmpdir / "profiles").mkdir()
+    return str(tmpdir)
+
+
+class TestProfiles:
+
+    def _compute_final_config(self, spec, hydra_test_env):
+        cli = parse_test_spec(spec, hydra_test_env=hydra_test_env)
+        blames = apply_profiles_and_cli(
+            config_name="root",
+            config_dir=hydra_test_env,
+            profile_path=f"{hydra_test_env}/profiles/profile.yaml",
+            cli={ClassyBlame("cli"): cli},
+            output_config_name="merged"
+        )
+        with hydra.initialize_config_dir(config_dir=hydra_test_env, job_name="test"):
+            cfg = hydra.compose(config_name="merged")
+        return cfg, blames
+
+    def test_profile_overrides_value(self, hydra_test_env):
+        """
+        Check profile can override values
+        """
+        spec = """
+        root.yaml               | a: 0
+
+        profiles/profile.yaml   | a: 1
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a == 1
+
+    def test_profile_adds_value(self, hydra_test_env):
+        """
+        Check profile can add values
+        """
+        spec = """
+        root.yaml               | a: 0
+
+        profiles/profile.yaml   | b: 1
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg == {"a": 0, "b": 1}
+
+    def test_profile_and_cli_override_value(self, hydra_test_env):
+        """
+        Check cli edits are prioritized over profile's
+        """
+        spec = """
+        root.yaml               | a: 0
+
+        profiles/profile.yaml   | a: 1
+
+        cli                     | a: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a == 2
+
+    def test_profile_override_target(self, hydra_test_env):
+        """
+        Check that, when profile changes _target_, profile's node entirely overwrites the original one
+        (b is discarded and _target_ replaced)
+        """
+        spec = """
+        root.yaml               | defaults:
+                                | - a: a1
+
+        a/a1.yaml               | _target_: A1
+                                | b: 1
+
+        profiles/profile.yaml   | a:
+                                |   _target_: A2
+                                |   c: 1
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a._target_ == "A2"
+        assert "b" not in cfg.a
+        assert cfg.a.c == 1
+
+    def test_profile_override_not_on_target(self, hydra_test_env):
+        """
+        Check that, when profile does not change _target_, profile's node only updates the original one
+        (_target_ is unchanged and b updated)
+        """
+        spec = """
+        root.yaml               | defaults:
+                                | - a: a1
+
+        a/a1.yaml               | _target_: A1
+                                | b: 1
+
+        profiles/profile.yaml   | a:
+                                |   b: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a._target_ == "A1"
+        assert cfg.a.b == 2
+
+    def test_profile_override_not_on_target_from_file(self, hydra_test_env):
+        """
+        Check that, when profile specifies a dict override from file (not on _target), profile's node only updates
+        original one (_target_ is unchanged and b updated)
+        """
+        spec = """
+        root.yaml               | defaults:
+                                | - a: a1
+
+        a/a1.yaml               | _target_: A1
+                                | b: 1
+
+        a/a2.yaml               | b: 2
+
+        profiles/profile.yaml   | a: a2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a.b == 2
+
+    def test_profile_override_list(self, hydra_test_env):
+        """
+        Check that, when profile specifies a list, it is appended to the corresponding original node
+        """
+        spec = """
+        root.yaml               | defaults:
+                                | - a: a1
+
+        a/a1.yaml               | - a: 0
+                                | - b: 0
+
+        profiles/profile.yaml   | a:
+                                |   - c: 0
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a == [{"a": 0}, {"b": 0}, {"c": 0}]
+
+    def test_profile_override_list_from_file(self, hydra_test_env):
+        """
+        Check that, when profile specifies a list override from file, it is appended to the corresponding original node
+        """
+        spec = """
+        root.yaml               | defaults:
+                                | - a: a1
+
+        a/a1.yaml               | - a: 0
+                                | - b: 0
+
+        a/a2.yaml               | - c: 0
+
+        profiles/profile.yaml   | a: a2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a == [{"a": 0}, {"b": 0}, {"c": 0}]
+
+    def test_interpolation_on_value(self, hydra_test_env):
+        """
+        Check interpolation is preserved after profile application
+        """
+        spec = """
+        root.yaml               | a: 0
+                                | b: ${a}
+
+        profiles/profile.yaml   | a: 1
+
+        cli                     | a: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.b == 2
+
+    def test_interpolation_node_with_value_override(self, hydra_test_env):
+        """
+        Check that standard profile rules apply also on interpolated nodes (value override)
+        """
+        spec = """
+        root.yaml               | a: 0
+                                | b:
+                                |   _target_: B1
+                                |   c: 1
+                                |   d: ${a}
+                                | e: ${b}
+
+        profiles/profile.yaml   | e:
+                                |   c: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.b.c == 1
+        assert cfg.e.c == 2
+
+    def test_interpolation_node_with_interpolation_preserved(self, hydra_test_env):
+        """
+        Check that standard profile rules apply also on interpolated nodes (interpolation preserved)
+        """
+        spec = """
+        root.yaml               | a: 0
+                                | b:
+                                |   _target_: B1
+                                |   c: 1
+                                |   d: ${a}
+                                | e: ${b}
+
+        profiles/profile.yaml   | e:
+                                |   c: 2
+
+        cli                     | a: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.b.d == 2
+        assert cfg.e.d == 2
+
+    def test_interpolation_node_with_interpolation_discarded(self, hydra_test_env):
+        """
+        Check that standard profile rules apply also on interpolated nodes (interpolation discarded)
+        """
+        spec = """
+        root.yaml               | a: 0
+                                | b:
+                                |   _target_: B1
+                                |   c: 1
+                                |   d: ${a}
+                                | e: ${b}
+
+        profiles/profile.yaml   | e:
+                                |   c: 2
+                                |   d: 4
+
+        cli                     | a: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.b.d == 2
+        assert cfg.e.d == 4
+
+    def test_file_on_cli(self, hydra_test_env):
+        """
+        Check that, when cli overrides contain a file, it is handled as if it was written as part of profile
+        """
+        spec = """
+        root.yaml               | v: 0
+                                |
+                                | defaults:
+                                | - a: a1
+
+        a/a1.yaml               | b: 1
+                                | c: 1
+
+        a/a2.yaml               | c: 3
+
+        profiles/profile.yaml   | v: 1
+                                | a:
+                                |   b: 2
+
+        cli                     | a: a2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.v == 1
+        assert cfg.a == {"b": 1, "c": 3}
+
+    def test_profile_override_to_null(self, hydra_test_env):
+        """
+        Check that, when cli overrides something to null, it is properly handled
+        """
+        spec = """
+        root.yaml               | a:
+                                |   b: 1
+                                |   c: 2
+
+        profiles/profile.yaml   | a: null
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a is None
+
+    def test_profile_override_from_null(self, hydra_test_env):
+        """
+        Check that, when cli overrides null to something, it is properly handled
+        """
+        spec = """
+        root.yaml               | a: null
+
+        profiles/profile.yaml   | a:
+                                |   b: 1
+                                |   c: 2
+        """
+        cfg, blames = self._compute_final_config(spec, hydra_test_env)
+        assert cfg.a == {"b": 1, "c": 2}