Repository: obss/trapper (language: python)

Issue:
 Implement CONLL2003 POS tagging as an example for trapper
This will be used for showing how trapper can be used as a library in dowstream tasks


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index b57c33e..bf59bc9 100644
--- a/README.md
+++ b/README.md
@@ -1,21 +1,112 @@
 # TRAPPER (Transformer wRAPPER)
 
-A lightweight framework that aims to make it easier to train transformer based
-models on downstream tasks. It wraps the `HuggingFace`'s
+An NLP library that aims to make it easier to train transformer based models on
+downstream tasks. `trapper` wraps the HuggingFace's
 `transformers` library to provide the transformer model implementations and training
-mechanisms conveniently.
+mechanisms. It defines abstractions with base classes for common tasks encountered
+while using transformer models. Additionally, it provides a dependency-injection
+mechanism and allows defining training and/or evaluation experiments via
+configuration files. By this way, you can replicate your experiment with different
+models, optimizers etc by only changing their values inside the configuration file
+without writing any new code or changing the existing code. These features foster
+code reuse, less boiler-plate code, as well as repeatable and better documented
+training experiments which is crucial in machine learning.
 
-* `allennlp`'s registry mechanism is used to provide dependency injection and enable
-  reading the experiment details from training configuration files which are
-  typically `json` or `jsonnet` files. Moreover, registrable base classes are
-  implemented to abstract away the common operations for data processing and model
-  training.
+## Key Features
 
-* Auto classes from `transformers` are used to provide polymorphism and make it
-  possible to instantiate the actual task-specific classes (e.g. for models and
-  tokenizers) from the configuration files dynamically.
+### Compatibility with HuggingFace Transformers
 
-## Table of Currently Supported Tasks and Models From Transformers
+**trapper extends transformers!**
+
+We implement the trapper components by trying to use the available components of
+the `transformers` library as much as we can. For example, trapper uses the models
+and trainer as they are in transformers. This make it easy to use the models trained
+with trapper on another projects or libraries that depend on transformers library (
+or pytorch in general).
+
+We strive to keep trapper fully compatible with transformers so you can always use
+some of our components to write a script for your own needs while not using the full
+pipeline (e.g. for training).
+
+### Dependency Injection and Config File based Training
+
+We use `allennlp`'s registry mechanism to provide dependency injection and enable
+reading the experiment details from training configuration files in `json`
+or `jsonnet` format. You can look at the
+[allennlp guide on dependency injection](https://guide.allennlp.org/using-config-files)
+to learn more about how the registry system and dependency injection works as well
+as how to write configuration files. In addition, we strongly recommend reading the
+remaining parts of the [allennlp guide](https://guide.allennlp.org/)
+to learn more about its design philosophy, the importance of abstractions etc.
+(especially Part2: Abstraction, Design and Testing). As a warning, please note that
+we do not use allennlp's abstractions and base classes in general, which means you
+can not mix and match the trapper's and allennlp's components. Instead, we just use
+the class registry and dependency injection mechanisms and only adapt its very
+limited set of components, first by wrapping and registering them as trapper
+components. For example, we use the optimizers from allennlp since we can
+conveniently do so without hindering our full compatibility with transformers.
+
+### Full Integration with HuggingFace datasets
+
+In trapper, we officially use the format of the datasets from the HuggingFace's
+`datasets` library and provide full integration with it. You can directly use all
+datasets published in [datasets hub](https://huggingface.co/datasets) without doing
+any extra work. You can write the dataset name and extra loading arguments (if there
+are any) in your training config file, and trapper will automatically download the
+dataset and pass it to the trainer. If you have a local or private dataset, you can
+still use it after converting it to the HuggingFace `datasets` format by writing a
+dataset loading script as explained
+[here](https://huggingface.co/docs/datasets/dataset_script.html).
+
+### Abstractions and Base Classes
+
+Following allennlp, we implement our own registrable base classes to abstract away
+the common operations for data processing and model training.
+
+* Data reading and preprocessing base classes including
+
+    - The classes to be used directly: `DatasetReader`, `DatasetLoader`
+      and `DataCollator`.
+
+    - The classes that you may need to extend: `LabelMapper`,`DataProcessor`,
+      and `DataAdapter`.
+
+    - `TokenizerWrapper` classes utilizing `AutoTokenizer` from transformers are
+      used as factories to instantiate wrapped tokenizers into which task-specific
+      special tokens are registered automatically.
+
+
+* `ModelWrapper` classes utilizing the `AutoModelFor...` classes from transformers
+  are used as factories to instantiate the actual task-specific models from the
+  configuration files dynamically.
+
+
+* Optimizers from allennlp: Implemented as children of the base `Optimizer` class.
+
+## Usage
+
+AutoModel, task name etc.
+
+### Modeling the Problem
+
+The first step in using trapper is to decide on how to model the problem.
+
+### Modeling the Input
+
+Special BOS, EOS tokens and task-specific extra tokens.
+
+## Currently Supported Tasks and Models From Transformers
+
+Hypothetically, nearly all models should work on any task if it has an entry in the
+AutoModel table for that task. However, since some models require more (
+or less) parameters compared to most models, you might get errors in some models. In
+these edge cases, we try to support them by adding the extra parameters they
+require. Feel free to open an issue/PR if you encounter/solve that in a model-task
+combination. We have used trapper on a limited set of model-task combinations so
+far. We list these combinations below to indicate that they have been tested and
+validated to work without problems.
+
+### Table of Model-task Combinations Tested so far
 
 | model       | question_answering | token_classification |
 |-------------|--------------------|----------------------|
@@ -24,9 +115,49 @@ mechanisms conveniently.
 | ELECTRA     | &#10004;           | &#10004;             |
 | RoBERTa     | &#10004;           | &#10004;             |
 
-## Usage
+To use trapper on training, evaluation on a task that is not readily supported in
+transformers library, you need to extend the provided base classes according to your
+own needs. These are as follows:
+
+**For Training & Evaluation**: DataProcessor, DataAdapter, TokenizerFactory.
+
+**For Inference**: In addition to the ones listed above, you may need to implement
+a `transformers.Pipeline` or directly use form the transformers library if they
+already implemented one that matches your need.
+
+**Typically Extended Classes**
+
+1) **LabelMapper**:
+   Used in tasks that require mapping between categorical labels and integer ids
+   such as token classification.
+
+
+2) **DataProcessor**:
+   This class is responsible for taking a single instance in dict format, typically
+   coming from a `datasets.Dataset`, extracting the information fields suitable for
+   the task and hand, and converting their values to integers or collections of
+   integers. This includes, tokenizing the string fields, and getting the token ids,
+   converting the categoric labels to integer ids and so on.
 
-WIP
+
+3) **DataAdapter**:
+   This is responsible for converting the information fields inside an instance dict
+   that was previously processed by a `DataProcessor` to a format suitable for
+   feeding into a transformer model. This also includes handling the special tokens
+   signaling the start or end of a sequence, the separation of tho sequence for a
+   sequence-pair task as well as chopping excess tokens etc.
+
+
+4) **TokenizerWrapper**:
+   This class wraps a pretrained tokenizer from the transformers library while also
+   recording the special tokens needed for the task to the wrapped tokenizer. It
+   also stores the missing values from BOS - CLS, EOS - SEP token pairs for the
+   tokenizers that only support one of them. This means you can model your input
+   sequence by using the bos_token for start and eos_token for end without thinking
+   which model you are working with. If your task and input modeling needs extra
+   special tokens e.g. the `<CONTEXT>` for a context dependent task, you can store
+   these tokens by setting the `_TASK_SPECIFIC_SPECIAL_TOKENS` class variable in
+   your TokenizerWrapper subclass. Otherwise, you can directly use TokenizerWrapper.
 
 #### Registering classes from custom modules to the library
 
@@ -90,11 +221,35 @@ trapper run test_experiment.jsonnet \
 --include-package ner.data.dataset_readers
 ```
 
-## Contributing
+#### Running a training and/or evaluation experiment
 
-PRs are always welcome :)
+##### Config File Based Training Using the CLI
 
-### Installation
+Go to your project root and execute the `trapper run` command with a config file
+specifying the details of the training experiment. E.g.
+
+```shell
+trapper run SOME_DIRECTORY/experiment.jsonnet
+```
+
+Don't forget to provide the args["output_dir"] and args["result_dir"] values in your
+experiment file. Please look at the `examples/pos_tagging/README.md` for a detailed
+example.
+
+##### Scrip Based Training
+
+Go to your project root and execute the `trapper run` command with a config file
+specifying the details of the training experiment. E.g.
+
+```shell
+trapper run SOME_DIRECTORY/experiment.jsonnet
+```
+
+Don't forget to provide the args["output_dir"] and args["result_dir"] values in your
+experiment file. Please look at the `examples/pos_tagging/README.md` for a detailed
+example.
+
+## Installation
 
 #### Environment Creation
 
@@ -119,6 +274,11 @@ show below or any other branch / tag / commit you would like.
 pip install git+ssh://github.com/obss/trapper.git
 ```
 
+## Contributing
+
+If you want to open a PR, please create a fresh environment as described before,
+clone the repo and install trapper in editable mode.
+
 #### Installing in Editable Mode
 
 You can clone the repo locally and install the package to your environment as shown
@@ -130,13 +290,12 @@ cd trapper
 pip install -e .[dev]
 ```
 
-### Tests
+### Testing trapper
 
 #### Caching the test fixtures from the HuggingFace's datasets library
 
-To speed up the data-related tests and enable accessing the fixture datasets by
-their folder name, we cache the test dataset fixtures from HuggingFace's datasets
-library using the following command.
+To speed up the data-related tests, we cache the test dataset fixtures from
+HuggingFace's datasets library using the following command.
 
 ```console
 python -m scripts.cache_hf_datasets_fixtures
@@ -148,10 +307,11 @@ Then, you can simply test with the following command:
 python -m scripts.run_tests
 ```
 
-**NOTE:** To significantly speed up the tests, you can set the following environment
-variables which makes HuggingFace's transformers and datasets libraries work in
-offline mode. However, beware that you may need to run the tests once first without
-setting these environment variables so that the models, tokenizers etc. are cached.
+**NOTE:** To significantly speed up the tests depending on HuggingFace's
+transformers and datasets libraries, you can set the following environment variables
+to make them work in offline mode. However, beware that you may need to run the
+tests once first without setting these environment variables so that the pretrained
+models, tokenizers etc. are downloaded and cached.
 
 ```shell
 export TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1
diff --git a/examples/pos_tagging/.gitignore b/examples/pos_tagging/.gitignore
new file mode 100644
index 0000000..33a96f3
--- /dev/null
+++ b/examples/pos_tagging/.gitignore
@@ -0,0 +1,12 @@
+# Ignore the following folders
+results/
+outputs/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# datasets library's dataset artifacts
+*.lock
+*dataset_infos.json
diff --git a/examples/pos_tagging/.trapper_plugins b/examples/pos_tagging/.trapper_plugins
new file mode 100644
index 0000000..85de9cf
--- /dev/null
+++ b/examples/pos_tagging/.trapper_plugins
@@ -0,0 +1,1 @@
+src
diff --git a/examples/pos_tagging/README.md b/examples/pos_tagging/README.md
new file mode 100644
index 0000000..648e801
--- /dev/null
+++ b/examples/pos_tagging/README.md
@@ -0,0 +1,189 @@
+# POS Tagging in CONLL2003 Dataset
+
+This project show how to train a transformer model from on CONLL2003 dataset
+from `HuggingFace datasets`. You can explore the dataset
+from [its page](https://huggingface.co/datasets/conll2003). This project is intended
+to serve as a demo for using trapper as a library to train and evaluate a
+transformer model on a custom task/dataset as well as perform inference using it. We
+start by creating a fresh python environment and install the dependencies.
+
+## Environment Creation and Dependency Installation
+
+It is strongly recommended creating a virtual environment using conda or virtualenv
+etc. before installing the dependencies for this example. For example, the following
+code creates a conda environment with name post_tagging and python version 3.7.10,
+and activates it.
+
+```console
+conda create --name post_tagging python=3.7.10
+conda activate post_tagging
+```
+
+Then, you can install the dependencies using pip as follows:
+
+```console
+pip install -r requirements.txt
+```
+
+## Task Modeling
+
+To use trapper, you need to select the common NLP formulation of the problem you are
+tackling as well as decide on its input representation, including the special
+tokens.
+
+First, you need to model your problem as one of the common modeling tasks in NLP
+such as seq-to-seq, sequence classification etc. We stick with the transformers' way
+of dividing the tasks into common categories as it does in its `AutoModelFor...`
+classes. To be compatible with transformers and reuse its model factories, trapper
+formalizes the tasks by wrapping the `AutoModelFor...` classes and matching them to
+a name that represent a common task in NLP. In our case, the natural choice for POS
+tagging is to model it as a token classification (i.e. sequence labeling) task.
+
+To form our input sequence, we will convert the words into tokens to get the
+sequence. Then, prepend it with the special `bos_token` and append the special
+`eos_token` to the end of the sequence. So the format of our instances wil
+be `BOS ... tokens ... EOS`.
+
+### Custom Class Implementations
+
+We implement four data-related classes under `src` directory and register them to
+trapper's registry. In fact, the tokenizer wrapper class is the same as the base
+class, we just implemented it to show how you can implement and register a custom
+tokenizer wrapper in your own tasks.
+
+#### ExampleLabelMapperForPosTagging
+
+Implemented in `src/label_mapper.py`. It is responsible for handling the mapping
+between POS tags in CONLL2003  (such as 'FW', 'IN' etc) and their integer ids which
+are fed to the model during training. Registered with the name of
+`"conll2003_pos_tagging_example"` to the registry.
+
+#### ExamplePosTaggingTokenizerWrapper
+
+Implemented in `src/tokenizer_wrapper.py`. Although we could have used the
+`TokenizerWrapper` directly, this class is implemented to demonstrate how to
+subclass the TokenizerWrapper and register the new class.
+
+#### ExampleConll2003PosTaggingDataProcessor
+
+Implemented in `src/data_processor.py`. It extracts the `tokens`,
+`pos_tags` and `id` fields from an input data instance. It tokenizes the
+`tokens` field since it actually consists of words which may need further
+tokenization. Then, it generates the corresponding token ids and store them.
+Finally, the`pos_tags` are stored directly without any processing since this field
+consists of integer labels ids instead of categorical labels.
+
+#### ExampleDataAdapterForPosTagging
+
+Implemented in `src/data_adapter.py`. It takes a processed instance dict from the
+data processor and creates a new dict that has the `input_ids` and `labels` keys
+required by the models. It also takes care of the special BOS and EOS tokens while
+constructing these fields.
+
+## Training
+
+We encourage writing configuration files and start the training/evaluation
+experiments from the trapper CLI although you can write your custom training script
+as well. For this project, we will show how to use trapper using the CLI and a
+config file.
+
+### Configuration File
+
+We use `allennlp`'s registry mechanism to provide dependency injection and enable
+reading the experiment details from training configuration files in `json`
+or `jsonnet` format. You can look at the
+[allennlp guide on dependency injection](https://guide.allennlp.org/using-config-files)
+to learn more about how the registry system and dependency injection works as well
+as how to write configuration files.
+
+In our config file, we specify all the details of an experiment, including model
+type, training arguments, optimizer etc. The config file we have written is
+`experiments/roberta/experiment.json`. As can be seen, we put it under `roberta`
+directory and specify the model type as `"roberta-base"` (in
+`pretrained_model_name_or_path` argument). You can use whatever directory structure
+you like for your experiments.
+
+#### Config File Field Explanations
+
+`pretrained_model_name_or_path`: The pretrained model name or path from the
+transformers library. This value is used for creating the pretrained model and
+tokenizer inside their wrapper objects.
+
+`train_split_name`: The split name that will be used while getting the training
+dataset from the full dataset loaded via dataset reader.
+
+`dev_split_name`: The split name that will be used while getting the validation
+dataset from the full dataset loaded via dataset reader.
+
+`tokenizer_wrapper`: The TokenizerWrapper parameters. The generated TokenizerWrapper
+object will be reused by other objects that needs it via the dependency-injection
+mechanism.
+
+`dataset_loader`: We use the base DatasetLoader class since it is sufficient for
+most purposes. However, we need to specify its internal objects which we have
+implemented before. These include `dataset_reader`, `data_processor` and
+`data_adapter`.
+
+`dataset_reader`: The DatasetReader construction parameters. As the argument of
+`"path"` parameter, provide `conll2003_test_fixture` local variable while testing
+the project whereas `"conll2003"` for the actual training.
+
+`data_processor`: The DatasetReader construction parameters. Only specifying the
+registered type is sufficient.
+
+`data_adapter`: The DatasetAdapter construction parameters. Only specifying the
+registered type is sufficient.
+
+`data_collator`: The DataCollator construction parameters. We use the base class
+which is the default registered one without any extra parameter.
+
+`model_wrapper`: The ModelWrapper construction parameters. We specify the type of
+the wrapper as  `"token_classification"` since we model the POS tagging problem as a
+token classification task. Moreover, we provide a `num_labels`
+argument with value of 47 as there are 47 labels in CONLL2003 POS tagging dataset.
+This value will be used for modifying the output layer of the pretrained model to
+make its output dimension size equal to the number of labels in the dataset if
+needed.
+
+`args`: The TransformerTrainingArguments construction parameters. Includes batch
+sizes, experiment output/result and results directories, flag variables to decide on
+performing train and/or evaluation and so on.
+
+`optimizer`: The Optimizer construction parameters. We use Huggingface's AdamW
+Optimizer via allennlp.
+
+### Starting the Experiment via trapper CLI
+
+From the project root, go to the POS tagging project directory with the following
+command:
+
+```shell
+cd examples/pos_tagging
+```
+
+Then, execute the `trapper run` command with a config file specifying the details of
+the training experiment.
+
+```shell
+trapper run experiments/roberta/experiment.jsonnet
+```
+
+If you use the provided experiment config file, two directories, named `outputs`
+and `results` will be created corresponding to the args["output_dir"] and
+args["result_dir"] values in the config file. You can change their values to point
+them to different directories, but must specify both of them (to non-existent or
+empty directories). `outputs` should have the `logs` subdirectory containing the
+training logs and `checkpoints` subdirectory containing the saved model weights.
+The `results` directory stores the train and/or evaluation results, the trainer
+state, as well as the model and the tokenizer.
+
+## Testing Our Project
+
+We have implemented tests checking the correctness of custom class we wrote. From
+the root of the example project (i.e. the `pos_tagging` directory), you can run the
+`run_tests` script as follows.
+
+```console
+cd examples/pos_tagging
+python -m scripts.run_tests
+```
diff --git a/examples/pos_tagging/experiments/roberta/experiment.jsonnet b/examples/pos_tagging/experiments/roberta/experiment.jsonnet
new file mode 100644
index 0000000..7804ec3
--- /dev/null
+++ b/examples/pos_tagging/experiments/roberta/experiment.jsonnet
@@ -0,0 +1,59 @@
+local output_dir = "experiments/roberta/outputs";
+local result_dir = "experiments/roberta/results";
+local conll2003_test_fixture="test_fixtures/hf_datasets/conll2003_test_fixture";
+{
+        "pretrained_model_name_or_path": "roberta-base",
+        "train_split_name": "train",
+        "dev_split_name": "validation",
+        "tokenizer_wrapper": {
+            "type": "pos_tagging_example",
+            "model_max_sequence_length": 512,
+            "add_prefix_space": true
+        },
+        "dataset_loader": {
+            "dataset_reader": {
+                "path": "conll2003",  # actual dataset
+//                "path": conll2003_test_fixture,  # for testing the project
+            },
+            "data_processor": {"type": "conll2003_pos_tagging_example"},
+            "data_adapter": {"type": "conll2003_pos_tagging_example"},
+        },
+        "data_collator": {"type": "default"},
+        "model_wrapper": {"type": "token_classification", "num_labels": 47},
+        "compute_metrics": {"type": "seqeval",
+                            "return_entity_level_metrics": false},
+        "label_mapper": {"type": "conll2003_pos_tagging_example"},
+        "args": {
+            "type": "default",
+            "output_dir": output_dir + "/checkpoints",
+            "result_dir": result_dir,
+            "num_train_epochs": 3,
+            "per_device_train_batch_size": 1,
+            "per_device_eval_batch_size": 1,
+            "logging_dir": output_dir + "/logs",
+            "no_cuda": true,
+            "logging_steps": 1,
+            "evaluation_strategy": "steps",
+            "save_steps": 2,
+            "label_names": ["labels"],
+            "lr_scheduler_type": "linear",
+            "warmup_steps": 2,
+            "do_train": true,
+            "do_eval": true,
+            "save_total_limit": 1,
+            "metric_for_best_model": "eval_loss",
+            "greater_is_better": false,
+        },
+        "optimizer": {
+            "type": "huggingface_adamw",
+            "weight_decay": 0.01,
+            "parameter_groups": [
+                [
+                    ["bias", "LayerNorm\\\\.weight", "layer_norm\\\\.weight"],
+                    {"weight_decay": 0},
+                ]
+            ],
+            "lr": 5e-5,
+            "eps": 1e-8,
+        },
+    }
diff --git a/examples/pos_tagging/requirements.txt b/examples/pos_tagging/requirements.txt
new file mode 100644
index 0000000..56fac74
--- /dev/null
+++ b/examples/pos_tagging/requirements.txt
@@ -0,0 +1,1 @@
+trapper==0.0.2
diff --git a/examples/pos_tagging/scripts/__init__.py b/examples/pos_tagging/scripts/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/pos_tagging/scripts/cache_hf_datasets_fixtures.py b/examples/pos_tagging/scripts/cache_hf_datasets_fixtures.py
new file mode 100644
index 0000000..a5fc083
--- /dev/null
+++ b/examples/pos_tagging/scripts/cache_hf_datasets_fixtures.py
@@ -0,0 +1,11 @@
+"""
+Caches the tests dataset to HuggingFace's `datasets` library's cache so that the
+interpreter can find it when we try to load it through the `datasets` library.
+"""
+from examples.pos_tagging.src import POS_TAGGING_FIXTURES_ROOT
+from trapper.common.testing_utils.hf_datasets_caching import (
+    renew_hf_datasets_fixtures_cache,
+)
+
+if __name__ == "__main__":
+    renew_hf_datasets_fixtures_cache(POS_TAGGING_FIXTURES_ROOT / "hf_datasets")
diff --git a/examples/pos_tagging/src/__init__.py b/examples/pos_tagging/src/__init__.py
new file mode 100644
index 0000000..3f67b16
--- /dev/null
+++ b/examples/pos_tagging/src/__init__.py
@@ -0,0 +1,13 @@
+from pathlib import Path
+
+from examples.pos_tagging.src.data_adapter import ExampleDataAdapterForPosTagging
+from examples.pos_tagging.src.data_processor import (
+    ExampleConll2003PosTaggingDataProcessor,
+)
+from examples.pos_tagging.src.tokenizer_wrapper import (
+    ExamplePosTaggingTokenizerWrapper,
+)
+
+POS_TAGGING_PROJECT_ROOT = Path(__file__).parent.parent.resolve()
+POS_TAGGING_TESTS_ROOT = POS_TAGGING_PROJECT_ROOT / "tests"
+POS_TAGGING_FIXTURES_ROOT = POS_TAGGING_PROJECT_ROOT / "test_fixtures"
diff --git a/examples/pos_tagging/src/data_adapter.py b/examples/pos_tagging/src/data_adapter.py
new file mode 100644
index 0000000..b1ea63d
--- /dev/null
+++ b/examples/pos_tagging/src/data_adapter.py
@@ -0,0 +1,35 @@
+from examples.pos_tagging.src.label_mapper import ExampleLabelMapperForPosTagging
+from trapper.common.constants import IGNORED_LABEL_ID
+from trapper.data import TokenizerWrapper
+from trapper.data.data_adapters.data_adapter import DataAdapter
+from trapper.data.data_processors import IndexedInstance
+
+
+@DataAdapter.register("conll2003_pos_tagging_example")
+class ExampleDataAdapterForPosTagging(DataAdapter):
+    """
+    This class takes the processed instance dict from the data processor and
+    creates a new dict that has the "input_ids" and "labels" keys required by the
+    models. It also takes care of the special BOS and EOS tokens while constructing
+    these fields.
+    """
+
+    def __init__(
+            self,
+            tokenizer_wrapper: TokenizerWrapper,
+    ):
+        label_mapper = ExampleLabelMapperForPosTagging.from_labels()
+        super().__init__(tokenizer_wrapper, label_mapper)
+
+    def __call__(self, raw_instance: IndexedInstance) -> IndexedInstance:
+        """
+        Create a sequence with the following field:
+        input_ids: <bos> ...tokens... <eos>
+        """
+        # We return an instance having the keys fields specified in
+        # `trapper.models.auto_wrappers._TASK_TO_INPUT_FIELDS["token_classification"]`
+        input_ids = [self._bos_token_id] + raw_instance["tokens"]
+        input_ids.append(self._eos_token_id)
+        labels = [IGNORED_LABEL_ID] + raw_instance["pos_tags"]
+        labels.append(IGNORED_LABEL_ID)
+        return {"input_ids": input_ids, "labels": labels}
diff --git a/examples/pos_tagging/src/data_processor.py b/examples/pos_tagging/src/data_processor.py
new file mode 100644
index 0000000..1b94b84
--- /dev/null
+++ b/examples/pos_tagging/src/data_processor.py
@@ -0,0 +1,49 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+from trapper.data.data_processors import DataProcessor
+from trapper.data.data_processors.data_processor import IndexedInstance
+
+logger = logging.getLogger(__file__)
+
+
+@DataProcessor.register("conll2003_pos_tagging_example")
+class ExampleConll2003PosTaggingDataProcessor(DataProcessor):
+    """
+    This class extracts the "tokens", "pos_tags" and "id" fields from from an input
+    data instance. It tokenizes the `tokens` field since it actually consists of
+    words which may need further tokenization. Then, it generates the corresponding
+    token ids and store them. Finally, the`pos_tags` are stored directly without
+    any processing since this field consists of integer labels ids instead of
+    categorical labels.
+    """
+    NUM_EXTRA_SPECIAL_TOKENS_IN_SEQUENCE = 2  # <bos> tokens <eos>
+
+    def process(self, instance_dict: Dict[str, Any]) -> Optional[IndexedInstance]:
+        return self.text_to_instance(
+            id_=instance_dict["id"],
+            tokens=instance_dict["tokens"],
+            pos_tags=instance_dict["pos_tags"],
+        )
+
+    def text_to_instance(
+            self,
+            id_: str,
+            tokens: List[str],
+            pos_tags: Optional[List[int]] = None,
+    ) -> IndexedInstance:
+        expanded_tokens = []
+        expanded_pos_tags = []
+        for token, pos_tag in zip(tokens, pos_tags):
+            expanded_token = self.tokenizer.tokenize(token)
+            expanded_pos_tag = [pos_tag] * len(expanded_token)
+            expanded_tokens.extend(expanded_token)
+            expanded_pos_tags.extend(expanded_pos_tag)
+
+        for seq in (expanded_tokens, expanded_pos_tags):
+            self._chop_excess_tokens(seq, len(seq))
+
+        return {
+            "tokens": self.tokenizer.convert_tokens_to_ids(expanded_tokens),
+            "pos_tags": expanded_pos_tags
+        }
diff --git a/examples/pos_tagging/src/label_mapper.py b/examples/pos_tagging/src/label_mapper.py
new file mode 100644
index 0000000..be362b6
--- /dev/null
+++ b/examples/pos_tagging/src/label_mapper.py
@@ -0,0 +1,13 @@
+from trapper.data.label_mapper import LabelMapper
+
+
+@LabelMapper.register("conll2003_pos_tagging_example", constructor="from_labels")
+class ExampleLabelMapperForPosTagging(LabelMapper):
+    # Obtained by executing `dataset["train"].features["pos_tags"].feature.names`
+    _LABELS = (
+        '"', "''", '#', '$', '(', ')', ',', '.', ':', '``', 'CC', 'CD', 'DT', 'EX',
+        'FW', 'IN', 'JJ', 'JJR', 'JJS', 'LS', 'MD', 'NN', 'NNP', 'NNPS', 'NNS',
+        'NN|SYM', 'PDT', 'POS', 'PRP', 'PRP$', 'RB', 'RBR', 'RBS', 'RP', 'SYM',
+        'TO', 'UH', 'VB', 'VBD', 'VBG', 'VBN', 'VBP', 'VBZ', 'WDT', 'WP', 'WP$',
+        'WRB'
+    )
diff --git a/examples/pos_tagging/src/tokenizer_wrapper.py b/examples/pos_tagging/src/tokenizer_wrapper.py
new file mode 100644
index 0000000..25b473e
--- /dev/null
+++ b/examples/pos_tagging/src/tokenizer_wrapper.py
@@ -0,0 +1,11 @@
+from trapper.data import TokenizerWrapper
+
+
+@TokenizerWrapper.register("pos_tagging_example", constructor="from_pretrained")
+class ExamplePosTaggingTokenizerWrapper(TokenizerWrapper):
+    """ A `tokenizer wrapper` that is used for demonstrating how to implement
+    your own by extending the base class. """
+    #  Although we could have used the `TokenizerWrapper` directly, this class is
+    #  implemented for demonstration purposes. You can override
+    #  `_TASK_SPECIFIC_SPECIAL_TOKENS` here if your task requires custom extra
+    #  special tokens.
diff --git a/examples/pos_tagging/train.py b/examples/pos_tagging/train.py
new file mode 100644
index 0000000..7e72d09
--- /dev/null
+++ b/examples/pos_tagging/train.py
@@ -0,0 +1,7 @@
+"""
+This scripts shows how to train a POS tagger using trapper. However,
+intentionally, the dependency injection mechanism is not used, instead the
+objects are instantiated directly from concrete classes suitable for the job.
+Although we recommended config file based training using the trapper CLI instead,
+you can use this file to create a custom training script for your needs.
+"""
diff --git a/examples/experiment.py b/examples/question_answering/experiment.py
similarity index 96%
rename from examples/experiment.py
rename to examples/question_answering/experiment.py
index 3685652..6151255 100644
--- a/examples/experiment.py
+++ b/examples/question_answering/experiment.py
@@ -5,7 +5,11 @@
 
 import requests
 
-from examples.util import DATASET_DIR, DEFAULT_EXTRA_VARIABLES, get_dir_from_task
+from examples.question_answering.util import (
+    DATASET_DIR,
+    DEFAULT_EXTRA_VARIABLES,
+    get_dir_from_task,
+)
 from trapper.training.train import run_experiment
 
 __arguments__ = ["config", "task", "experiment_name"]
diff --git a/examples/experiments/question-answering/experiment.jsonnet b/examples/question_answering/experiments/question-answering/experiment.jsonnet
similarity index 94%
rename from examples/experiments/question-answering/experiment.jsonnet
rename to examples/question_answering/experiments/question-answering/experiment.jsonnet
index 0078eae..3037151 100644
--- a/examples/experiments/question-answering/experiment.jsonnet
+++ b/examples/question_answering/experiments/question-answering/experiment.jsonnet
@@ -4,14 +4,14 @@ local result_dir = std.extVar("OUTPUT_PATH");
     "train_split_name": "train",
     "dev_split_name": "validation",
     "pretrained_model_name_or_path": "roberta-base",
-    "tokenizer": {
+    "tokenizer_wrapper": {
         "type": "question-answering"
     },
     "dataset_loader": {
         "type": "default",
         "dataset_reader": {
             "type": "default",
-            "path": "squad"
+            "path": "squad_qa_test_fixture"
         },
         "data_processor": {
             "type": "squad-question-answering"
@@ -23,7 +23,7 @@ local result_dir = std.extVar("OUTPUT_PATH");
     "data_collator":{
         "type": "default"
     },
-    "model": {
+    "model_wrapper": {
         "type": "question_answering"
     },
     "args": {
diff --git a/examples/question_answering.ipynb b/examples/question_answering/question_answering.ipynb
similarity index 100%
rename from examples/question_answering.ipynb
rename to examples/question_answering/question_answering.ipynb
diff --git a/examples/util.py b/examples/question_answering/util.py
similarity index 100%
rename from examples/util.py
rename to examples/question_answering/util.py
diff --git a/requirements.txt b/requirements.txt
index dfb6aa8..55a9e74 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 allennlp>=2.3.0,<2.8
-datasets>=1.8.0,<2.0
+datasets>=1.13.0,<2.0
 numpy>=1.21.2
 tensorboardX==2.1
 transformers>=4.5.0,<4.12
+seqeval>=1.2.0
 deepdiff>=5.2.0
diff --git a/trapper/common/constants.py b/trapper/common/constants.py
index 203b33a..170911f 100644
--- a/trapper/common/constants.py
+++ b/trapper/common/constants.py
@@ -13,12 +13,16 @@
 else:
     from typing_extensions import TypedDict
 
-PAD_TOKEN_LABEL_ID = -100  # automatically ignored by PyTorch loss functions
+IGNORED_LABEL_ID = -100  # automatically ignored by PyTorch loss functions
+# Common special tokens
 CLS_TOKEN = "[CLS]"
 SEP_TOKEN = "[SEP]"
 BOS_TOKEN = "<bos>"
 EOS_TOKEN = "<eos>"
 PAD_TOKEN = "<pad>"
+MASK_TOKEN = "<mask>"
+UNK_TOKEN = "<unk>"
+# Task-specific special tokens
 CONTEXT_TOKEN = "<context>"
 ANSWER_TOKEN = "<ans>"
 
diff --git a/trapper/data/__init__.py b/trapper/data/__init__.py
index 7142d02..4f79d05 100644
--- a/trapper/data/__init__.py
+++ b/trapper/data/__init__.py
@@ -8,4 +8,5 @@
 from trapper.data.data_processors.data_processor import IndexedInstance
 from trapper.data.dataset_loader import DatasetLoader
 from trapper.data.dataset_reader import DatasetReader
-from trapper.data.tokenizers import TransformerTokenizer
+from trapper.data.label_mapper import LabelMapper
+from trapper.data.tokenizers import TokenizerWrapper
diff --git a/trapper/data/data_adapters/data_adapter.py b/trapper/data/data_adapters/data_adapter.py
index 352886d..be4c926 100644
--- a/trapper/data/data_adapters/data_adapter.py
+++ b/trapper/data/data_adapters/data_adapter.py
@@ -1,8 +1,12 @@
 from abc import ABC, abstractmethod
+from typing import Optional
+
+from transformers import PreTrainedTokenizerBase
 
 from trapper.common import Registrable
 from trapper.data.data_processors.data_processor import IndexedInstance
-from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+from trapper.data.label_mapper import LabelMapper
+from trapper.data.tokenizers import TokenizerWrapper
 
 
 class DataAdapter(ABC, Registrable):
@@ -13,22 +17,45 @@ class DataAdapter(ABC, Registrable):
     and adapts the input fields by renaming them to the names accepted by the
     models e.g. "input_ids", "token_type_ids" etc. Moreover, it also handles
     the insertion of the special tokens signaling the beginning or ending of a
-    sequence such as `[CLS]`, `[SEP]` etc. You need to implement the `__call__`
-    method suitable for your task when you subclass it. See
+    sequence such as `[CLS]`, `[SEP]` etc. To extend this class, you need to
+    implement the `__call__` method as suitable for your task. See
     `DataAdapterForQuestionAnswering` for an example.
 
     Args:
-        tokenizer (): Required to access the ids of special tokens
+        tokenizer_wrapper (): Required to access the ids of special tokens such
+            as BOS, EOS etc
+        label_mapper (): Only used in some tasks that require mapping between
+            categorical labels and integer ids such as token classification.
     """
 
-    def __init__(self, tokenizer: TransformerTokenizer):
-        self._tokenizer = tokenizer
+    def __init__(
+        self,
+        tokenizer_wrapper: TokenizerWrapper,
+        label_mapper: Optional[LabelMapper] = None,
+    ):
+        tokenizer = tokenizer_wrapper.tokenizer
+        self._bos_token_id: int = tokenizer.bos_token_id
+        self._eos_token_id: int = tokenizer.eos_token_id
+        self._tokenizer: PreTrainedTokenizerBase = tokenizer
+        self._label_mapper = label_mapper
+
+    @property
+    def tokenizer(self) -> PreTrainedTokenizerBase:
+        return self._tokenizer
 
     @abstractmethod
     def __call__(self, instance: IndexedInstance) -> IndexedInstance:
         """
         Takes a raw `IndexedInstance`, performs some processing on it, and returns
-        an `IndexedInstance` again. Look at
+        an `IndexedInstance` again. Typically, you should return an instance having
+        the fields specified in
+        `trapper.models.auto_wrappers._TASK_TO_INPUT_FIELDS[task]` (except for
+        "attention_mask" since it will be generated by the DataCollator for you)
+        where `task` is your choice of modeling the downstream task you want to
+        solve using `transformers`. Each task in the `_TASK_TO_INPUT_FIELDS`
+        corresponds to an `AutoModel...` class implemented in `transformers`.
+
+        You may look at
         `DataAdapterForQuestionAnswering.__call__` for an example.
 
         Args:
diff --git a/trapper/data/data_adapters/question_answering_adapter.py b/trapper/data/data_adapters/question_answering_adapter.py
index 1bd38af..e47f3f5 100644
--- a/trapper/data/data_adapters/question_answering_adapter.py
+++ b/trapper/data/data_adapters/question_answering_adapter.py
@@ -2,7 +2,7 @@
 
 from trapper.data.data_adapters.data_adapter import DataAdapter
 from trapper.data.data_processors import IndexedInstance
-from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+from trapper.data.tokenizers import TokenizerWrapper
 
 
 @DataAdapter.register("question-answering")
@@ -12,7 +12,7 @@ class DataAdapterForQuestionAnswering(DataAdapter):
     answering tasks that involves a context, question and answer.
 
     Args:
-        tokenizer (): Required to access the ids BOS and EOS tokens
+        tokenizer_wrapper (): Required to access the ids of BOS and EOS tokens
     """
 
     CONTEXT_TOKEN_TYPE_ID = 0
@@ -20,11 +20,9 @@ class DataAdapterForQuestionAnswering(DataAdapter):
 
     def __init__(
         self,
-        tokenizer: TransformerTokenizer,
+        tokenizer_wrapper: TokenizerWrapper,
     ):
-        super().__init__(tokenizer)
-        self._bos_token_id: int = self._tokenizer.bos_token_id
-        self._eos_token_id: int = self._tokenizer.eos_token_id
+        super().__init__(tokenizer_wrapper)
 
     def __call__(self, raw_instance: IndexedInstance) -> IndexedInstance:
         """
diff --git a/trapper/data/data_collator.py b/trapper/data/data_collator.py
index 46dccca..fefc0dd 100644
--- a/trapper/data/data_collator.py
+++ b/trapper/data/data_collator.py
@@ -3,11 +3,12 @@
 
 import torch
 from torch import Tensor
+from transformers import PreTrainedTokenizerBase
 
 from trapper.common import Registrable
-from trapper.common.constants import PAD_TOKEN_LABEL_ID
+from trapper.common.constants import IGNORED_LABEL_ID
 from trapper.data.data_processors.data_processor import IndexedInstance
-from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+from trapper.data.tokenizers import TokenizerWrapper
 
 InputBatch = Dict[str, List[Union[int, List[int]]]]
 InputBatchTensor = Dict[str, Tensor]
@@ -31,10 +32,10 @@ class DataCollator(Registrable):
 
     def __init__(
         self,
-        tokenizer: TransformerTokenizer,
+        tokenizer_wrapper: TokenizerWrapper,
         model_forward_params: Tuple[str, ...],
     ):
-        self._tokenizer = tokenizer
+        self._tokenizer: PreTrainedTokenizerBase = tokenizer_wrapper.tokenizer
         self._model_forward_params: Tuple[str, ...] = model_forward_params
 
     def __call__(
@@ -123,7 +124,7 @@ def _pad_id(self, padded_field: str) -> int:
         elif padded_field == "token_type_ids":
             pad_id = self._tokenizer.pad_token_type_id
         elif padded_field == "labels":
-            pad_id = PAD_TOKEN_LABEL_ID
+            pad_id = IGNORED_LABEL_ID
         elif padded_field == "attention_mask":
             pad_id = 0
         elif padded_field == "special_tokens_mask":
diff --git a/trapper/data/data_processors/data_processor.py b/trapper/data/data_processors/data_processor.py
index e3faf3f..2daf396 100644
--- a/trapper/data/data_processors/data_processor.py
+++ b/trapper/data/data_processors/data_processor.py
@@ -1,10 +1,13 @@
 import logging
 from abc import ABCMeta, abstractmethod
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
+
+from transformers import PreTrainedTokenizerBase
 
 from trapper.common import Registrable
 from trapper.common.constants import PositionDict
-from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+from trapper.data.label_mapper import LabelMapper
+from trapper.data.tokenizers import TokenizerWrapper
 
 logger = logging.getLogger(__file__)
 
@@ -32,18 +35,45 @@ class DataProcessor(Registrable, metaclass=ABCMeta):
             tokens (do not have to be unique) in the input ids.
 
     Args:
-        tokenizer ():
+        tokenizer_wrapper ():
+        model_max_sequence_length (): The maximum length of the processed
+            sequence. Actually, the maximum sequence will be the minimum of this
+            value and the `model_max_length` value of the tokenizer.
+        label_mapper (): Only used in some tasks that require mapping between
+            categorical labels and integer ids such as token classification.
     """
 
     NUM_EXTRA_SPECIAL_TOKENS_IN_SEQUENCE = 0
 
-    def __init__(self, tokenizer: TransformerTokenizer):
-        self._tokenizer = tokenizer
+    def __init__(
+        self,
+        tokenizer_wrapper: TokenizerWrapper,
+        model_max_sequence_length: Optional[int] = None,
+        label_mapper: Optional[LabelMapper] = None,
+    ):
+        self._tokenizer: PreTrainedTokenizerBase = tokenizer_wrapper.tokenizer
+        self._model_max_sequence_length = self._find_model_max_seq_length(
+            model_max_sequence_length
+        )
+        self._label_mapper = label_mapper
 
     @property
     def tokenizer(self):
         return self._tokenizer
 
+    @property
+    def model_max_sequence_length(self) -> int:
+        return self._model_max_sequence_length
+
+    def _find_model_max_seq_length(
+        self,
+        provided_model_max_sequence_length: int = None,
+    ) -> int:
+        model_max_length = self.tokenizer.model_max_length
+        if provided_model_max_sequence_length is None:
+            return model_max_length
+        return min(model_max_length, provided_model_max_sequence_length)
+
     @abstractmethod
     def text_to_instance(self, *inputs) -> IndexedInstance:
         """
@@ -112,10 +142,24 @@ def _total_seq_len(cls, *sequences):
         total_seq_len += cls.NUM_EXTRA_SPECIAL_TOKENS_IN_SEQUENCE
         return total_seq_len
 
-    def _chop_excess_tokens(self, tokens: List, max_len: int):
+    def _chop_excess_tokens(self, sequence: List, total_len: int):
         """
-        Utilizes a  heuristic of chopping off the excess tokens in-place
-         from the end
+        Chops the excess tokens in a sequence from the right side. `total_len`
+        is the current total length computed in some way before calling this
+        function. If the caller deals with a single sequence,
+        `total_len == len(sequence)` should hold. However, if the caller will
+        append two or more sequences next to each other, `total_len` should be the
+        current total length of the sequences, and the `sequence` argument should be
+        the first one. This is useful in contextual tasks such as question
+        answering where we typically append the context and question next to each
+        other. In this case, the context should be supplied as the `sequence`
+        argument so that we chop toward the end of it (right side) without chopping
+        the question.
+
+        Args:
+            sequence (): the input sequence. If the caller deals with multiple
+                sequences, the first (leftmost) one should be passed
+            total_len (): total length before calling this function
         """
-        excess = max_len - self._tokenizer.model_max_sequence_length
-        del tokens[-1 * excess :]
+        excess = total_len - self.model_max_sequence_length
+        del sequence[-1 * excess :]
diff --git a/trapper/data/data_processors/squad/squad_processor.py b/trapper/data/data_processors/squad/squad_processor.py
index 4e1a9cc..b3dcbae 100644
--- a/trapper/data/data_processors/squad/squad_processor.py
+++ b/trapper/data/data_processors/squad/squad_processor.py
@@ -90,7 +90,7 @@ def _chop_excess_context_tokens(
     ):
         subsequences = [tokenized_context, *other_tokenized_subsequences]
         seq_len = self._total_seq_len(*subsequences)
-        if seq_len > self._tokenizer.model_max_sequence_length:
+        if seq_len > self.model_max_sequence_length:
             self._chop_excess_tokens(tokenized_context, seq_len)
             seq_len = self._total_seq_len(*subsequences)
-            assert seq_len <= self._tokenizer.model_max_sequence_length
+            assert seq_len <= self.model_max_sequence_length
diff --git a/trapper/data/dataset_reader.py b/trapper/data/dataset_reader.py
index 2dd8044..b511a48 100644
--- a/trapper/data/dataset_reader.py
+++ b/trapper/data/dataset_reader.py
@@ -1,4 +1,17 @@
-# pylint: disable=unused-argument
+# Copyright 2021 Open Business Software Solutions, the HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import inspect
 import logging
 from typing import Mapping, Optional, Sequence, Union
@@ -27,33 +40,76 @@
 
 class DatasetReader(Registrable):
     """
-    This class is used for reading a raw dataset from the `datasets` library. The
-    dataset is read and stored during `DatasetReader` object instantiation so
-    that the available splits (train, validation etc) can be queried later without
-    reading it again. The constructor arguments are directly transferred to the
-    `datasets.load_dataset` function. See `datasets.load_dataset`_ for the
-    details of these parameters and how the dataset is loaded.
+    This class is used for reading a raw dataset from the `datasets` library.
+    The constructor arguments are stored so that they can be used from the `read`
+    method to read the dataset using the `datasets.load_dataset` function . Notice
+    that the constructor does not have `split` parameter. Instead, the split is
+    specified inside the `read` method as the only parameter. See
+    `datasets.load_dataset`_ for the details of how the dataset is loaded.
 
-    .. _datasets.load_dataset: https://github.com/huggingface/datasets/blob/b057846bcc1ccfe0fda4d8d42f190d146a70ee64/src/datasets/load.py#L996
+    .. _datasets.load_dataset: https://github.com/huggingface/datasets/blob/ec824222c227ea5c9c75568d6f357a819599a6c7/src/datasets/load.py#L1472
 
     Args:
-        path ():
-        name ():
-        data_dir ():
-        data_files ():
-        split ():
-        cache_dir ():
-        features ():
-        download_config ():
-        download_mode ():
-        ignore_verifications ():
-        keep_in_memory ():
-        save_infos ():
-        script_version ():
-        use_auth_token ():
-        task ():
-        streaming ():
-        **config_kwargs ():
+
+        path (:obj:`str`): Path or name of the dataset.
+            Depending on ``path``, the dataset builder that is returned id either generic dataset builder (csv, json, text etc.) or a dataset builder defined defined a dataset script (a python file).
+
+            For local datasets:
+
+            - if ``path`` is a local directory (but doesn't contain a dataset script)
+              -> load a generic dataset builder (csv, json, text etc.) based on the content of the directory
+              e.g. ``'./path/to/directory/with/my/csv/data'``.
+            - if ``path`` is a local dataset script or a directory containing a local dataset script (if the script has the same name as the directory):
+              -> load the dataset builder from the dataset script
+              e.g. ``'./dataset/squad'`` or ``'./dataset/squad/squad.py'``.
+
+            For datasets on the Hugging Face Hub (list all available datasets and ids with ``datasets.list_datasets()``)
+
+            - if ``path`` is a canonical dataset on the HF Hub (ex: `glue`, `squad`)
+              -> load the dataset builder from the dataset script in the github repository at huggingface/datasets
+              e.g. ``'squad'`` or ``'glue'``.
+            - if ``path`` is a dataset repository on the HF hub (without a dataset script)
+              -> load a generic dataset builder (csv, text etc.) based on the content of the repository
+              e.g. ``'username/dataset_name'``, a dataset repository on the HF hub containing your data files.
+            - if ``path`` is a dataset repository on the HF hub with a dataset script (if the script has the same name as the directory)
+              -> load the dataset builder from the dataset script in the dataset repository
+              e.g. ``'username/dataset_name'``, a dataset repository on the HF hub containing a dataset script `'dataset_name.py'`.
+
+        name (:obj:`str`, optional): Defining the name of the dataset configuration.
+        data_dir (:obj:`str`, optional): Defining the data_dir of the dataset configuration.
+        data_files (:obj:`str` or :obj:`Sequence` or :obj:`Mapping`, optional): Path(s) to source data file(s).
+        split (:class:`Split` or :obj:`str`): Which split of the data to load.
+            If None, will return a `dict` with all splits (typically `datasets.Split.TRAIN` and `datasets.Split.TEST`).
+            If given, will return a single Dataset.
+            Splits can be combined and specified like in tensorflow-datasets.
+        cache_dir (:obj:`str`, optional): Directory to read/write data. Defaults to "~/.cache/huggingface/datasets".
+        features (:class:`Features`, optional): Set the features type to use for this dataset.
+        download_config (:class:`~utils.DownloadConfig`, optional): Specific download configuration parameters.
+        download_mode (:class:`GenerateMode`, default ``REUSE_DATASET_IF_EXISTS``): Download/generate mode.
+        ignore_verifications (:obj:`bool`, default ``False``): Ignore the verifications of the downloaded/processed dataset information (checksums/size/splits/...).
+        keep_in_memory (:obj:`bool`, default ``None``): Whether to copy the dataset in-memory. If `None`, the dataset
+            will not be copied in-memory unless explicitly enabled by setting `datasets.config.IN_MEMORY_MAX_SIZE` to
+            nonzero. See more details in the :ref:`load_dataset_enhancing_performance` section.
+        save_infos (:obj:`bool`, default ``False``): Save the dataset information (checksums/size/splits/...).
+        revision (:class:`~utils.Version` or :obj:`str`, optional): Version of the dataset script to load:
+
+            - For canonical datasets in the `huggingface/datasets` library like "squad", the default version of the module is the local version of the lib.
+              You can specify a different version from your local version of the lib (e.g. "master" or "1.2.0") but it might cause compatibility issues.
+            - For community provided datasets like "lhoestq/squad" that have their own git repository on the Datasets Hub, the default version "main" corresponds to the "main" branch.
+              You can specify a different version that the default "main" by using a commit sha or a git tag of the dataset repository.
+        use_auth_token (``str`` or ``bool``, optional): Optional string or boolean to use as Bearer token for remote files on the Datasets Hub.
+            If True, will get token from `"~/.huggingface"`.
+        task (``str``): The task to prepare the dataset for during training and evaluation. Casts the dataset's :class:`Features` to standardized column names and types as detailed in :py:mod:`datasets.tasks`.
+        streaming (``bool``, default ``False``): If set to True, don't download the data files. Instead, it streams the data progressively while
+            iterating on the dataset. An IterableDataset or IterableDatasetDict is returned instead in this case.
+
+            Note that streaming works for datasets that use data formats that support being iterated over like txt, csv, jsonl for example.
+            Json files may be downloaded completely. Also streaming from remote zip or gzip files is supported but other compressed formats
+            like rar and xz are not yet supported. The tgz format doesn't allow streaming.
+        script_version:
+            .. deprecated:: 1.13
+                'script_version' was renamed to 'revision' in version 1.13 and will be removed in 1.15.
+        **config_kwargs: Keyword arguments to be passed to the :class:`BuilderConfig` and used in the :class:`DatasetBuilder`.
     """
 
     default_implementation = "default"
@@ -74,10 +130,11 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         ignore_verifications: bool = False,
         keep_in_memory: Optional[bool] = None,
         save_infos: bool = False,
-        script_version: Optional[Union[str, Version]] = None,
+        revision: Optional[Union[str, Version]] = None,
         use_auth_token: Optional[Union[bool, str]] = None,
         task: Optional[Union[str, TaskTemplate]] = None,
         streaming: bool = False,
+        script_version="deprecated",
         **config_kwargs,
     ):
         self._init_params = locals()
@@ -87,12 +144,27 @@ def read(
         self, split: Optional[Union[str, Split]] = None
     ) -> Union[TrapperDataset, TrapperDatasetDict]:
         """
-        Returns the specified split of the dataset.
+        Combines the `split` argument with the previously stored parameters
+        inside the `__init__` method to load the specified split of the dataset.
+        See `datasets.load_dataset`_ for the details of how the dataset is loaded.
+
+        .. _datasets.load_dataset: https://github.com/huggingface/datasets/blob/ec824222c227ea5c9c75568d6f357a819599a6c7/src/datasets/load.py#L1472
 
         Args:
-            split (): Which split of the data to load. If None, will return a
-                `dict` with all splits e.g. "train" and "validation". If specified,
-                will return a single Dataset
+            split (:class:`Split` or :obj:`str`): Which split of the data to load.
+                If None, will return a `dict` with all splits (typically `datasets.Split.TRAIN` and `datasets.Split.TEST`).
+                If given, will return a single Dataset.
+                Splits can be combined and specified like in tensorflow-datasets.
+
+        Returns:
+            :class:`Dataset` or :class:`DatasetDict`:
+            - if `split` is not None: the dataset requested,
+            - if `split` is None, a ``datasets.DatasetDict`` with each split.
+
+            or :class:`IterableDataset` or :class:`IterableDatasetDict`: if streaming=True
+
+            - if `split` is not None: the dataset requested,
+            - if `split` is None, a ``datasets.streaming.IterableDatasetDict`` with each split.
         """
         init_params = self._init_params
         init_params["split"] = split
diff --git a/trapper/data/label_mapper.py b/trapper/data/label_mapper.py
new file mode 100644
index 0000000..644c055
--- /dev/null
+++ b/trapper/data/label_mapper.py
@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from copy import copy
+from typing import Dict, Optional, Sequence, Tuple
+
+from trapper.common import Registrable
+
+
+class LabelMapper(Registrable):
+    """
+    Used in tasks that require mapping between categorical labels and integer
+    ids such as token classification. The list of labels can be provided in the
+    constructor as well as the class variable `_LABELS`. The reason for the
+    latter is to document the labels and enable writing the labels only once for
+    some tasks that have a large number of labels, both of which aim to
+    reduce the possibility of errors while specifying the labels.
+
+    Optional class variables:
+
+    _LABELS (Tuple[str]): An optional class variable that may be used for setting
+        the list of labels in `LabelMapper.from_labels` method. Note
+        that the order you provide will be preserved while assigning ids to the
+        labels.
+
+    Args:
+        label_to_id_map (): The mapping from string labels to integer ids
+    """
+
+    default_implementation = "from_labels"
+    _LABELS: Tuple[str] = None
+
+    def __init__(self, label_to_id_map: Optional[Dict[str, int]] = None):
+        #  We need to make `label_to_id_map` optional with default of None,
+        #  since otherwise allennlp tries to invoke __init__ although we register
+        #  a classmethod as a default constructor and demand it via the "type"
+        #  parameter inside the from_params method or a config file.
+        if label_to_id_map is None:
+            raise ValueError("`label_to_id_map` can not be None!")
+        self._label_to_id_map = label_to_id_map
+        self._id_to_label = {
+            id_: label for label, id_ in self._label_to_id_map.items()
+        }
+
+    @classmethod
+    def from_labels(
+        cls, labels: Optional[Sequence[str]] = None, start_id: Optional[int] = 0
+    ) -> LabelMapper:
+        """
+        Create a `LabelMapper` from a list of labels. If `labels` is left as None,
+        the class variable `_LABELS` will be used instead. The indices will be
+        the enumeration of labels starting from `start_ind`.
+
+        Args:
+            labels (): The list of labels. Note that the order you provide will
+                be preserved while assigning ids to the labels.
+            start_id (): The start value for enumeration of label ids. By
+                default, we start from 0 and increment.
+        """
+        if labels is None:
+            labels = cls._LABELS
+        labels = tuple(labels)
+        ids = tuple(range(start_id, start_id + len(labels)))
+        label_to_id = {label: id_ for label, id_ in zip(labels, ids)}
+        return cls(label_to_id)
+
+    @property
+    def labels(self) -> Tuple[str]:
+        return tuple(self._label_to_id_map.keys())
+
+    @property
+    def ids(self) -> Tuple[int]:
+        return tuple(self._label_to_id_map.values())
+
+    @property
+    def label_to_id_map(self) -> Dict[str, int]:
+        return copy(self._label_to_id_map)
+
+    @property
+    def id_to_label_map(self) -> Dict[int, str]:
+        """This method may be used by the pipeline object for inference."""
+        return copy(self._id_to_label)
+
+    def get_id(self, label: str) -> int:
+        return self._label_to_id_map[label]
+
+    def get_label(self, id_: int) -> str:
+        return self._id_to_label[id_]
+
+
+LabelMapper.register("from_label_to_id_map")(LabelMapper)
+LabelMapper.register("from_labels", constructor="from_labels")(LabelMapper)
diff --git a/trapper/data/tokenizers/__init__.py b/trapper/data/tokenizers/__init__.py
index 2b820be..c28a2c7 100644
--- a/trapper/data/tokenizers/__init__.py
+++ b/trapper/data/tokenizers/__init__.py
@@ -1,2 +1,2 @@
-from trapper.data.tokenizers.squad import QuestionAnsweringTokenizer
-from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+from trapper.data.tokenizers.squad import QuestionAnsweringTokenizerWrapper
+from trapper.data.tokenizers.tokenizer_wrapper import TokenizerWrapper
diff --git a/trapper/data/tokenizers/squad.py b/trapper/data/tokenizers/squad.py
index 337cf8d..93f7f54 100644
--- a/trapper/data/tokenizers/squad.py
+++ b/trapper/data/tokenizers/squad.py
@@ -1,17 +1,12 @@
-from trapper.common.constants import BOS_TOKEN, CONTEXT_TOKEN, EOS_TOKEN, PAD_TOKEN
-from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+from trapper.common.constants import CONTEXT_TOKEN
+from trapper.data.tokenizers.tokenizer_wrapper import TokenizerWrapper
 
 
-@TransformerTokenizer.register("question-answering", constructor="from_pretrained")
-class QuestionAnsweringTokenizer(TransformerTokenizer):
+@TokenizerWrapper.register("question-answering", constructor="from_pretrained")
+class QuestionAnsweringTokenizerWrapper(TokenizerWrapper):
     """
     This tokenizer can be used in SQuAD style question answering tasks that
     utilizes a context, question and answer.
     """
 
-    _attr_to_special_token = {"additional_special_tokens": [CONTEXT_TOKEN]}
-    _TASK_SPECIFIC_SPECIAL_TOKENS = {
-        "bos_token": BOS_TOKEN,
-        "eos_token": EOS_TOKEN,
-        "pad_token": PAD_TOKEN,
-    }
+    _TASK_SPECIFIC_SPECIAL_TOKENS = [CONTEXT_TOKEN]
diff --git a/trapper/data/tokenizers/tokenizer.py b/trapper/data/tokenizers/tokenizer.py
deleted file mode 100644
index d9b05e6..0000000
--- a/trapper/data/tokenizers/tokenizer.py
+++ /dev/null
@@ -1,110 +0,0 @@
-from copy import deepcopy
-from typing import Dict, List, Tuple, Union
-
-from transformers import (
-    AutoTokenizer,
-    PreTrainedTokenizer,
-    PreTrainedTokenizerBase,
-    PreTrainedTokenizerFast,
-)
-
-from trapper.common import Registrable
-from trapper.common.utils import add_property
-
-_PreTrainedTransformerTokenizer = Union[
-    PreTrainedTokenizerFast, PreTrainedTokenizer
-]
-
-
-class TransformerTokenizer(Registrable, PreTrainedTokenizerBase):
-    """
-    The base tokenizer class for trapper that extends `PreTrainedTokenizerBase`
-    with further attributes such as task-specific special tokens and the maximum
-    sequence length accepted by the model for that tokenizer. To extend this for
-    your custom tasks, just override the `_attr_to_special_token` and
-    `_TASK_SPECIFIC_SPECIAL_TOKENS` attributes. You can look at
-    `QuestionAnsweringTokenizer` for an example.
-    """
-
-    default_implementation = "from_pretrained"
-    _BOS_TOKEN_KEYS = ("bos_token", "cls_token")
-    _EOS_TOKEN_KEYS = ("eos_token", "sep_token")
-    _attr_to_special_token: Dict[str, List[str]] = {}
-    _TASK_SPECIFIC_SPECIAL_TOKENS: Dict[str, str] = {}
-
-    def __init__(self):
-        raise EnvironmentError(
-            "`TransformerTokenizer` is designed to be instantiated "
-            "using the `TransformerTokenizer.from_pretrained` method."
-        )
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
-        tokenizer = AutoTokenizer.from_pretrained(
-            pretrained_model_name_or_path, *inputs, **kwargs
-        )
-        cls._post_init(tokenizer, **kwargs)
-        return tokenizer
-
-    @classmethod
-    def _post_init(cls, tokenizer: _PreTrainedTransformerTokenizer, **kwargs):
-        tokenizer.__num_added_special_tokens = cls._add_task_specific_tokens(
-            tokenizer
-        )
-        tokenizer.__model_max_sequence_length = cls._find_model_max_seq_length(
-            tokenizer, kwargs.get("model_max_sequence_length")
-        )
-        add_property(
-            tokenizer,
-            {
-                "num_added_special_tokens": lambda self: self.__num_added_special_tokens,
-                "model_max_sequence_length": lambda self: self.__model_max_sequence_length,
-                "num_tokens": len,
-            },
-        )
-
-    @classmethod
-    def _add_task_specific_tokens(
-        cls, tokenizer: _PreTrainedTransformerTokenizer
-    ) -> int:
-        _attr_to_special_token = deepcopy(cls._attr_to_special_token)
-        for tok_name, tok_value in cls._TASK_SPECIFIC_SPECIAL_TOKENS.items():
-            if getattr(tokenizer, tok_name) is None:
-                for alternative_pair in (cls._BOS_TOKEN_KEYS, cls._EOS_TOKEN_KEYS):
-                    if tok_name in alternative_pair:
-                        tok_value = cls._find_alternative_token_value(
-                            tok_name, tok_value, alternative_pair
-                        )
-                        break
-                _attr_to_special_token[tok_name] = tok_value
-
-        return tokenizer.add_special_tokens(_attr_to_special_token)
-
-    @classmethod
-    def _find_alternative_token_value(
-        cls, token_name: str, token_value: str, alternative_pair: Tuple[str, str]
-    ) -> str:
-        if token_name == alternative_pair[0]:
-            return cls._TASK_SPECIFIC_SPECIAL_TOKENS.get(
-                alternative_pair[1], token_value
-            )
-        else:
-            return cls._TASK_SPECIFIC_SPECIAL_TOKENS.get(
-                alternative_pair[0], token_value
-            )
-
-    @classmethod
-    def _find_model_max_seq_length(
-        cls,
-        tokenizer: _PreTrainedTransformerTokenizer,
-        provided_model_max_sequence_length: int = None,
-    ) -> int:
-        model_max_length = tokenizer.model_max_length
-        if provided_model_max_sequence_length is None:
-            return model_max_length
-        return min(model_max_length, provided_model_max_sequence_length)
-
-
-TransformerTokenizer.register("from_pretrained", constructor="from_pretrained")(
-    TransformerTokenizer
-)
diff --git a/trapper/data/tokenizers/tokenizer_wrapper.py b/trapper/data/tokenizers/tokenizer_wrapper.py
new file mode 100644
index 0000000..8ccca21
--- /dev/null
+++ b/trapper/data/tokenizers/tokenizer_wrapper.py
@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import os
+from copy import deepcopy
+from typing import Dict, List, Optional, Tuple, Union
+
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from trapper.common import Registrable
+from trapper.common.constants import (
+    BOS_TOKEN,
+    CLS_TOKEN,
+    EOS_TOKEN,
+    MASK_TOKEN,
+    PAD_TOKEN,
+    SEP_TOKEN,
+    UNK_TOKEN,
+)
+
+
+class TokenizerWrapper(Registrable):
+    """
+    The base tokenizer class for trapper that acts as a factory which returns a
+    `PreTrainedTokenizerBase` instance after adding the task-specific tokens to it.
+    Internally, it uses `transformers.AutoTokenizer` for creating the pretrained
+    tokenizer objects. In addition to the tokenizer, the wrapper object also holds
+    the maximum sequence length accepted by the model of that tokenizer. This class
+    also handles the differences between the special start and end of sequence
+    tokens in different models. By utilizing a tokenizer wrapped this class, you can
+    use `tokenizer.bos_token` to access the start token without thinking which model
+    you are working with. Otherwise, you would have to use `tokenizer.cls_token`
+    when you are working with BERT, whereas `tokenizer.bos_token` if you are working
+    with GPT2 for example. We fill the missing value from the (cls_token, bos_token)
+    and (eos_token, sep_token) token pairs by saving the other's value if the
+    pretrained tokenizer does not have only one of them. If neither were present,
+    they get recorded with separate values. For instance, sep_token is saved with
+    the value of eos_token in the GPT2 tokenizer since it has only eos_token
+    normally. This is done to make the BOS-CLS and EOS-SEP tokens interchangeable.
+    Finally, pad_token, mask_token and unk_token values are also set if they
+    were not already present.
+
+    You may need to override the `_TASK_SPECIFIC_SPECIAL_TOKENS` class variable to
+    specify the extra special tokens needed for your task.
+
+    Class variables that can be overridden:
+
+    _TASK_SPECIFIC_SPECIAL_TOKENS (List[str]): A list of extra special tokens that
+    is needed for the task at hand. E.g. `CONTEXT` token for SQuAD style question
+    answering tasks that utilizes a context.  You can look at
+    `QuestionAnsweringTokenizerWrapper` for that example.
+
+    Args:
+        pretrained_tokenizer (): The pretrained tokenizer to be wrapped
+    """
+
+    default_implementation = "from_pretrained"
+    _BOS_TOKEN_KEYS = ("bos_token", "cls_token")
+    _EOS_TOKEN_KEYS = ("eos_token", "sep_token")
+    _SPECIAL_TOKENS_DICT: Dict[str, str] = {
+        "bos_token": BOS_TOKEN,
+        "eos_token": EOS_TOKEN,
+        "cls_token": CLS_TOKEN,
+        "sep_token": SEP_TOKEN,
+        "pad_token": PAD_TOKEN,
+        "mask_token": MASK_TOKEN,
+        "unk_token": UNK_TOKEN,
+    }
+    _TASK_SPECIFIC_SPECIAL_TOKENS: List[str] = []
+
+    def __init__(
+        self, pretrained_tokenizer: Optional[PreTrainedTokenizerBase] = None
+    ):
+        #  We need to make `pretrained_tokenizer` optional with default of None,
+        #  since otherwise allennlp tries to invoke __init__ although we
+        #  register a classmethod as a default constructor and demand it via the
+        #  "type" parameter inside the from_params method or a config file.
+        if pretrained_tokenizer is None:
+            raise ValueError("`pretrained_tokenizer` can not be None!")
+        self._pretrained_tokenizer = pretrained_tokenizer
+        self._num_added_special_tokens = self._add_task_specific_tokens()
+
+    @property
+    def tokenizer(self) -> PreTrainedTokenizerBase:
+        return self._pretrained_tokenizer
+
+    @property
+    def num_added_special_tokens(self) -> int:
+        return self._num_added_special_tokens
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        *inputs,
+        **kwargs,
+    ) -> TokenizerWrapper:
+        pretrained_tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path, *inputs, **kwargs
+        )
+        return cls(pretrained_tokenizer)
+
+    def _add_task_specific_tokens(self) -> int:
+        tokenizer = self._pretrained_tokenizer
+        special_tokens_dict = {
+            "additional_special_tokens": deepcopy(
+                self._TASK_SPECIFIC_SPECIAL_TOKENS
+            )
+        }
+        for tok_name, tok_value in self._SPECIAL_TOKENS_DICT.items():
+            if getattr(tokenizer, tok_name) is None:
+                for alternative_pair in (
+                    self._BOS_TOKEN_KEYS,
+                    self._EOS_TOKEN_KEYS,
+                ):
+                    if tok_name in alternative_pair:
+                        tok_value = self._find_alternative_token_value(
+                            tok_name, tok_value, alternative_pair
+                        )
+                        break
+                special_tokens_dict[tok_name] = tok_value
+        num_added_special_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+        return num_added_special_tokens
+
+    def _find_alternative_token_value(
+        self,
+        original_token_name: str,
+        original_token_value: str,
+        alternative_pair: Tuple[str, str],
+    ) -> str:
+        if original_token_name == alternative_pair[0]:
+            alternative_token_name = alternative_pair[1]
+        else:
+            alternative_token_name = alternative_pair[0]
+
+        alternative_token_value = getattr(
+            self._pretrained_tokenizer, alternative_token_name
+        )
+        return alternative_token_value or original_token_value
+
+
+TokenizerWrapper.register("from_pretrained", constructor="from_pretrained")(
+    TokenizerWrapper
+)
diff --git a/trapper/models/__init__.py b/trapper/models/__init__.py
index d17f05e..548706d 100644
--- a/trapper/models/__init__.py
+++ b/trapper/models/__init__.py
@@ -1,11 +1,11 @@
 from trapper.models.auto_wrappers import (
-    TransformerModelForCausalLM,
-    TransformerModelForMaskedLM,
-    TransformerModelForMultipleChoice,
-    TransformerModelForNextSentencePrediction,
-    TransformerModelForQuestionAnswering,
-    TransformerModelForSeq2SeqLM,
-    TransformerModelForSequenceClassification,
-    TransformerModelForTokenClassification,
+    ModelWrapperForCausalLM,
+    ModelWrapperForMaskedLM,
+    ModelWrapperForMultipleChoice,
+    ModelWrapperForNextSentencePrediction,
+    ModelWrapperForQuestionAnswering,
+    ModelWrapperForSeq2SeqLM,
+    ModelWrapperForSequenceClassification,
+    ModelWrapperForTokenClassification,
 )
-from trapper.models.model import TransformerModel
+from trapper.models.model_wrapper import ModelWrapper
diff --git a/trapper/models/auto_wrappers.py b/trapper/models/auto_wrappers.py
index 75bcad2..96b0280 100644
--- a/trapper/models/auto_wrappers.py
+++ b/trapper/models/auto_wrappers.py
@@ -17,13 +17,13 @@
     AutoModelForTokenClassification,
 )
 
-from trapper.models.model import TransformerModel
+from trapper.models.model_wrapper import ModelWrapper
 
 _COMMON_INPUT_FIELDS = ["input_ids", "attention_mask"]
 _TASK_TO_INPUT_FIELDS = OrderedDict(
     [
         # Input fields are required by the data collators and task-specific
-        # TransformerModel subclasses.
+        # ModelWrapper subclasses.
         (
             "question_answering",
             (
@@ -52,49 +52,47 @@
 
 def _create_and_register_transformer_subclass(
     auto_cls: Type, task: str
-) -> Type[TransformerModel]:
+) -> Type[ModelWrapper]:
     """
-    Dynamically creates a TransformerModel subclass by wrapping an `auto` model
-    from the Transformers library. Then, the subclass is registered to
-    the framework with the `task` argument and returned.
+    Dynamically creates a ModelWrapper subclass by wrapping an `AutoModelFor...`
+    factory from the Transformers library. Then, the subclass is
+    registered to the framework with the `task` argument and returned.
     Args:
         auto_cls (Type): an `auto` class from `Transformers`
         task (str): registered name of the subclass
 
     Returns:
-        A registered task-specific `TransformerModel` subclass
+        A registered task-specific `ModelWrapper` subclass
     """
     cls = _create_transformer_subclass(auto_cls, task)
-    TransformerModel.register(task, constructor="from_pretrained")(cls)
+    ModelWrapper.register(task, constructor="from_pretrained")(cls)
     cls._TASK_SPECIFIC_AUTO_CLASS = auto_cls
     cls._TASK_SPECIFIC_FORWARD_PARAMS = _TASK_TO_INPUT_FIELDS[task]
     return cls
 
 
-def _create_transformer_subclass(
-    auto_cls: Type, task: str
-) -> Type[TransformerModel]:
+def _create_transformer_subclass(auto_cls: Type, task: str) -> Type[ModelWrapper]:
     """
-    Dynamically creates a TransformerModel subclass by wrapping an `auto` model
-    from the Transformers library.
+    Dynamically creates a ModelWrapper subclass by wrapping an `AutoModelFor...`
+    factory from the Transformers library.
     Args:
         auto_cls (Type): an `auto` class from `Transformers`
         task (str): the task name inserted to the docstring
 
     Returns:
-        A task-specific `TransformerModel` subclass
+        A task-specific `ModelWrapper` subclass
     """
     auto_cls_name = auto_cls.__name__
-    subcls_name = auto_cls_name.replace("Auto", "Transformer")
+    subcls_name = auto_cls_name.replace("AutoModel", "ModelWrapper")
     attr_dict = {"__doc__": _get_transformer_subclass_doc(auto_cls_name, task)}
-    cls: Any = type(subcls_name, (TransformerModel,), attr_dict)
+    cls: Any = type(subcls_name, (ModelWrapper,), attr_dict)
     return cls
 
 
 def _get_transformer_subclass_doc(auto_model_name: str, task: str):
     return f"""
-    Wrapper for `transformers.{auto_model_name}`. Registered as
-    the `TransformerModel` factory for `{task}` style tasks.
+    Wrapper for `transformers.{auto_model_name}`. Registered as the `ModelWrapper`
+    factory for `{task}` style tasks.
     """
 
 
@@ -105,13 +103,15 @@ def _get_transformer_subclass_doc(auto_model_name: str, task: str):
 
 # The classes that have been tested are below
 
-# The base model with a question answering head
-TransformerModelForQuestionAnswering = _create_and_register_transformer_subclass(
+# The model wrapper factory that yields a wrapped base model with a question
+# answering head
+ModelWrapperForQuestionAnswering = _create_and_register_transformer_subclass(
     AutoModelForQuestionAnswering, "question_answering"
 )
 
-# The base model with a token classification head
-TransformerModelForTokenClassification = _create_and_register_transformer_subclass(
+# The model wrapper factory that yields a wrapped base model with a token
+# classification head
+ModelWrapperForTokenClassification = _create_and_register_transformer_subclass(
     AutoModelForTokenClassification, "token_classification"
 )
 # --------------------------------------------------------------------------
@@ -119,36 +119,38 @@ def _get_transformer_subclass_doc(auto_model_name: str, task: str):
 # Experimental classes that have not been tested yet are below. Note that some of
 # them may be removed in the future.
 
-# The base model with a causal language modeling head
-TransformerModelForCausalLM = _create_and_register_transformer_subclass(
+# The model wrapper factory that yields a wrapped base model with a causal language
+# modeling head
+ModelWrapperForCausalLM = _create_and_register_transformer_subclass(
     AutoModelForCausalLM, "causal_lm"
 )
 
-# The base model with a masked language modeling head
-TransformerModelForMaskedLM = _create_and_register_transformer_subclass(
+# The model wrapper factory that yields a wrapped base model with a masked language
+# modeling head
+ModelWrapperForMaskedLM = _create_and_register_transformer_subclass(
     AutoModelForMaskedLM, "masked_lm"
 )
 
-# The base model with a seq-to-seq language modeling head
-TransformerModelForSeq2SeqLM = _create_and_register_transformer_subclass(
+# The model wrapper factory that yields a wrapped base model with a seq-to-seq
+# language modeling head
+ModelWrapperForSeq2SeqLM = _create_and_register_transformer_subclass(
     AutoModelForSeq2SeqLM, "seq2seq_lm"
 )
 
-# The base model with a sequence classification head
-TransformerModelForSequenceClassification = (
-    _create_and_register_transformer_subclass(
-        AutoModelForSequenceClassification, "sequence_classification"
-    )
+# The model wrapper factory that yields a wrapped base model with a sequence
+# classification head
+ModelWrapperForSequenceClassification = _create_and_register_transformer_subclass(
+    AutoModelForSequenceClassification, "sequence_classification"
 )
 
-# The base model with a multiple choice head
-TransformerModelForMultipleChoice = _create_and_register_transformer_subclass(
+# The model wrapper factory that yields a wrapped base model with a multiple choice
+# head
+ModelWrapperForMultipleChoice = _create_and_register_transformer_subclass(
     AutoModelForMultipleChoice, "multiple_choice"
 )
 
-# The base model with a next sentence prediction head
-TransformerModelForNextSentencePrediction = (
-    _create_and_register_transformer_subclass(
-        AutoModelForNextSentencePrediction, "next_sentence_prediction"
-    )
+# The model wrapper factory that yields a wrapped base model with a next sentence
+# prediction head
+ModelWrapperForNextSentencePrediction = _create_and_register_transformer_subclass(
+    AutoModelForNextSentencePrediction, "next_sentence_prediction"
 )
diff --git a/trapper/models/model.py b/trapper/models/model_wrapper.py
similarity index 57%
rename from trapper/models/model.py
rename to trapper/models/model_wrapper.py
index c8d0022..8e6a24b 100644
--- a/trapper/models/model.py
+++ b/trapper/models/model_wrapper.py
@@ -1,29 +1,35 @@
 # pylint: disable=protected-access
+from __future__ import annotations
+
 import inspect
 from pathlib import Path
-from typing import Dict, Optional, Tuple, Type, Union
+from typing import Dict, Optional, Tuple, Union
 
 from torch import nn
-from transformers import AutoConfig, PreTrainedModel
+from transformers import (
+    AutoConfig,
+    AutoModelForTokenClassification,
+    PreTrainedModel,
+)
+from transformers.models.auto.auto_factory import _BaseAutoModelClass
 from transformers.utils import logging
 
 from trapper.common import Registrable
-from trapper.common.utils import add_property
 
 logger = logging.get_logger(__name__)
 
 
-class TransformerModel(PreTrainedModel, Registrable):
+class ModelWrapper(Registrable):
     """
-    The `TransformerModel` is the base registrable model factory that is responsible
-    for creating transformer models from the `Transformers` library. It does so by
-    keeping a mapping from task name to the `auto_class_factory` (i.e. `auto)
-    classes defined in the `Transformers` such as
-    `AutoModelForSequenceClassification`. To make this work, we register the `auto`
-    classes with their task name to the registry as a subclass of
-    `TransformerModel`.
-
-    This class does extra things depending on the `model_type` of the created
+    The `ModelWrapper` is the registrable base factory that is responsible
+    for creating transformer models from the `transformers` library using
+    the `auto` classes . It does so by keeping a mapping from task name to the
+    `AutoModelFor...`  classes defined in `transformers` such as
+    `AutoModelForSequenceClassification`.
+    To make this work, we register the wrapped `auto` classes with their task name
+    to the registry as a subclass of `ModelWrapper`.
+
+    This class does some extra things depending on the `model_type` of the created
     task-specific class. These include:
         - inspecting the `forward` parameters of the model,
         - checking the invalid parameters i.e. the parameters that are required
@@ -33,7 +39,7 @@ class TransformerModel(PreTrainedModel, Registrable):
         but are not required by other models in general.
 
     Then, the legitimate parameters are stored as `forward_params` attribute of
-    the returned model. Thanks to that, `trapper.training.TransformerTrainer`
+    the model wrapper object. Thanks to that, `trapper.training.TransformerTrainer`
     can access that attribute to use as argument for the `model_forward_params`
     parameter while creating the dataset collator.
 
@@ -58,14 +64,13 @@ class TransformerModel(PreTrainedModel, Registrable):
         extra parameters that are peculiar to them. E.g. `longformer` utilizes
         "global_attention_mask".
 
-    Below are the constructor parameters, which are left generic.
-        *inputs ():
-        **kwargs ():
+    Args:
+        pretrained_model (): The pretrained model to be wrapped
     """
 
     default_implementation = "from_pretrained"
-    _TASK_SPECIFIC_AUTO_CLASS: Optional[Type] = None
-    _TASK_SPECIFIC_FORWARD_PARAMS: Optional[Tuple[str, ...]] = None
+    _TASK_SPECIFIC_AUTO_CLASS: _BaseAutoModelClass = None
+    _TASK_SPECIFIC_FORWARD_PARAMS: Tuple[str, ...] = None
 
     _MODEL_TYPE_TO_INVALID_PARAMS: Dict[str, Tuple[str]] = {
         "roberta": ("token_type_ids",),
@@ -78,45 +83,60 @@ class TransformerModel(PreTrainedModel, Registrable):
         "longformer": ("global_attention_mask",),
     }
 
-    def __init__(self, *inputs, **kwargs):  # pylint: disable=super-init-not-called
-        if self.__class__ == TransformerModel:
-            raise EnvironmentError(
-                "TransformerModel is designed to be a factory that can "
-                "instantiate concrete models using `TransformerModel.from_params` "
-                "method."
-            )
-        raise EnvironmentError(
-            "Task-specific `TransformerModel` subclasses are designed to be "
-            "instantiated using the `from_pretrained` method."
-        )
+    def __init__(self, pretrained_model: Optional[PreTrainedModel] = None):
+        #  We need to make `pretrained_model` optional with default of None,
+        #  since otherwise allennlp tries to invoke __init__ although we
+        #  register a classmethod as a default constructor and demand it via the
+        #  "type" parameter inside the from_params method or a config file.
+        if pretrained_model is None:
+            raise ValueError("`pretrained_model` can not be None!")
+        self._pretrained_model = pretrained_model
+        self._forward_params = self._get_forward_params(pretrained_model)
+
+    @property
+    def forward_params(self) -> Tuple[str, ...]:
+        return self._forward_params
+
+    @property
+    def model(self) -> PreTrainedModel:
+        return self._pretrained_model
 
     @classmethod
     def from_pretrained(
         cls, pretrained_model_name_or_path: Union[str, Path], *model_args, **kwargs
-    ) -> PreTrainedModel:
-        # Handles architectural changes (e.g. for further training an already
-        # fine-tuned downstream model on a new dataset) for the required
-        # task-specific `auto` models.
-        if cls == TransformerModel:
+    ) -> ModelWrapper:
+        """
+        Creates and returns a transformer model wrapper from a task-specific wrapper
+        factory. Additionally, handles the architectural changes (e.g. when further
+        training an already fine-tuned model on a new dataset with different labels)
+        for the task-specific models that require such changes.
+
+        Args:
+            pretrained_model_name_or_path ():
+            *model_args ():
+            **kwargs ():
+
+        Returns:
+
+        """
+        if cls is ModelWrapper:
             raise EnvironmentError(
-                "TransformerModel is designed to be a factory that can "
-                "instantiate concrete models using `TransformerModel.from_params` "
-                "method."
+                "ModelWrapper is designed to be a factory that can "
+                "instantiate concrete models using `ModelWrapper.from_params` "
+                "method. If you want to instantiate a ModelWrapper object using "
+                "`from_pretrained` method, you need to use a subclass of "
+                "ModelWrapper instead."
             )
-        if (
-            cls._TASK_SPECIFIC_AUTO_CLASS.__name__
-            == "AutoModelForTokenClassification"
-        ):
-            model = cls._create_token_classification_model(
+        if cls._TASK_SPECIFIC_AUTO_CLASS is AutoModelForTokenClassification:
+            pretrained_model = cls._create_token_classification_model(
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
         else:
-            model = cls._TASK_SPECIFIC_AUTO_CLASS.from_pretrained(
+            pretrained_model = cls._TASK_SPECIFIC_AUTO_CLASS.from_pretrained(
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
 
-        cls._post_init(model)
-        return model
+        return cls(pretrained_model=pretrained_model)
 
     @classmethod
     def _create_token_classification_model(
@@ -140,14 +160,14 @@ def _create_token_classification_model(
                 )
             if pretrained_num_labels != provided_num_labels:
                 logger.warning(
-                    (
-                        "Provided `num_labels` value (%s) is different from the one "
-                        "found in the archived config (%s). The classifier head will "
-                        "have %s labels and be initialized randomly!",
-                        provided_num_labels,
-                        pretrained_num_labels,
-                        provided_num_labels,
-                    )
+                    "Provided `num_labels` value (%(provided)d) is different from "
+                    "the one found in the archived config (%(pretrained)d). "
+                    "The classifier head will have %(provided)d labels and be"
+                    " initialized randomly!",
+                    {
+                        "provided": provided_num_labels,
+                        "pretrained": pretrained_num_labels,
+                    },
                 )
                 kwargs.pop("num_labels")
                 pretrained_weights = cls._TASK_SPECIFIC_AUTO_CLASS.from_pretrained(
@@ -171,27 +191,16 @@ def _create_token_classification_model(
             pretrained_model_name_or_path, *model_args, **kwargs
         )
 
-    @classmethod
-    def _post_init(cls, model: PreTrainedModel):
-        model.__forward_params = cls._get_forward_params(model)
-        add_property(
-            model,
-            {
-                "forward_params": lambda self: self.__forward_params,
-            },
-        )
-
-    @classmethod
-    def _get_forward_params(cls, model: PreTrainedModel) -> Tuple[str, ...]:
+    def _get_forward_params(self, model: PreTrainedModel) -> Tuple[str, ...]:
         forward_method_signature = inspect.signature(model.forward)
-        model_specific_invalid_params = cls._MODEL_TYPE_TO_INVALID_PARAMS.get(
+        model_specific_invalid_params = self._MODEL_TYPE_TO_INVALID_PARAMS.get(
             model.config.model_type, ()
         )
-        model_specific_extra_params = cls._MODEL_TYPE_TO_EXTRA_PARAMS.get(
+        model_specific_extra_params = self._MODEL_TYPE_TO_EXTRA_PARAMS.get(
             model.config.model_type, ()
         )
         task_and_model_params = {
-            *cls._TASK_SPECIFIC_FORWARD_PARAMS,
+            *self._TASK_SPECIFIC_FORWARD_PARAMS,
             *model_specific_extra_params,
         }
         return tuple(
@@ -202,6 +211,6 @@ def _get_forward_params(cls, model: PreTrainedModel) -> Tuple[str, ...]:
         )
 
 
-TransformerModel.register("from_pretrained", constructor="from_pretrained")(
-    TransformerModel
+ModelWrapper.register("from_pretrained", constructor="from_pretrained")(
+    ModelWrapper
 )
diff --git a/trapper/pipelines/pipeline.py b/trapper/pipelines/pipeline.py
index 604b308..a2ab640 100644
--- a/trapper/pipelines/pipeline.py
+++ b/trapper/pipelines/pipeline.py
@@ -1,13 +1,14 @@
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 from transformers import AutoConfig, Pipeline
 from transformers.pipelines import pipeline
 
 from trapper.common.params import Params
-from trapper.data import DataAdapter, DataProcessor, TransformerTokenizer
+from trapper.data import DataAdapter, DataProcessor, TokenizerWrapper
 from trapper.data.data_collator import DataCollator
-from trapper.models import TransformerModel
+from trapper.data.label_mapper import LabelMapper
+from trapper.models import ModelWrapper
 
 
 def create_pipeline_from_checkpoint(
@@ -22,17 +23,18 @@ def create_pipeline_from_checkpoint(
     return pipeline_
 
 
-def _create_pipeline(checkpoint_path, params, task: str, **kwargs):
-    model = _create_model(checkpoint_path, params)
-    tokenizer = _create_tokenizer(checkpoint_path, params)
-    data_processor = _create_data_processor(params, tokenizer)
-    data_adapter = _create_data_adapter(params, tokenizer)
-    data_collator = _create_data_collator(model, tokenizer)
+def _create_pipeline(checkpoint_path, params: Params, task: str, **kwargs):
+    model_wrapper = _create_model_wrapper(checkpoint_path, params)
+    tokenizer_wrapper = _create_tokenizer(checkpoint_path, params)
+    label_mapper = _create_label_mapper(params)
+    data_processor = _create_data_processor(params, tokenizer_wrapper, label_mapper)
+    data_adapter = _create_data_adapter(params, tokenizer_wrapper, label_mapper)
+    data_collator = _create_data_collator(model_wrapper, tokenizer_wrapper)
     config = AutoConfig.from_pretrained(checkpoint_path)
     pipeline_ = pipeline(
         task=task,
-        model=model,
-        tokenizer=tokenizer,
+        model=model_wrapper.model,
+        tokenizer=tokenizer_wrapper.tokenizer,
         config=config,
         framework="pt",
         data_processor=data_processor,
@@ -43,31 +45,34 @@ def _create_pipeline(checkpoint_path, params, task: str, **kwargs):
     return pipeline_
 
 
-def _create_data_collator(model, tokenizer):
+def _create_data_collator(
+    model_wrapper: ModelWrapper, tokenizer_wrapper: TokenizerWrapper
+) -> DataCollator:
     return DataCollator(
-        tokenizer=tokenizer, model_forward_params=model.forward_params
+        tokenizer_wrapper=tokenizer_wrapper,
+        model_forward_params=model_wrapper.forward_params,
     )
 
 
-def _validate_checkpoint_dir(path: Union[str, Path]):
+def _validate_checkpoint_dir(path: Union[str, Path]) -> None:
     path = Path(path)
     if not path.is_dir():
         raise ValueError("Input path must be an existing directory")
 
 
-def _create_model(checkpoint_path, params):
-    return TransformerModel.from_params(
+def _create_model_wrapper(checkpoint_path, params: Params) -> ModelWrapper:
+    return ModelWrapper.from_params(
         Params(
             {
-                "type": params["model"]["type"],
+                "type": params["model_wrapper"]["type"],
                 "pretrained_model_name_or_path": checkpoint_path,
             }
         )
     )
 
 
-def _create_tokenizer(checkpoint_path, params):
-    return TransformerTokenizer.from_params(
+def _create_tokenizer(checkpoint_path, params: Params) -> TokenizerWrapper:
+    return TokenizerWrapper.from_params(
         Params(
             {
                 "type": params["tokenizer"]["type"],
@@ -77,13 +82,36 @@ def _create_tokenizer(checkpoint_path, params):
     )
 
 
-def _create_data_processor(params, tokenizer):
-    return DataProcessor.by_name(
-        params["dataset_loader"]["data_processor"]["type"]
-    )(tokenizer)
+def _create_label_mapper(params: Params) -> Optional[LabelMapper]:
+    label_mapper_params = params.get("label_mapper")
+    if label_mapper_params is None:
+        return None
+    constructor = LabelMapper.by_name(label_mapper_params["type"])
+    del label_mapper_params["type"]
+    return constructor(label_mapper_params)
 
 
-def _create_data_adapter(params, tokenizer):
-    return DataAdapter.by_name(params["dataset_loader"]["data_adapter"]["type"])(
-        tokenizer
+def _create_data_processor(
+    params: Params, tokenizer_wrapper: TokenizerWrapper, label_mapper: LabelMapper
+) -> DataProcessor:
+    data_processor_params = params["dataset_loader"]["data_processor"]
+    constructor = DataProcessor.by_name(data_processor_params["type"])
+    model_max_sequence_length = data_processor_params.get(
+        "model_max_sequence_length"
+    )
+    return constructor(
+        tokenizer_wrapper=tokenizer_wrapper,
+        model_max_sequence_length=model_max_sequence_length,
+        label_mapper=label_mapper,
+    )
+
+
+def _create_data_adapter(
+    params: Params, tokenizer_wrapper: TokenizerWrapper, label_mapper: LabelMapper
+) -> DataAdapter:
+    constructor = DataAdapter.by_name(
+        params["dataset_loader"]["data_adapter"]["type"]
+    )
+    return constructor(
+        tokenizer_wrapper=tokenizer_wrapper, label_mapper=label_mapper
     )
diff --git a/trapper/pipelines/question_answering_pipeline.py b/trapper/pipelines/question_answering_pipeline.py
index b8c7a68..509d39e 100644
--- a/trapper/pipelines/question_answering_pipeline.py
+++ b/trapper/pipelines/question_answering_pipeline.py
@@ -46,7 +46,7 @@
     DataAdapterForQuestionAnswering,
 )
 from trapper.data.data_collator import DataCollator
-from trapper.models import TransformerModel
+from trapper.models import ModelWrapper
 
 
 class QuestionAnsweringArgumentHandler(ArgumentHandler):
@@ -431,5 +431,5 @@ def postprocess_answer(raw_answer: List[Dict]) -> SpanDict:
 
 SUPPORTED_TASKS["squad-question-answering"] = {
     "impl": SquadQuestionAnsweringPipeline,
-    "pt": TransformerModel,
+    "pt": ModelWrapper,
 }
diff --git a/trapper/training/metrics.py b/trapper/training/metrics.py
index 09e6d73..8f44be2 100644
--- a/trapper/training/metrics.py
+++ b/trapper/training/metrics.py
@@ -1,12 +1,13 @@
 from abc import ABCMeta, abstractmethod
-from typing import Dict, List
+from typing import Dict, Optional
 
 import numpy as np
 from datasets import load_metric
 from transformers import EvalPrediction
 
 from trapper.common import Registrable
-from trapper.common.constants import PAD_TOKEN_LABEL_ID
+from trapper.common.constants import IGNORED_LABEL_ID
+from trapper.data.label_mapper import LabelMapper
 
 
 class TransformerMetric(Registrable, metaclass=ABCMeta):
@@ -15,8 +16,15 @@ class TransformerMetric(Registrable, metaclass=ABCMeta):
     evaluating the models. The subclasses should be implemented as callables
     that accepts a `transformers.EvalPrediction` in their `__call__` method and
     compute score for that prediction.
+
+    Args:
+        label_mapper (): Only used in some tasks that require mapping between
+            categorical labels and integer ids such as token classification.
     """
 
+    def __init__(self, label_mapper: Optional[LabelMapper] = None):
+        self._label_mapper = label_mapper
+
     @abstractmethod
     def __call__(self, prediction: EvalPrediction) -> Dict[str, float]:
         raise NotImplementedError
@@ -28,19 +36,30 @@ class SeqEvalMetric(TransformerMetric):
     Creates a token classification task metric that returns of precision,
     recall, f1 and accuracy scores. Internally, uses the "seqeval" metric
     from the HuggingFace's `datasets` library.
+
     Args:
+        label_mapper (): Used for mapping between categorical labels and integer
+            ids.
         return_entity_level_metrics (bool): Set True to return all the
             entity levels during evaluation. Otherwise, returns overall
             results.
     """
 
     def __init__(
-        self, label_list: List[str], return_entity_level_metrics: bool = True
+        self, label_mapper: LabelMapper, return_entity_level_metrics: bool = True
     ):
+        if label_mapper is None:
+            raise ValueError(
+                f"`SeqEvalMetric` can not be instantiated without a `LabelMapper`"
+            )
+
+        super().__init__(label_mapper=label_mapper)
         self._metric = load_metric("seqeval")
-        self._label_list = label_list
         self._return_entity_level_metrics = return_entity_level_metrics
 
+    def _id_to_label(self, id_: int) -> str:
+        return self._label_mapper.get_label(id_)
+
     def __call__(self, pred: EvalPrediction) -> Dict[str, float]:
         all_predicted_ids = np.argmax(pred.predictions, axis=2)
         all_label_ids = pred.label_ids
@@ -50,9 +69,9 @@ def __call__(self, pred: EvalPrediction) -> Dict[str, float]:
             actual_prediction = []
             actual_label = []
             for (p, l) in zip(predicted_ids, label_ids):
-                if l != PAD_TOKEN_LABEL_ID:
-                    actual_prediction.append(self._label_list[p])
-                    actual_label.append(self._label_list[l])
+                if l != IGNORED_LABEL_ID:
+                    actual_prediction.append(self._id_to_label(p))
+                    actual_label.append(self._id_to_label(l))
 
             actual_predictions.append(actual_prediction)
             actual_labels.append(actual_label)
diff --git a/trapper/training/train.py b/trapper/training/train.py
index ea5c47a..5feb939 100644
--- a/trapper/training/train.py
+++ b/trapper/training/train.py
@@ -175,8 +175,8 @@ def main():
     """
     if not len(sys.argv) == 2:
         raise ValueError(
-            "`train` scripts accept a single argument i.e. the "
-            "config file in json or jsonnet format!"
+            "`train` script accepts a single argument i.e. the config file in json"
+            " or jsonnet format!"
         )
     config_path = sys.argv[1]
     return run_experiment(config_path)
diff --git a/trapper/training/trainer.py b/trapper/training/trainer.py
index 2f77b38..ab0ce4c 100644
--- a/trapper/training/trainer.py
+++ b/trapper/training/trainer.py
@@ -4,16 +4,17 @@
 import datasets
 import torch
 from torch.optim.lr_scheduler import LambdaLR
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
 from transformers.trainer import Trainer as _Trainer
 from transformers.trainer_utils import EvalPrediction
 
 from trapper.common import Lazy, Registrable
 from trapper.common.plugins import import_plugins
 from trapper.common.utils import append_parent_docstr
-from trapper.data import DataAdapter, DatasetLoader, TransformerTokenizer
+from trapper.data import DatasetLoader, TokenizerWrapper
 from trapper.data.data_collator import DataCollator
-from trapper.models import TransformerModel
+from trapper.data.label_mapper import LabelMapper
+from trapper.models import ModelWrapper
 from trapper.training.callbacks import TrainerCallback
 from trapper.training.metrics import TransformerMetric
 from trapper.training.optimizers import Optimizer
@@ -36,8 +37,8 @@ def __init__(
         data_collator: Optional[DataCollator] = None,
         train_dataset: Optional[datasets.Dataset] = None,
         eval_dataset: Optional[datasets.Dataset] = None,
-        tokenizer: Optional[TransformerTokenizer] = None,
-        model_init: Callable[[], TransformerModel] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Callable[[], PreTrainedModel] = None,
         compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
         callbacks: Optional[List[TrainerCallback]] = None,
         optimizers: Tuple[torch.optim.Optimizer, Optional[LambdaLR]] = (None, None),
@@ -61,11 +62,12 @@ def from_partial_objects(
         pretrained_model_name_or_path: str,
         train_split_name: str,
         dev_split_name: str,
-        model: Lazy[TransformerModel],
-        tokenizer: Lazy[TransformerTokenizer],
+        model_wrapper: Lazy[ModelWrapper],
+        tokenizer_wrapper: Lazy[TokenizerWrapper],
         dataset_loader: Lazy[DatasetLoader],
         data_collator: Lazy[DataCollator],
         optimizer: Lazy[Optimizer],
+        label_mapper: Optional[LabelMapper] = None,
         compute_metrics: Optional[Lazy[TransformerMetric]] = None,
         no_grad: List[str] = None,
         args: TransformerTrainingArguments = None,
@@ -75,74 +77,94 @@ def from_partial_objects(
         #  To find the registrable components from the user-defined packages
         import_plugins()
 
-        tokenizer_ = tokenizer.construct(
+        model_wrapper_ = model_wrapper.construct(
             pretrained_model_name_or_path=pretrained_model_name_or_path
         )
-        model_ = model.construct(
+        model = model_wrapper_.model
+        model_forward_params = model_wrapper_.forward_params
+
+        tokenizer_wrapper_ = tokenizer_wrapper.construct(
             pretrained_model_name_or_path=pretrained_model_name_or_path
         )
-        cls._resize_token_embeddings(model=model_, tokenizer=tokenizer_)
-        cls.mark_params_with_no_grads(model_, no_grad)
-        params_with_grad = [
-            [n, p] for n, p in model_.named_parameters() if p.requires_grad
-        ]
-        optimizer_ = optimizer.construct(model_parameters=params_with_grad)
+
+        cls._resize_token_embeddings(
+            model=model, tokenizer_wrapper=tokenizer_wrapper_
+        )
+
+        optimizer_ = cls._create_optimizer(model, optimizer, no_grad)
+
         dataset_loader_ = dataset_loader.construct(
-            tokenizer=tokenizer_, model_forward_params=model_.forward_params
+            tokenizer_wrapper=tokenizer_wrapper_,
+            model_forward_params=model_forward_params,
         )
         train_dataset_ = dataset_loader_.load(train_split_name)
         eval_dataset_ = dataset_loader_.load(dev_split_name)
+
         data_collator_ = data_collator.construct(
-            tokenizer=tokenizer_, model_forward_params=model_.forward_params
+            tokenizer_wrapper=tokenizer_wrapper_,
+            model_forward_params=model_forward_params,
         )
+
         compute_metrics_ = cls._create_compute_metrics(
-            compute_metrics, dataset_loader_.data_adapter
+            compute_metrics, label_mapper
         )
         return cls(
-            model=model_,
+            model=model_wrapper_.model,
             args=args,
             data_collator=data_collator_,
             train_dataset=train_dataset_,
             eval_dataset=eval_dataset_,
-            tokenizer=tokenizer_,
+            tokenizer=tokenizer_wrapper_.tokenizer,
             compute_metrics=compute_metrics_,
             callbacks=callbacks,
             optimizers=(optimizer_, None),
         )
 
+    @classmethod
+    def _create_optimizer(
+        cls,
+        model: PreTrainedModel,
+        optimizer: Lazy[Optimizer],
+        no_grad: List[str] = None,
+    ) -> Optimizer:
+        cls.mark_params_with_no_grads(model, no_grad)
+        params_with_grad = [
+            [n, p] for n, p in model.named_parameters() if p.requires_grad
+        ]
+        optimizer = optimizer.construct(model_parameters=params_with_grad)
+        return optimizer
+
     @classmethod
     def _create_compute_metrics(
         cls,
         compute_metrics: Optional[Lazy[TransformerMetric]],
-        data_adapter: DataAdapter,
+        label_mapper: Optional[LabelMapper] = None,
     ) -> Optional[TransformerMetric]:
         if compute_metrics is None:
             return None
-        label_list = getattr(data_adapter, "label_list", None)
-        if label_list is None:
-            raise ValueError(
-                f"The data adapter {str(type(data_adapter))}"
-                + " must implement the label_list attribute."
-            )
-        return compute_metrics.construct(label_list=data_adapter.label_list)
+        return compute_metrics.construct(label_mapper=label_mapper)
 
     @classmethod
-    def mark_params_with_no_grads(cls, model_, no_grad):
+    def mark_params_with_no_grads(
+        cls, model: PreTrainedModel, no_grad: List[str] = None
+    ):
         if no_grad:
-            for name, parameter in model_.named_parameters():
+            for name, parameter in model.named_parameters():
                 if any(re.search(regex, name) for regex in no_grad):
                     parameter.requires_grad_(False)
 
     @classmethod
     def _resize_token_embeddings(
-        cls, model: PreTrainedModel, tokenizer: TransformerTokenizer
+        cls, model: PreTrainedModel, tokenizer_wrapper: TokenizerWrapper
     ):
         """
         Update the token embedding layer of the model to accommodate
         for the special tokens added to the tokenizer
         """
-        if tokenizer.num_added_special_tokens > 0:
-            model.resize_token_embeddings(new_num_tokens=tokenizer.num_tokens)
+        if tokenizer_wrapper.num_added_special_tokens > 0:
+            model.resize_token_embeddings(
+                new_num_tokens=len(tokenizer_wrapper.tokenizer)
+            )
 
 
 TransformerTrainer.register("default", constructor="from_partial_objects")(
diff --git a/trapper/version.py b/trapper/version.py
index 1e599f0..e7fdd5e 100644
--- a/trapper/version.py
+++ b/trapper/version.py
@@ -1,5 +1,5 @@
 _MAJOR = "0"
 _MINOR = "0"
-_PATCH = "1"
+_PATCH = "2"
 
 VERSION = "{0}.{1}.{2}".format(_MAJOR, _MINOR, _PATCH)


Test patch:
diff --git a/examples/pos_tagging/scripts/run_tests.py b/examples/pos_tagging/scripts/run_tests.py
new file mode 100644
index 0000000..256900c
--- /dev/null
+++ b/examples/pos_tagging/scripts/run_tests.py
@@ -0,0 +1,7 @@
+from trapper.common.testing_utils.shell_utils import shell, validate_and_exit
+
+if __name__ == "__main__":
+    sts_tests = shell(
+        "pytest --cov trapper --cov-report term-missing --cov-report xml -vvv tests"
+    )
+    validate_and_exit(tests=sts_tests)
diff --git a/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/README.md b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/README.md
new file mode 100644
index 0000000..6382401
--- /dev/null
+++ b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/README.md
@@ -0,0 +1,198 @@
+---
+annotations_creators:
+- crowdsourced
+language_creators:
+- found
+languages:
+- en
+licenses:
+- unknown
+multilinguality:
+- monolingual
+size_categories:
+- 10K<n<100K
+source_datasets:
+- extended|other-reuters-corpus
+task_categories:
+- structure-prediction
+task_ids:
+- named-entity-recognition
+- part-of-speech-tagging
+paperswithcode_id: conll-2003
+---
+
+# Dataset Card for "conll2003"
+
+## Table of Contents
+- [Dataset Description](#dataset-description)
+  - [Dataset Summary](#dataset-summary)
+  - [Supported Tasks and Leaderboards](#supported-tasks-and-leaderboards)
+  - [Languages](#languages)
+- [Dataset Structure](#dataset-structure)
+  - [Data Instances](#data-instances)
+  - [Data Fields](#data-fields)
+  - [Data Splits](#data-splits)
+- [Dataset Creation](#dataset-creation)
+  - [Curation Rationale](#curation-rationale)
+  - [Source Data](#source-data)
+  - [Annotations](#annotations)
+  - [Personal and Sensitive Information](#personal-and-sensitive-information)
+- [Considerations for Using the Data](#considerations-for-using-the-data)
+  - [Social Impact of Dataset](#social-impact-of-dataset)
+  - [Discussion of Biases](#discussion-of-biases)
+  - [Other Known Limitations](#other-known-limitations)
+- [Additional Information](#additional-information)
+  - [Dataset Curators](#dataset-curators)
+  - [Licensing Information](#licensing-information)
+  - [Citation Information](#citation-information)
+  - [Contributions](#contributions)
+
+## Dataset Description
+
+- **Homepage:** [https://www.aclweb.org/anthology/W03-0419/](https://www.aclweb.org/anthology/W03-0419/)
+- **Repository:** [More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+- **Paper:** [More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+- **Point of Contact:** [More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+- **Size of downloaded dataset files:** 4.63 MB
+- **Size of the generated dataset:** 9.78 MB
+- **Total amount of disk used:** 14.41 MB
+
+### Dataset Summary
+
+The shared task of CoNLL-2003 concerns language-independent named entity recognition. We will concentrate on
+four types of named entities: persons, locations, organizations and names of miscellaneous entities that do
+not belong to the previous three groups.
+
+The CoNLL-2003 shared task data files contain four columns separated by a single space. Each word has been put on
+a separate line and there is an empty line after each sentence. The first item on each line is a word, the second
+a part-of-speech (POS) tag, the third a syntactic chunk tag and the fourth the named entity tag. The chunk tags
+and the named entity tags have the format I-TYPE which means that the word is inside a phrase of type TYPE. Only
+if two phrases of the same type immediately follow each other, the first word of the second phrase will have tag
+B-TYPE to show that it starts a new phrase. A word with tag O is not part of a phrase. Note the dataset uses IOB2
+tagging scheme, whereas the original dataset uses IOB1.
+
+For more details see https://www.clips.uantwerpen.be/conll2003/ner/ and https://www.aclweb.org/anthology/W03-0419
+
+### Supported Tasks and Leaderboards
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Languages
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+## Dataset Structure
+
+### Data Instances
+
+#### conll2003
+
+- **Size of downloaded dataset files:** 4.63 MB
+- **Size of the generated dataset:** 9.78 MB
+- **Total amount of disk used:** 14.41 MB
+
+An example of 'train' looks as follows.
+
+```
+{
+    "chunk_tags": [11, 12, 12, 21, 13, 11, 11, 21, 13, 11, 12, 13, 11, 21, 22, 11, 12, 17, 11, 21, 17, 11, 12, 12, 21, 22, 22, 13, 11, 0],
+    "id": "0",
+    "ner_tags": [0, 3, 4, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    "pos_tags": [12, 22, 22, 38, 15, 22, 28, 38, 15, 16, 21, 35, 24, 35, 37, 16, 21, 15, 24, 41, 15, 16, 21, 21, 20, 37, 40, 35, 21, 7],
+    "tokens": ["The", "European", "Commission", "said", "on", "Thursday", "it", "disagreed", "with", "German", "advice", "to", "consumers", "to", "shun", "British", "lamb", "until", "scientists", "determine", "whether", "mad", "cow", "disease", "can", "be", "transmitted", "to", "sheep", "."]
+}
+```
+
+The original data files have `-DOCSTART-` lines used to separate documents, but these lines are removed here.
+Indeed `-DOCSTART-` is a special line that acts as a boundary between two different documents, and it is filtered out in this implementation.
+
+### Data Fields
+
+The data fields are the same among all splits.
+
+#### conll2003
+- `id`: a `string` feature.
+- `tokens`: a `list` of `string` features.
+- `pos_tags`: a `list` of classification labels, with possible values including `"` (0), `''` (1), `#` (2), `$` (3), `(` (4).
+- `chunk_tags`: a `list` of classification labels, with possible values including `O` (0), `B-ADJP` (1), `I-ADJP` (2), `B-ADVP` (3), `I-ADVP` (4).
+- `ner_tags`: a `list` of classification labels, with possible values including `O` (0), `B-PER` (1), `I-PER` (2), `B-ORG` (3), `I-ORG` (4) `B-LOC` (5), `I-LOC` (6) `B-MISC` (7), `I-MISC` (8).
+
+### Data Splits
+
+|  name   |train|validation|test|
+|---------|----:|---------:|---:|
+|conll2003|14041|      3250|3453|
+
+## Dataset Creation
+
+### Curation Rationale
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Source Data
+
+#### Initial Data Collection and Normalization
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+#### Who are the source language producers?
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Annotations
+
+#### Annotation process
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+#### Who are the annotators?
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Personal and Sensitive Information
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+## Considerations for Using the Data
+
+### Social Impact of Dataset
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Discussion of Biases
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Other Known Limitations
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+## Additional Information
+
+### Dataset Curators
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Licensing Information
+
+[More Information Needed](https://github.com/huggingface/datasets/blob/master/CONTRIBUTING.md#how-to-contribute-to-the-dataset-cards)
+
+### Citation Information
+
+```
+@inproceedings{tjong-kim-sang-de-meulder-2003-introduction,
+    title = "Introduction to the {C}o{NLL}-2003 Shared Task: Language-Independent Named Entity Recognition",
+    author = "Tjong Kim Sang, Erik F.  and
+      De Meulder, Fien",
+    booktitle = "Proceedings of the Seventh Conference on Natural Language Learning at {HLT}-{NAACL} 2003",
+    year = "2003",
+    url = "https://www.aclweb.org/anthology/W03-0419",
+    pages = "142--147",
+}
+
+```
+
+
+### Contributions
+
+Thanks to [@jplu](https://github.com/jplu), [@vblagoje](https://github.com/vblagoje), [@lhoestq](https://github.com/lhoestq) for adding this dataset.
\ No newline at end of file
diff --git a/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/conll2003_test_fixture.py b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/conll2003_test_fixture.py
new file mode 100644
index 0000000..61ec312
--- /dev/null
+++ b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/conll2003_test_fixture.py
@@ -0,0 +1,242 @@
+# coding=utf-8
+# Copyright 2020 HuggingFace Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""Introduction to the CoNLL-2003 Shared Task: Language-Independent Named Entity Recognition"""
+
+import datasets
+
+logger = datasets.logging.get_logger(__name__)
+
+_CITATION = """\
+@inproceedings{tjong-kim-sang-de-meulder-2003-introduction,
+    title = "Introduction to the {C}o{NLL}-2003 Shared Task: Language-Independent Named Entity Recognition",
+    author = "Tjong Kim Sang, Erik F.  and
+      De Meulder, Fien",
+    booktitle = "Proceedings of the Seventh Conference on Natural Language Learning at {HLT}-{NAACL} 2003",
+    year = "2003",
+    url = "https://www.aclweb.org/anthology/W03-0419",
+    pages = "142--147",
+}
+"""
+
+_DESCRIPTION = """\
+The shared task of CoNLL-2003 concerns language-independent named entity recognition. We will concentrate on
+four types of named entities: persons, locations, organizations and names of miscellaneous entities that do
+not belong to the previous three groups.
+
+The CoNLL-2003 shared task data files contain four columns separated by a single space. Each word has been put on
+a separate line and there is an empty line after each sentence. The first item on each line is a word, the second
+a part-of-speech (POS) tag, the third a syntactic chunk tag and the fourth the named entity tag. The chunk tags
+and the named entity tags have the format I-TYPE which means that the word is inside a phrase of type TYPE. Only
+if two phrases of the same type immediately follow each other, the first word of the second phrase will have tag
+B-TYPE to show that it starts a new phrase. A word with tag O is not part of a phrase. Note the dataset uses IOB2
+tagging scheme, whereas the original dataset uses IOB1.
+
+For more details see https://www.clips.uantwerpen.be/conll2003/ner/ and https://www.aclweb.org/anthology/W03-0419
+"""
+
+_TRAINING_FILE = "train.txt"
+_DEV_FILE = "valid.txt"
+_TEST_FILE = "test.txt"
+
+
+class Conll2003Config(datasets.BuilderConfig):
+    """BuilderConfig for Conll2003"""
+
+    def __init__(self, **kwargs):
+        """BuilderConfig forConll2003.
+
+        Args:
+          **kwargs: keyword arguments forwarded to super.
+        """
+        super(Conll2003Config, self).__init__(**kwargs)
+
+
+class Conll2003(datasets.GeneratorBasedBuilder):
+    """Conll2003 dataset."""
+
+    BUILDER_CONFIGS = [
+        Conll2003Config(name="conll2003", version=datasets.Version("1.0.0"),
+                        description="Conll2003 dataset"),
+    ]
+
+    def _info(self):
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=datasets.Features(
+                {
+                    "id": datasets.Value("string"),
+                    "tokens": datasets.Sequence(datasets.Value("string")),
+                    "pos_tags": datasets.Sequence(
+                        datasets.features.ClassLabel(
+                            names=[
+                                '"',
+                                "''",
+                                "#",
+                                "$",
+                                "(",
+                                ")",
+                                ",",
+                                ".",
+                                ":",
+                                "``",
+                                "CC",
+                                "CD",
+                                "DT",
+                                "EX",
+                                "FW",
+                                "IN",
+                                "JJ",
+                                "JJR",
+                                "JJS",
+                                "LS",
+                                "MD",
+                                "NN",
+                                "NNP",
+                                "NNPS",
+                                "NNS",
+                                "NN|SYM",
+                                "PDT",
+                                "POS",
+                                "PRP",
+                                "PRP$",
+                                "RB",
+                                "RBR",
+                                "RBS",
+                                "RP",
+                                "SYM",
+                                "TO",
+                                "UH",
+                                "VB",
+                                "VBD",
+                                "VBG",
+                                "VBN",
+                                "VBP",
+                                "VBZ",
+                                "WDT",
+                                "WP",
+                                "WP$",
+                                "WRB",
+                            ]
+                        )
+                    ),
+                    "chunk_tags": datasets.Sequence(
+                        datasets.features.ClassLabel(
+                            names=[
+                                "O",
+                                "B-ADJP",
+                                "I-ADJP",
+                                "B-ADVP",
+                                "I-ADVP",
+                                "B-CONJP",
+                                "I-CONJP",
+                                "B-INTJ",
+                                "I-INTJ",
+                                "B-LST",
+                                "I-LST",
+                                "B-NP",
+                                "I-NP",
+                                "B-PP",
+                                "I-PP",
+                                "B-PRT",
+                                "I-PRT",
+                                "B-SBAR",
+                                "I-SBAR",
+                                "B-UCP",
+                                "I-UCP",
+                                "B-VP",
+                                "I-VP",
+                            ]
+                        )
+                    ),
+                    "ner_tags": datasets.Sequence(
+                        datasets.features.ClassLabel(
+                            names=[
+                                "O",
+                                "B-PER",
+                                "I-PER",
+                                "B-ORG",
+                                "I-ORG",
+                                "B-LOC",
+                                "I-LOC",
+                                "B-MISC",
+                                "I-MISC",
+                            ]
+                        )
+                    ),
+                }
+            ),
+            supervised_keys=None,
+            homepage="https://www.aclweb.org/anthology/W03-0419/",
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        """Returns SplitGenerators."""
+        urls_to_download = {
+            "train": f"{_TRAINING_FILE}",
+            "dev": f"{_DEV_FILE}",
+            "test": f"{_TEST_FILE}",
+        }
+        downloaded_files = dl_manager.download_and_extract(urls_to_download)
+
+        return [
+            datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={
+                "filepath": downloaded_files["train"]}),
+            datasets.SplitGenerator(name=datasets.Split.VALIDATION, gen_kwargs={
+                "filepath": downloaded_files["dev"]}),
+            datasets.SplitGenerator(name=datasets.Split.TEST, gen_kwargs={
+                "filepath": downloaded_files["test"]}),
+        ]
+
+    def _generate_examples(self, filepath):
+        logger.info("⏳ Generating examples from = %s", filepath)
+        with open(filepath, encoding="utf-8") as f:
+            guid = 0
+            tokens = []
+            pos_tags = []
+            chunk_tags = []
+            ner_tags = []
+            for line in f:
+                if line.startswith("-DOCSTART-") or line == "" or line == "\n":
+                    if tokens:
+                        yield guid, {
+                            "id": str(guid),
+                            "tokens": tokens,
+                            "pos_tags": pos_tags,
+                            "chunk_tags": chunk_tags,
+                            "ner_tags": ner_tags,
+                        }
+                        guid += 1
+                        tokens = []
+                        pos_tags = []
+                        chunk_tags = []
+                        ner_tags = []
+                else:
+                    # conll2003 tokens are space separated
+                    splits = line.split(" ")
+                    tokens.append(splits[0])
+                    pos_tags.append(splits[1])
+                    chunk_tags.append(splits[2])
+                    ner_tags.append(splits[3].rstrip())
+            # last example
+            yield guid, {
+                "id": str(guid),
+                "tokens": tokens,
+                "pos_tags": pos_tags,
+                "chunk_tags": chunk_tags,
+                "ner_tags": ner_tags,
+            }
diff --git a/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/test.txt b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/test.txt
new file mode 100644
index 0000000..7487d1e
--- /dev/null
+++ b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/test.txt
@@ -0,0 +1,25 @@
+Japan NNP B-NP B-LOC
+began VBD B-VP O
+the DT B-NP O
+defence NN I-NP O
+of IN B-PP O
+their PRP$ B-NP O
+Asian JJ I-NP B-MISC
+Cup NNP I-NP I-MISC
+title NN I-NP O
+with IN B-PP O
+a DT B-NP O
+lucky JJ I-NP O
+2-1 CD I-NP O
+win VBP B-VP O
+against IN B-PP O
+Syria NNP B-NP B-LOC
+in IN B-PP O
+a DT B-NP O
+Group NNP I-NP O
+C NNP I-NP O
+championship NN I-NP O
+match NN I-NP O
+on IN B-PP O
+Friday NNP B-NP O
+. . O O
diff --git a/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/train.txt b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/train.txt
new file mode 100644
index 0000000..de77c96
--- /dev/null
+++ b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/train.txt
@@ -0,0 +1,30 @@
+The DT B-NP O
+European NNP I-NP B-ORG
+Commission NNP I-NP I-ORG
+said VBD B-VP O
+on IN B-PP O
+Thursday NNP B-NP O
+it PRP B-NP O
+disagreed VBD B-VP O
+with IN B-PP O
+German JJ B-NP B-MISC
+advice NN I-NP O
+to TO B-PP O
+consumers NNS B-NP O
+to TO B-VP O
+shun VB I-VP O
+British JJ B-NP B-MISC
+lamb NN I-NP O
+until IN B-SBAR O
+scientists NNS B-NP O
+determine VBP B-VP O
+whether IN B-SBAR O
+mad JJ B-NP O
+cow NN I-NP O
+disease NN I-NP O
+can MD B-VP O
+be VB I-VP O
+transmitted VBN I-VP O
+to TO B-PP O
+sheep NN B-NP O
+. . O O
diff --git a/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/valid.txt b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/valid.txt
new file mode 100644
index 0000000..a43d6a7
--- /dev/null
+++ b/examples/pos_tagging/test_fixtures/hf_datasets/conll2003_test_fixture/valid.txt
@@ -0,0 +1,35 @@
+West NNP B-NP B-MISC
+Indian NNP I-NP I-MISC
+all-rounder NN I-NP O
+Phil NNP I-NP B-PER
+Simmons NNP I-NP I-PER
+took VBD B-VP O
+four CD B-NP O
+for IN B-PP O
+38 CD B-NP O
+on IN B-PP O
+Friday NNP B-NP O
+as IN B-PP O
+Leicestershire NNP B-NP B-ORG
+beat VBD B-VP O
+Somerset NNP B-NP B-ORG
+by IN B-PP O
+an DT B-NP O
+innings NN I-NP O
+and CC O O
+39 CD B-NP O
+runs NNS I-NP O
+in IN B-PP O
+two CD B-NP O
+days NNS I-NP O
+to TO B-VP O
+take VB I-VP O
+over IN B-PP O
+at IN B-PP O
+the DT B-NP O
+head NN I-NP O
+of IN B-PP O
+the DT B-NP O
+county NN I-NP O
+championship NN I-NP O
+. . O O
diff --git a/examples/pos_tagging/tests/__init__.py b/examples/pos_tagging/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/pos_tagging/tests/conftest.py b/examples/pos_tagging/tests/conftest.py
new file mode 100644
index 0000000..63a5c0d
--- /dev/null
+++ b/examples/pos_tagging/tests/conftest.py
@@ -0,0 +1,25 @@
+import pytest
+
+from examples.pos_tagging.src import POS_TAGGING_FIXTURES_ROOT
+
+# noinspection PyUnresolvedReferences
+# pylint: disable=unused-import
+from trapper.common.testing_utils.pytest_fixtures import (
+    create_data_collator_args,
+    create_data_processor_args,
+    get_raw_dataset,
+    make_data_collator,
+    make_sequential_sampler,
+    temp_output_dir,
+    temp_result_dir,
+)
+
+_HF_DATASETS_FIXTURES_ROOT = POS_TAGGING_FIXTURES_ROOT / "hf_datasets"
+
+
+@pytest.fixture(scope="package")
+def get_hf_datasets_fixture_path():
+    def _get_hf_datasets_fixture_path(dataset: str) -> str:
+        return str(_HF_DATASETS_FIXTURES_ROOT / dataset)
+
+    return _get_hf_datasets_fixture_path
diff --git a/examples/pos_tagging/tests/test_data_adapter.py b/examples/pos_tagging/tests/test_data_adapter.py
new file mode 100644
index 0000000..a670144
--- /dev/null
+++ b/examples/pos_tagging/tests/test_data_adapter.py
@@ -0,0 +1,89 @@
+import pytest
+
+from examples.pos_tagging.src.data_adapter import ExampleDataAdapterForPosTagging
+from examples.pos_tagging.src.data_processor import (
+    ExampleConll2003PosTaggingDataProcessor,
+)
+from examples.pos_tagging.src.tokenizer_wrapper import (
+    ExamplePosTaggingTokenizerWrapper,
+)
+from trapper.common.constants import IGNORED_LABEL_ID
+from trapper.data import InputBatch
+
+
+@pytest.fixture(scope="module")
+def data_collator_args(create_data_collator_args):
+    return create_data_collator_args(
+        task_type="token_classification",
+        train_batch_size=1,
+        validation_batch_size=1,
+        is_distributed=False,
+        model_max_sequence_length=512,
+        tokenizer_factory=ExamplePosTaggingTokenizerWrapper,
+        tokenizer_model_name="roberta-base",
+        add_prefix_space=True
+    )
+
+
+@pytest.fixture(scope="module")
+def raw_conll03_postagging_dataset(get_raw_dataset, get_hf_datasets_fixture_path):
+    return get_raw_dataset(
+        path=get_hf_datasets_fixture_path("conll2003_test_fixture"),
+        split="train"
+    )
+
+
+@pytest.fixture(scope="module")
+def adapted_conll03_postagging_dataset(raw_conll03_postagging_dataset,
+                                       data_collator_args):
+    data_adapter = ExampleDataAdapterForPosTagging(
+        data_collator_args.tokenizer_wrapper)
+    data_processor = ExampleConll2003PosTaggingDataProcessor(
+        data_collator_args.tokenizer_wrapper
+    )
+    processed_dataset = raw_conll03_postagging_dataset.map(data_processor)
+    return processed_dataset.map(data_adapter)
+
+
+@pytest.fixture(scope="module")
+def data_collator(make_data_collator, data_collator_args):
+    return make_data_collator(data_collator_args)
+
+
+def test_batch_content_on_squad_dev_dataset(
+        raw_conll03_postagging_dataset,
+        adapted_conll03_postagging_dataset,
+        data_collator_args,
+        data_collator,
+):
+    expected_sentence = "The European Commission said on Thursday it disagreed with German advice to consumers to shun British lamb until scientists determine whether mad cow disease can be transmitted to sheep."
+    if data_collator_args.is_tokenizer_uncased:
+        expected_sentence = expected_sentence.lower()
+
+    collated_batch = data_collator.build_model_inputs(
+        adapted_conll03_postagging_dataset)
+    input_ids = collated_batch["input_ids"][0]
+    labels = collated_batch["labels"][0]
+
+    tokenizer = data_collator_args.tokenizer_wrapper.tokenizer
+    decoded_sentence = tokenizer.decode(
+        input_ids, skip_special_tokens=True).lstrip()
+    assert expected_sentence == decoded_sentence
+    assert len(input_ids) == len(labels)
+
+    encoding = tokenizer(expected_sentence, add_special_tokens=False)
+    raw_pos_tags = [12, 22, 22, 38, 15, 22, 28, 38, 15, 16, 21, 35, 24, 35, 37, 16,
+                    21, 15, 24, 41, 15, 16, 21, 21, 20, 37, 40, 35, 21, 7]
+    expected_labels = [raw_pos_tags[ind] for ind in encoding.word_ids()]
+    expected_labels.insert(0, IGNORED_LABEL_ID)  # BOS
+    expected_labels.append(IGNORED_LABEL_ID)  # EOS
+    assert expected_labels == labels
+    validate_attention_mask(collated_batch)
+
+
+def validate_attention_mask(instance_batch: InputBatch):
+    for input_ids, attention_mask in zip(
+            instance_batch["input_ids"], instance_batch["attention_mask"]
+    ):
+        assert len(attention_mask) == len(input_ids)
+        assert all(val == 1 for val in attention_mask)
diff --git a/examples/pos_tagging/tests/test_data_processor.py b/examples/pos_tagging/tests/test_data_processor.py
new file mode 100644
index 0000000..1bc26f1
--- /dev/null
+++ b/examples/pos_tagging/tests/test_data_processor.py
@@ -0,0 +1,41 @@
+import pytest
+
+from examples.pos_tagging.src.data_processor import (
+    ExampleConll2003PosTaggingDataProcessor,
+)
+from examples.pos_tagging.src.tokenizer_wrapper import (
+    ExamplePosTaggingTokenizerWrapper,
+)
+
+
+@pytest.fixture(scope="module")
+def args(create_data_processor_args):
+    return create_data_processor_args(
+        tokenizer_factory=ExamplePosTaggingTokenizerWrapper,
+        tokenizer_model_name="roberta-base",
+        model_max_sequence_length=512,
+        add_prefix_space=True
+    )
+
+
+def test_data_processor(get_raw_dataset, args, get_hf_datasets_fixture_path):
+    expected_sentence = "The European Commission said on Thursday it disagreed with German advice to consumers to shun British lamb until scientists determine whether mad cow disease can be transmitted to sheep."
+    if args.is_tokenizer_uncased:
+        expected_sentence = expected_sentence.lower()
+    data_processor = ExampleConll2003PosTaggingDataProcessor(
+        args.tokenizer_wrapper)
+    raw_dataset = get_raw_dataset(
+        path=get_hf_datasets_fixture_path("conll2003_test_fixture"),
+        split="train"
+    )
+    processed_instance = raw_dataset.map(data_processor)[0]
+    tokenizer = args.tokenizer_wrapper.tokenizer
+    decoded_sentence = tokenizer.decode(processed_instance["tokens"]).lstrip()
+    assert expected_sentence == decoded_sentence
+    assert len(processed_instance["tokens"]) == len(processed_instance["pos_tags"])
+
+    encoding = tokenizer(expected_sentence, add_special_tokens=False)
+    raw_pos_tags = [12, 22, 22, 38, 15, 22, 28, 38, 15, 16, 21, 35, 24, 35, 37, 16,
+                    21, 15, 24, 41, 15, 16, 21, 21, 20, 37, 40, 35, 21, 7]
+    expected_pos_tags = [raw_pos_tags[ind] for ind in encoding.word_ids()]
+    assert expected_pos_tags == processed_instance["pos_tags"]
diff --git a/examples/pos_tagging/tests/test_trainer.py b/examples/pos_tagging/tests/test_trainer.py
new file mode 100644
index 0000000..99d7484
--- /dev/null
+++ b/examples/pos_tagging/tests/test_trainer.py
@@ -0,0 +1,96 @@
+import datasets
+import pytest
+
+# needed for registering the data-related classes
+# noinspection PyUnresolvedReferences
+# pylint: disable=unused-import
+import src
+from transformers import DistilBertForTokenClassification, DistilBertTokenizerFast
+
+from trapper.common import Params
+from trapper.data.data_collator import DataCollator
+from trapper.training import TransformerTrainer, TransformerTrainingArguments
+from trapper.training.optimizers import HuggingfaceAdamWOptimizer
+from trapper.training.train import run_experiment_using_trainer
+
+
+@pytest.fixture(scope="module")
+def trainer_params(temp_output_dir, temp_result_dir, get_hf_datasets_fixture_path):
+    params_dict = {
+        "pretrained_model_name_or_path": "distilbert-base-uncased",
+        "train_split_name": "train",
+        "dev_split_name": "validation",
+        "tokenizer_wrapper": {
+            "type": "pos_tagging_example",
+            "add_prefix_space": True
+        },
+        "dataset_loader": {
+            "dataset_reader": {
+                "path": get_hf_datasets_fixture_path("conll2003_test_fixture"),
+            },
+            "data_processor": {
+                "type": "conll2003_pos_tagging_example",
+                "model_max_sequence_length": 512,
+            },
+            "data_adapter": {"type": "conll2003_pos_tagging_example"},
+        },
+        "data_collator": {},
+        "model_wrapper": {"type": "token_classification", "num_labels": 47},
+        "compute_metrics": {"type": "seqeval",
+                            "return_entity_level_metrics": False},
+        "label_mapper": {"type": "conll2003_pos_tagging_example"},
+        "args": {
+            "type": "default",
+            "output_dir": temp_output_dir + "/checkpoints",
+            "result_dir": temp_result_dir,
+            "num_train_epochs": 3,
+            "per_device_train_batch_size": 1,
+            "per_device_eval_batch_size": 1,
+            "logging_dir": temp_output_dir + "/logs",
+            "no_cuda": True,
+            "logging_steps": 1,
+            "evaluation_strategy": "steps",
+            "save_steps": 2,
+            "label_names": ["labels"],
+            "lr_scheduler_type": "linear",
+            "warmup_steps": 2,
+            "do_train": True,
+            "do_eval": True,
+            "save_total_limit": 1,
+            "metric_for_best_model": "eval_loss",
+            "greater_is_better": False,
+        },
+        "optimizer": {
+            "type": "huggingface_adamw",
+            "weight_decay": 0.01,
+            "parameter_groups": [
+                [
+                    ["bias", "LayerNorm\\\\.weight", "layer_norm\\\\.weight"],
+                    {"weight_decay": 0},
+                ]
+            ],
+            "lr": 5e-5,
+            "eps": 1e-8,
+        },
+    }
+    return Params(params_dict)
+
+
+@pytest.fixture(scope="module")
+def trainer(trainer_params) -> TransformerTrainer:
+    return TransformerTrainer.from_params(trainer_params)
+
+
+def test_trainer_fields(trainer):
+    assert isinstance(trainer, TransformerTrainer)
+    assert isinstance(trainer.model, DistilBertForTokenClassification)
+    assert isinstance(trainer.args, TransformerTrainingArguments)
+    assert isinstance(trainer.data_collator, DataCollator)
+    assert isinstance(trainer.train_dataset, datasets.Dataset)
+    assert isinstance(trainer.eval_dataset, datasets.Dataset)
+    assert isinstance(trainer.tokenizer, DistilBertTokenizerFast)
+    assert isinstance(trainer.optimizer, HuggingfaceAdamWOptimizer)
+
+
+def test_trainer_can_train(trainer):
+    run_experiment_using_trainer(trainer)
diff --git a/test_fixtures/hf_datasets/squad_qa_test_fixture/squad_qa_test_fixture.py b/test_fixtures/hf_datasets/squad_qa_test_fixture/squad_qa_test_fixture.py
index 652508d..a4875b7 100644
--- a/test_fixtures/hf_datasets/squad_qa_test_fixture/squad_qa_test_fixture.py
+++ b/test_fixtures/hf_datasets/squad_qa_test_fixture/squad_qa_test_fixture.py
@@ -28,10 +28,9 @@
 _DESCRIPTION = """\
 Small dataset taken from SQuAD v1.1 and used for testing purposes.
 """
-_URL = "test_fixtures/hf_datasets/squad_qa_test_fixture/"
 _URLS = {
-    "train": _URL + "train.json",
-    "dev": _URL + "dev.json",
+    "train": "train.json",
+    "dev": "dev.json",
 }
 
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 47aa585..877df4b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,16 +2,24 @@
 
 import pytest
 
-_PROJECT_ROOT = Path(__file__).parent.parent.resolve()
-_TESTS_ROOT = _PROJECT_ROOT / "tests/trapper"
-_FIXTURES_ROOT = _PROJECT_ROOT / "test_fixtures"
+from trapper import FIXTURES_ROOT, TESTS_ROOT
+
+_HF_DATASETS_FIXTURES_ROOT = FIXTURES_ROOT / "hf_datasets"
 
 
 @pytest.fixture(scope="package")
 def tests_root():
-    return _TESTS_ROOT
+    return TESTS_ROOT
 
 
 @pytest.fixture(scope="package")
-def fixtures_root():
-    return _FIXTURES_ROOT
+def fixtures_root() -> Path:
+    return FIXTURES_ROOT
+
+
+@pytest.fixture(scope="package")
+def get_hf_datasets_fixture_path():
+    def _get_hf_datasets_fixture_path(dataset: str) -> str:
+        return str(_HF_DATASETS_FIXTURES_ROOT / dataset)
+
+    return _get_hf_datasets_fixture_path
diff --git a/tests/data/data_collators/squad/test_question_answering_collator.py b/tests/data/data_collators/squad/test_question_answering_collator_and_adapter.py
similarity index 82%
rename from tests/data/data_collators/squad/test_question_answering_collator.py
rename to tests/data/data_collators/squad/test_question_answering_collator_and_adapter.py
index d43e0ff..28ea4a0 100644
--- a/tests/data/data_collators/squad/test_question_answering_collator.py
+++ b/tests/data/data_collators/squad/test_question_answering_collator_and_adapter.py
@@ -1,18 +1,19 @@
 import pytest
 from torch.utils.data import DataLoader
+from transformers import PreTrainedTokenizerBase
 
 from trapper.data import SquadQuestionAnsweringDataProcessor
 from trapper.data.data_adapters.question_answering_adapter import (
     DataAdapterForQuestionAnswering,
 )
 from trapper.data.data_collator import InputBatch
-from trapper.data.tokenizers import QuestionAnsweringTokenizer
+from trapper.data.tokenizers import QuestionAnsweringTokenizerWrapper
 
 
 @pytest.fixture(scope="module")
 def data_collator_args(create_data_collator_args):
     return create_data_collator_args(
-        tokenizer_cls=QuestionAnsweringTokenizer,
+        tokenizer_factory=QuestionAnsweringTokenizerWrapper,
         train_batch_size=2,
         validation_batch_size=1,
         tokenizer_model_name="roberta-base",
@@ -22,16 +23,21 @@ def data_collator_args(create_data_collator_args):
 
 
 @pytest.fixture(scope="module")
-def processed_dataset(get_raw_dataset, data_collator_args):
+def processed_dataset(get_raw_dataset,
+                      data_collator_args,
+                      get_hf_datasets_fixture_path):
     data_processor = SquadQuestionAnsweringDataProcessor(
-        data_collator_args.tokenizer)
-    raw_dataset = get_raw_dataset(path="squad_qa_test_fixture")
+        data_collator_args.tokenizer_wrapper)
+    raw_dataset = get_raw_dataset(
+        path=get_hf_datasets_fixture_path("squad_qa_test_fixture")
+    )
     return raw_dataset.map(data_processor)
 
 
 @pytest.fixture(scope="module")
 def adapted_dataset(processed_dataset, data_collator_args):
-    data_adapter = DataAdapterForQuestionAnswering(data_collator_args.tokenizer)
+    data_adapter = DataAdapterForQuestionAnswering(
+        data_collator_args.tokenizer_wrapper)
     return processed_dataset.map(data_adapter)
 
 
@@ -88,7 +94,10 @@ def test_batch_content(
     if data_collator_args.is_tokenizer_uncased:
         expected_question = expected_question.lower()
     validate_target_question_positions_using_decoded_tokens(
-        expected_question, index, data_collator_args.tokenizer, collated_batch
+        expected_question,
+        index,
+        data_collator_args.tokenizer_wrapper.tokenizer,
+        collated_batch
     )
 
     instance = processed_dataset["validation"][index]
@@ -98,7 +107,10 @@ def test_batch_content(
 
 
 def validate_target_question_positions_using_decoded_tokens(
-        expected_question, index, tokenizer, input_batch: InputBatch
+        expected_question,
+        index,
+        tokenizer: PreTrainedTokenizerBase,
+        input_batch: InputBatch
 ):
     input_ids = input_batch["input_ids"][index]
     question_start = -sum(input_batch["token_type_ids"][index])
diff --git a/tests/data/data_processors/squad/test_question_answering_processor.py b/tests/data/data_processors/squad/test_question_answering_processor.py
index dbbe9c7..0e8082c 100644
--- a/tests/data/data_processors/squad/test_question_answering_processor.py
+++ b/tests/data/data_processors/squad/test_question_answering_processor.py
@@ -1,22 +1,26 @@
 from typing import Dict
 
 import pytest
+from transformers import PreTrainedTokenizerBase
 
 from trapper.data import SquadQuestionAnsweringDataProcessor
-from trapper.data.tokenizers import QuestionAnsweringTokenizer
+from trapper.data.tokenizers import QuestionAnsweringTokenizerWrapper
 
 
 @pytest.fixture(scope="module")
 def args(create_data_processor_args):
     return create_data_processor_args(
-        tokenizer_cls=QuestionAnsweringTokenizer,
+        tokenizer_factory=QuestionAnsweringTokenizerWrapper,
         tokenizer_model_name="roberta-base")
 
 
 @pytest.fixture(scope="module")
-def processed_dev_dataset(get_raw_dataset, args):
-    data_processor = SquadQuestionAnsweringDataProcessor(args.tokenizer)
-    raw_dataset = get_raw_dataset(path="squad_qa_test_fixture", split="validation")
+def processed_dev_dataset(get_raw_dataset, args, get_hf_datasets_fixture_path):
+    data_processor = SquadQuestionAnsweringDataProcessor(args.tokenizer_wrapper)
+    raw_dataset = get_raw_dataset(
+        path=get_hf_datasets_fixture_path("squad_qa_test_fixture"),
+        split="validation"
+    )
     return raw_dataset.map(data_processor)
 
 
@@ -32,10 +36,11 @@ def test_data_processor(processed_dev_dataset, args, index, expected_question):
     if args.is_tokenizer_uncased:
         expected_question = expected_question.lower()
     processed_instance = processed_dev_dataset[index]
-    assert decode_question(processed_instance, args.tokenizer) == expected_question
+    assert decode_question(processed_instance,
+                           args.tokenizer_wrapper.tokenizer) == expected_question
 
 
-def decode_question(instance: Dict, tokenizer: QuestionAnsweringTokenizer
+def decode_question(instance: Dict, tokenizer: PreTrainedTokenizerBase
                     ) -> str:
     field_value = tokenizer.decode(instance["question"])
     return field_value.lstrip()
diff --git a/tests/data/data_processors/test_data_processor.py b/tests/data/data_processors/test_data_processor.py
index f0d7b60..059505a 100644
--- a/tests/data/data_processors/test_data_processor.py
+++ b/tests/data/data_processors/test_data_processor.py
@@ -8,6 +8,10 @@
 
 
 class MockTokenizer:
+    @property
+    def model_max_length(self):
+        return 1
+
     @staticmethod
     def convert_tokens_to_ids(text: str) -> List[int]:
         return [int(tok.split("token")[-1]) for tok in text]
@@ -17,6 +21,15 @@ def tokenize(text: str) -> List[str]:
         return text.split()
 
 
+class MockTokenizerWrapper:
+    def __init__(self):
+        self._tokenizer = MockTokenizer()
+
+    @property
+    def tokenizer(self):
+        return self._tokenizer
+
+
 class MockDataProcessor(DataProcessor):
     def text_to_instance(self, ind: int, info1: str, info2: str = None
                          ) -> IndexedInstance:
@@ -48,7 +61,7 @@ def dummy_dataset():
 
 @pytest.fixture
 def mock_processor():
-    mock_tokenizer = MockTokenizer()
+    mock_tokenizer = MockTokenizerWrapper()
     return MockDataProcessor(mock_tokenizer)  # type: ignore
 
 
diff --git a/tests/data/test_dataset_reader.py b/tests/data/test_dataset_reader.py
index 1c4d910..69815fe 100644
--- a/tests/data/test_dataset_reader.py
+++ b/tests/data/test_dataset_reader.py
@@ -5,9 +5,11 @@
 
 
 @pytest.fixture
-def dataset_reader():
-    return DatasetReader(path="squad_qa_test_fixture",
-                         download_config=DownloadConfig(local_files_only=True))
+def dataset_reader(get_hf_datasets_fixture_path):
+    return DatasetReader(
+        path=get_hf_datasets_fixture_path("squad_qa_test_fixture"),
+        download_config=DownloadConfig(local_files_only=True)
+    )
 
 
 def test_read(dataset_reader):
diff --git a/tests/data/test_label_mapper.py b/tests/data/test_label_mapper.py
new file mode 100644
index 0000000..24436c6
--- /dev/null
+++ b/tests/data/test_label_mapper.py
@@ -0,0 +1,46 @@
+import pytest
+
+from trapper.common.utils import is_equal
+from trapper.data.label_mapper import LabelMapper
+
+DUMMY_LABELS = ("label1", "label2", "label3")
+
+
+class MockLabelMapperWithClassVariable(LabelMapper):
+    _LABELS = DUMMY_LABELS
+
+
+class MockLabelMapperWithConstructor(LabelMapper):
+    def __init__(self, labels):
+        super().__init__(labels)
+
+
+@pytest.mark.parametrize(
+    "start_id",
+    [0, 7],
+)
+def test_mapper_created_from_labels(start_id):
+    dummy_ids = (start_id, start_id + 1, start_id + 2)
+    dummy_label_to_id_map = {"label1": dummy_ids[0],
+                             "label2": dummy_ids[1],
+                             "label3": dummy_ids[2]}
+    dummy_id_to_label_map = {
+        id_: label for label, id_ in dummy_label_to_id_map.items()}
+
+    mapper_from_label_to_id_map = MockLabelMapperWithConstructor(
+        dummy_label_to_id_map)
+    mapper_from_class_variable_labels = MockLabelMapperWithClassVariable.from_labels(
+        start_id=start_id)
+    mapper_from_labels = MockLabelMapperWithConstructor.from_labels(
+        labels=DUMMY_LABELS, start_id=start_id)
+
+    for mapper in (mapper_from_label_to_id_map,
+                   mapper_from_class_variable_labels,
+                   mapper_from_labels):
+        assert mapper.labels == DUMMY_LABELS
+        assert mapper.ids == dummy_ids
+        assert is_equal(mapper.label_to_id_map, dummy_label_to_id_map)
+        assert is_equal(mapper.id_to_label_map, dummy_id_to_label_map)
+        for label, id_ in dummy_label_to_id_map.items():
+            assert mapper.get_id(label) == id_
+            assert mapper.get_label(id_) == label
diff --git a/tests/data/test_tokenizer_wrapper.py b/tests/data/test_tokenizer_wrapper.py
new file mode 100644
index 0000000..1e214da
--- /dev/null
+++ b/tests/data/test_tokenizer_wrapper.py
@@ -0,0 +1,65 @@
+"""
+Test if the tokenizer wrapper handles correctly the general special tokens such
+as BOS and EOS as well as the task-specific special tokens such as <CONTEXT>.
+"""
+import pytest
+from transformers import PreTrainedTokenizerBase
+
+from trapper.common.constants import ANSWER_TOKEN, CONTEXT_TOKEN
+from trapper.data import TokenizerWrapper
+
+
+@pytest.fixture
+def bert_tokenizer_with_context_token():
+    class TokenizerWrapperWithContextToken(TokenizerWrapper):
+        _TASK_SPECIFIC_SPECIAL_TOKENS = [CONTEXT_TOKEN]
+
+    return TokenizerWrapperWithContextToken.from_pretrained("bert-base-uncased")
+
+
+@pytest.fixture
+def gpt2_tokenizer_with_context_and_answer_tokens():
+    class TokenizerWrapperWithContextAndAnswerTokens(TokenizerWrapper):
+        _TASK_SPECIFIC_SPECIAL_TOKENS = [CONTEXT_TOKEN, ANSWER_TOKEN]
+
+    return TokenizerWrapperWithContextAndAnswerTokens.from_pretrained(
+        "gpt2")
+
+
+def test_bert_tokenizer(bert_tokenizer_with_context_token):
+    assert (bert_tokenizer_with_context_token.num_added_special_tokens
+            == 1)  # CONTEXT_TOKEN
+    tokenizer = bert_tokenizer_with_context_token.tokenizer
+    assert tokenizer.bos_token == tokenizer.cls_token == "[CLS]"
+    assert tokenizer.bos_token_id == tokenizer.cls_token_id == 101
+    assert tokenizer.eos_token == tokenizer.sep_token == "[SEP]"
+    assert tokenizer.sep_token_id == tokenizer.eos_token_id == 102
+    assert_special_tokens_are_preserved(tokenizer, CONTEXT_TOKEN)
+    assert_all_common_special_tokens_are_present(tokenizer)
+
+
+def test_gpt2_tokenizer(gpt2_tokenizer_with_context_and_answer_tokens):
+    assert (gpt2_tokenizer_with_context_and_answer_tokens.num_added_special_tokens
+            == 4)  # PAD_TOKEN, MASK_TOKEN, CONTEXT_TOKEN, ANSWER_TOKEN.
+    tokenizer = gpt2_tokenizer_with_context_and_answer_tokens.tokenizer
+    assert tokenizer.bos_token == tokenizer.cls_token == "
"
+    assert tokenizer.bos_token_id == tokenizer.cls_token_id == 50256
+    assert tokenizer.eos_token == tokenizer.sep_token == "
"
+    assert tokenizer.eos_token_id == tokenizer.sep_token_id == 50256
+    for token in [CONTEXT_TOKEN, ANSWER_TOKEN]:
+        assert_special_tokens_are_preserved(tokenizer, token)
+    assert_all_common_special_tokens_are_present(tokenizer)
+
+
+def assert_special_tokens_are_preserved(tokenizer: PreTrainedTokenizerBase,
+                                        token: str):
+    assert len(tokenizer.tokenize(token)) == 1
+
+
+def assert_all_common_special_tokens_are_present(
+        tokenizer: PreTrainedTokenizerBase):
+    COMMON_TOKENS = (
+        "bos_token", "eos_token", "cls_token", "sep_token", "pad_token",
+        "mask_token", "unk_token"
+    )
+    assert all(hasattr(tokenizer, key) for key in COMMON_TOKENS)
diff --git a/tests/training/test_trainer_for_question_answering.py b/tests/training/test_trainer_for_question_answering.py
index a227cdc..35ad1e2 100644
--- a/tests/training/test_trainer_for_question_answering.py
+++ b/tests/training/test_trainer_for_question_answering.py
@@ -2,9 +2,6 @@
 import pytest
 from transformers import DistilBertForQuestionAnswering, DistilBertTokenizerFast
 
-# noinspection PyUnresolvedReferences
-# pylint: disable=unused-import
-import trapper.data  # needed for registering the data-related classes
 from trapper.common import Params
 from trapper.data.data_collator import DataCollator
 from trapper.training import TransformerTrainer, TransformerTrainingArguments
@@ -13,19 +10,21 @@
 
 
 @pytest.fixture(scope="module")
-def trainer_params(temp_output_dir, temp_result_dir):
+def trainer_params(temp_output_dir, temp_result_dir, get_hf_datasets_fixture_path):
     params_dict = {
         "pretrained_model_name_or_path": "distilbert-base-uncased",
         "train_split_name": "train",
         "dev_split_name": "validation",
-        "tokenizer": {"type": "question-answering"},
+        "tokenizer_wrapper": {"type": "question-answering"},
         "dataset_loader": {
-            "dataset_reader": {"path": "squad_qa_test_fixture"},
+            "dataset_reader": {
+                "path": get_hf_datasets_fixture_path("squad_qa_test_fixture")
+            },
             "data_processor": {"type": "squad-question-answering"},
             "data_adapter": {"type": "question-answering"},
         },
         "data_collator": {},
-        "model": {"type": "question_answering"},
+        "model_wrapper": {"type": "question_answering"},
         "args": {
             "type": "default",
             "output_dir": temp_output_dir + "/checkpoints",
diff --git a/trapper/common/testing_utils/pytest_fixtures/data.py b/trapper/common/testing_utils/pytest_fixtures/data.py
index 0acdb0b..9b68d2b 100644
--- a/trapper/common/testing_utils/pytest_fixtures/data.py
+++ b/trapper/common/testing_utils/pytest_fixtures/data.py
@@ -1,12 +1,12 @@
 from dataclasses import dataclass
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Type, Union
 
 import pytest
 from datasets import DownloadConfig
 from torch.utils.data import SequentialSampler
 from transformers.trainer_pt_utils import SequentialDistributedSampler
 
-from trapper.data import DataCollator, DatasetReader, TransformerTokenizer
+from trapper.data import DataCollator, DatasetReader, TokenizerWrapper
 from trapper.data.dataset_reader import TrapperDataset, TrapperDatasetDict
 from trapper.models.auto_wrappers import _TASK_TO_INPUT_FIELDS
 
@@ -67,7 +67,8 @@ def _get_raw_dataset(
 
 @dataclass
 class DataProcessorArguments:
-    tokenizer_cls: TransformerTokenizer
+    model_max_sequence_length: int = None
+    tokenizer_factory: Type[TokenizerWrapper] = TokenizerWrapper
     tokenizer_model_name: str = "roberta-base"
     tokenizer_kwargs: Dict = None
 
@@ -77,21 +78,23 @@ def __post_init__(self):
         else:
             self.is_tokenizer_uncased = False
         tokenizer_kwargs = self.tokenizer_kwargs or {}
-        self.tokenizer = self.tokenizer_cls.from_pretrained(
+        self.tokenizer_wrapper = self.tokenizer_factory.from_pretrained(
             self.tokenizer_model_name, **tokenizer_kwargs
         )
-        del self.tokenizer_cls, self.tokenizer_kwargs
+        del self.tokenizer_factory, self.tokenizer_kwargs
 
 
 @pytest.fixture(scope="session")
 def create_data_processor_args():
     def _create_data_processor_args(
-        tokenizer_cls: TransformerTokenizer,
+        model_max_sequence_length: int = None,
+        tokenizer_factory: Type[TokenizerWrapper] = TokenizerWrapper,
         tokenizer_model_name: str = "roberta-base",
         **tokenizer_kwargs,
     ) -> DataProcessorArguments:
         return DataProcessorArguments(
-            tokenizer_cls=tokenizer_cls,
+            model_max_sequence_length=model_max_sequence_length,
+            tokenizer_factory=tokenizer_factory,
             tokenizer_model_name=tokenizer_model_name,
             tokenizer_kwargs=tokenizer_kwargs,
         )
@@ -114,16 +117,18 @@ def __post_init__(self):
 @pytest.fixture(scope="session")
 def create_data_collator_args():
     def _create_data_collator_args(
-        tokenizer_cls: TransformerTokenizer,
         train_batch_size: int,
         validation_batch_size: int,
+        model_max_sequence_length: int = None,
+        tokenizer_factory: Type[TokenizerWrapper] = TokenizerWrapper,
         tokenizer_model_name: str = "roberta-base",
         task_type: str = "question_answering",
         is_distributed: bool = False,
         **tokenizer_kwargs,
     ) -> DataProcessorArguments:
         return DataCollatorArguments(
-            tokenizer_cls=tokenizer_cls,
+            model_max_sequence_length=model_max_sequence_length,
+            tokenizer_factory=tokenizer_factory,
             tokenizer_kwargs=tokenizer_kwargs,
             train_batch_size=train_batch_size,
             validation_batch_size=validation_batch_size,
@@ -138,7 +143,7 @@ def _create_data_collator_args(
 @pytest.fixture(scope="session")
 def make_data_collator():
     def _make_data_collator(args: DataCollatorArguments):
-        return DataCollator(args.tokenizer, args.model_forward_params)
+        return DataCollator(args.tokenizer_wrapper, args.model_forward_params)
 
     return _make_data_collator