Repository: obss/trapper (language: python)

Issue:
Improve the HF datasets integration using dataset.map() 
We can do the pre-processing using data processors as callable argument to dataset.map()


Discussion / hints:

Fix patch:
diff --git a/examples/experiment.py b/examples/experiment.py
index 9f12c26..43fde2c 100644
--- a/examples/experiment.py
+++ b/examples/experiment.py
@@ -37,8 +37,9 @@ def download_squad(
     paths = []
 
     for dataset in datasets:
+        dest_name = "train.json" if "train" in dataset else "dev.json"
         url = os.path.join(dataset_base_url, dataset)
-        dest = os.path.join(destination_dir, dataset)
+        dest = os.path.join(destination_dir, dest_name)
         paths.append(dest)
 
         if not overwrite and os.path.exists(dest):
@@ -120,12 +121,8 @@ def start_experiment(config: str, task: str, ext_vars: Dict[str, str]):
     CHECKPOINT_DIR = os.path.join(EXPERIMENT_DIR, "checkpoints")
     OUTPUT_DIR = os.path.join(EXPERIMENT_DIR, "outputs")
 
-    train_path, dev_path = download_squad(task=TASK)
-
     ext_vars = {
         # Used to feed the jsonnet config file with file paths
-        "TRAIN_DATA_PATH": train_path,
-        "DEV_DATA_PATH": dev_path,
         "OUTPUT_PATH": OUTPUT_DIR,
         "CHECKPOINT_PATH": CHECKPOINT_DIR,
     }
diff --git a/examples/experiments/question-answering/experiment.jsonnet b/examples/experiments/question-answering/experiment.jsonnet
index 344602b..0078eae 100644
--- a/examples/experiments/question-answering/experiment.jsonnet
+++ b/examples/experiments/question-answering/experiment.jsonnet
@@ -1,19 +1,27 @@
-local train_path = std.extVar("TRAIN_DATA_PATH");
-local dev_path = std.extVar("DEV_DATA_PATH");
 local checkpoint_dir = std.extVar("CHECKPOINT_PATH");
 local result_dir = std.extVar("OUTPUT_PATH");
 {
+    "train_split_name": "train",
+    "dev_split_name": "validation",
     "pretrained_model_name_or_path": "roberta-base",
-    "train_file_path": train_path,
-    "dev_file_path": dev_path,
     "tokenizer": {
         "type": "question-answering"
     },
-    "data_collator": {
-        "type": "question-answering"
+    "dataset_loader": {
+        "type": "default",
+        "dataset_reader": {
+            "type": "default",
+            "path": "squad"
+        },
+        "data_processor": {
+            "type": "squad-question-answering"
+        },
+        "data_adapter": {
+            "type": "question-answering"
+        }
     },
-    "dataset_reader": {
-        "type": "squad-question-answering"
+    "data_collator":{
+        "type": "default"
     },
     "model": {
         "type": "question_answering"
diff --git a/examples/question_answering.ipynb b/examples/question_answering.ipynb
index 90f53af..850f681 100644
--- a/examples/question_answering.ipynb
+++ b/examples/question_answering.ipynb
@@ -9,6 +9,17 @@
     "This notebook serves as a walkthrough for training with trapper package."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "cd .."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -61,7 +72,7 @@
     "EXPERIMENT_NAME = \"roberta-base-training-example\"\n",
     "\n",
     "WORKING_DIR = os.getcwd()\n",
-    "EXPERIMENTS_DIR = os.path.join(WORKING_DIR, \"experiments\")"
+    "EXPERIMENTS_DIR = os.path.join(WORKING_DIR, \"examples/experiments\")"
    ]
   },
   {
@@ -79,43 +90,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def download_squad(task: str, version: str = \"1.1\", overwrite: bool = False) -> Tuple[str, str]:\n",
-    "    \"\"\"\n",
-    "    Downloads SQuAD dataset with given version.\n",
-    "\n",
-    "    Args:\n",
-    "        version: SQuAD dataset version.\n",
-    "        overwrite: If true, overwrites the destination file.\n",
-    "\n",
-    "    Returns: (train set path, dev set path) local paths of downloaded dataset files.\n",
-    "\n",
-    "    \"\"\"\n",
-    "    destination_dir = DATASET_DIR.format(task=task)\n",
-    "    os.makedirs(destination_dir, exist_ok=True)\n",
-    "    dataset_base_url = \"https://rajpurkar.github.io/SQuAD-explorer/dataset/\"\n",
-    "\n",
-    "    train_set = f\"train-v{version}.json\"\n",
-    "    dev_set = f\"dev-v{version}.json\"\n",
-    "\n",
-    "    datasets = [train_set, dev_set]\n",
-    "    paths = []\n",
-    "\n",
-    "    for dataset in datasets:\n",
-    "        url = os.path.join(dataset_base_url, dataset)\n",
-    "        dest = os.path.join(destination_dir, dataset)\n",
-    "        paths.append(dest)\n",
-    "\n",
-    "        if not overwrite and os.path.exists(dest):\n",
-    "            warnings.warn(f\"{dest} already exists, not overwriting.\")\n",
-    "            continue\n",
-    "\n",
-    "        r = requests.get(url, allow_redirects=True)\n",
-    "\n",
-    "        with open(dest, 'wb') as out_file:\n",
-    "            out_file.write(r.content)\n",
-    "\n",
-    "    return paths[0], paths[1]\n",
-    "\n",
     "def get_dir_from_task(path: str, task: str):\n",
     "    task = \"unnamed-task\" if task is None else task\n",
     "    return path.format(task=task)\n",
@@ -144,9 +118,7 @@
     "EXPERIMENT_DIR = os.path.join(TASK_DIR, EXPERIMENT_NAME)\n",
     "MODEL_DIR = os.path.join(EXPERIMENT_DIR, \"model\")\n",
     "CHECKPOINT_DIR = os.path.join(EXPERIMENT_DIR, \"checkpoints\")\n",
-    "OUTPUT_DIR = os.path.join(EXPERIMENT_DIR, \"outputs\")\n",
-    "\n",
-    "train_path, dev_path = download_squad(task=TASK)"
+    "OUTPUT_DIR = os.path.join(EXPERIMENT_DIR, \"outputs\")"
    ]
   },
   {
@@ -159,10 +131,8 @@
    "source": [
     "ext_vars = {\n",
     "    # Used to feed the jsonnet config file with file paths\n",
-    "    \"TRAIN_DATA_PATH\": train_path,\n",
-    "    \"DEV_DATA_PATH\": dev_path,\n",
     "    \"OUTPUT_PATH\": OUTPUT_DIR,\n",
-    "    \"CHECKPOINT_PATH\": CHECKPOINT_DIR,\n",
+    "    \"CHECKPOINT_PATH\": CHECKPOINT_DIR\n",
     "}\n",
     "\n",
     "CONFIG_PATH = os.path.join(TASK_DIR, \"experiment.jsonnet\")  # default experiment params"
@@ -179,7 +149,7 @@
     "result = start_experiment(\n",
     "    config=CONFIG_PATH,\n",
     "    task=TASK,\n",
-    "    ext_vars=ext_vars\n",
+    "    ext_vars=ext_vars,\n",
     ")"
    ]
   },
@@ -358,7 +328,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
diff --git a/pyproject.toml b/pyproject.toml
index 5497a9a..fc2998e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,12 +1,19 @@
 [tool.black]
 line-length = 84
+include = 'trapper'
 exclude = '''
 (
   /(
     | .git
-    | venv
-    | .venv
+    | .github
+    | .idea
+    | trapper.egg-info
     | .pytest_cache
+    | .venv
+    | __pycache__
+    | examples
+    | test_data
+    | venv
   )/
 )
 '''
\ No newline at end of file
diff --git a/trapper/common/__init__.py b/trapper/common/__init__.py
index 54f1190..6302f52 100644
--- a/trapper/common/__init__.py
+++ b/trapper/common/__init__.py
@@ -1,3 +1,3 @@
-from allennlp.common.lazy import Lazy
-
+from trapper.common.lazy import Lazy
+from trapper.common.params import Params
 from trapper.common.registrable import Registrable
diff --git a/trapper/common/lazy.py b/trapper/common/lazy.py
new file mode 100644
index 0000000..a28703f
--- /dev/null
+++ b/trapper/common/lazy.py
@@ -0,0 +1,9 @@
+"""
+`Lazy` class is adapted from `allennlp` for use when constructing objects using
+`FromParams`, when an argument to a constructor has a sequential dependency with
+another argument to the same constructor. See :py:class:`allennlp.common.lazy.Lazy`
+for further details.
+"""
+from allennlp.common.lazy import Lazy as _Lazy
+
+Lazy = _Lazy
diff --git a/trapper/data/__init__.py b/trapper/data/__init__.py
index 1715fbc..d0a1fa0 100644
--- a/trapper/data/__init__.py
+++ b/trapper/data/__init__.py
@@ -1,12 +1,10 @@
-from trapper.data.data_collators import TransformerDataCollator
+from trapper.data.data_adapters import DataAdapter, DataAdapterForQuestionAnswering
 from trapper.data.data_processors import (
+    DataProcessor,
     SquadDataProcessor,
     SquadQuestionAnsweringDataProcessor,
-    TransformerDataProcessor,
-)
-from trapper.data.data_processors.data_processor import (
-    IndexedDataset,
-    IndexedInstance,
 )
+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
diff --git a/trapper/data/data_adapters/__init__.py b/trapper/data/data_adapters/__init__.py
new file mode 100644
index 0000000..8d79f8a
--- /dev/null
+++ b/trapper/data/data_adapters/__init__.py
@@ -0,0 +1,4 @@
+from trapper.data.data_adapters.data_adapter import DataAdapter
+from trapper.data.data_adapters.question_answering_adapter import (
+    DataAdapterForQuestionAnswering,
+)
diff --git a/trapper/data/data_adapters/data_adapter.py b/trapper/data/data_adapters/data_adapter.py
new file mode 100644
index 0000000..111abf9
--- /dev/null
+++ b/trapper/data/data_adapters/data_adapter.py
@@ -0,0 +1,43 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Union
+
+from torch import Tensor
+
+from trapper.common import Registrable
+from trapper.data.data_processors.data_processor import IndexedInstance
+from trapper.data.tokenizers.tokenizer import TransformerTokenizer
+
+InputBatch = Dict[str, List[Union[int, List[int]]]]
+InputBatchTensor = Dict[str, Tensor]
+
+
+class DataAdapter(ABC, Registrable):
+    """
+    This callable class is responsible from converting the data instances that
+    are already tokenized and indexed into a format suitable for feeding into a
+    transformer model. Typically, it receives its inputs from a `DataProcessor`
+    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
+    `DataAdapterForQuestionAnswering` for an example.
+
+    Args:
+        tokenizer (): Required to access the ids of special tokens
+    """
+
+    def __init__(self, tokenizer: TransformerTokenizer):
+        self._tokenizer = tokenizer
+
+    @abstractmethod
+    def __call__(self, instance: IndexedInstance) -> IndexedInstance:
+        """
+        Takes a raw `IndexedInstance`, performs some processing on it, and returns
+        an `IndexedInstance` again. Look at
+        `DataAdapterForQuestionAnswering.__call__` for an example.
+
+        Args:
+            instance ():
+        """
+        raise NotImplementedError
diff --git a/trapper/data/data_collators/squad/question_answering_collator.py b/trapper/data/data_adapters/question_answering_adapter.py
similarity index 80%
rename from trapper/data/data_collators/squad/question_answering_collator.py
rename to trapper/data/data_adapters/question_answering_adapter.py
index e0e30b2..0b2b497 100644
--- a/trapper/data/data_collators/squad/question_answering_collator.py
+++ b/trapper/data/data_adapters/question_answering_adapter.py
@@ -1,19 +1,18 @@
-from typing import List, Tuple
+from typing import List
 
-from trapper.common.constants import CONTEXT_TOKEN
-from trapper.data.data_collators import TransformerDataCollator
+from trapper.data.data_adapters.data_adapter import DataAdapter
 from trapper.data.data_processors import IndexedInstance
 from trapper.data.tokenizers.tokenizer import TransformerTokenizer
 
 
-@TransformerDataCollator.register("question-answering")
-class DataCollatorForQuestionAnswering(TransformerDataCollator):
+@DataAdapter.register("question-answering")
+class DataAdapterForQuestionAnswering(DataAdapter):
     """
-    `DataCollatorForQuestionAnswering` can be used in SQuAD style question
+    `DataAdapterForQuestionAnswering` can be used in SQuAD style question
     answering tasks that involves a context, question and answer.
+
     Args:
-        tokenizer ():
-        model_input_keys ():
+        tokenizer (): Required to access the ids BOS and EOS tokens
     """
 
     CONTEXT_TOKEN_TYPE_ID = 0
@@ -22,16 +21,12 @@ class DataCollatorForQuestionAnswering(TransformerDataCollator):
     def __init__(
         self,
         tokenizer: TransformerTokenizer,
-        model_input_keys: Tuple[str, ...],
     ):
-        super().__init__(tokenizer, model_input_keys)
-        (self.context_token_id,) = self._tokenizer.convert_tokens_to_ids(
-            [CONTEXT_TOKEN]
-        )
-        self._eos_token_id = self._tokenizer.eos_token_id
+        super().__init__(tokenizer)
         self._bos_token_id = self._tokenizer.bos_token_id
+        self._eos_token_id = self._tokenizer.eos_token_id
 
-    def _build_input_fields(self, raw_instance: IndexedInstance) -> IndexedInstance:
+    def __call__(self, raw_instance: IndexedInstance) -> IndexedInstance:
         """
         Create a sequence with the following fields:
         input_ids: <bos> ...context_toks... <eos> ...question_toks... <eos>
@@ -43,6 +38,21 @@ def _build_input_fields(self, raw_instance: IndexedInstance) -> IndexedInstance:
         self._append_ending_token(instance)
         return instance
 
+    def _build_context(self, raw_instance: IndexedInstance) -> IndexedInstance:
+        context_tokens = raw_instance["context"]
+        input_ids = [self._bos_token_id] + context_tokens
+        token_type_ids = self._context_token_type_ids(context_tokens)
+        instance = {"input_ids": input_ids, "token_type_ids": token_type_ids}
+        self._handle_answer_span(instance, raw_instance)
+        return instance
+
+    def _append_separator_token(self, instance: IndexedInstance):
+        self._extend_token_ids(
+            instance=instance,
+            token_type_id=self.CONTEXT_TOKEN_TYPE_ID,
+            input_ids=[self._eos_token_id],
+        )
+
     def _append_question_tokens(
         self, instance: IndexedInstance, raw_instance: IndexedInstance
     ):
@@ -59,20 +69,13 @@ def _append_ending_token(self, instance: IndexedInstance):
             input_ids=[self._eos_token_id],
         )
 
-    def _append_separator_token(self, instance: IndexedInstance):
-        self._extend_token_ids(
-            instance=instance,
-            token_type_id=self.CONTEXT_TOKEN_TYPE_ID,
-            input_ids=[self._eos_token_id],
-        )
-
-    def _build_context(self, raw_instance: IndexedInstance) -> IndexedInstance:
-        context_tokens = raw_instance["context"]
-        input_ids = [self._bos_token_id] + context_tokens
-        token_type_ids = self._context_token_type_ids(context_tokens)
-        instance = {"input_ids": input_ids, "token_type_ids": token_type_ids}
-        self._handle_answer_span(instance, raw_instance)
-        return instance
+    @staticmethod
+    def _extend_token_ids(
+        instance: IndexedInstance, token_type_id: int, input_ids: List[int]
+    ):
+        instance["input_ids"].extend(input_ids)
+        token_type_ids = [token_type_id] * len(input_ids)
+        instance["token_type_ids"].extend(token_type_ids)
 
     def _context_token_type_ids(self, context_tokens: List[int]) -> List[int]:
         # handle segment encoding of the tokens inside the context
diff --git a/trapper/data/data_collators/data_collator.py b/trapper/data/data_collator.py
similarity index 76%
rename from trapper/data/data_collators/data_collator.py
rename to trapper/data/data_collator.py
index b984c78..46dccca 100644
--- a/trapper/data/data_collators/data_collator.py
+++ b/trapper/data/data_collator.py
@@ -1,4 +1,3 @@
-from abc import ABC, abstractmethod
 from collections import defaultdict
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
@@ -14,28 +13,29 @@
 InputBatchTensor = Dict[str, Tensor]
 
 
-class TransformerDataCollator(ABC, Registrable):
+class DataCollator(Registrable):
     """
     This class takes a batch of `IndexedInstance`s, typically generated using a
-    dataset reader, and returns an `InputBatchTensor`. It is meant to be
-    used as a callable just like the `DataCollator` in Transformers library. To
-    create and register your own, you must implement the abstract
-    `_build_input_fields` method. There are some convenience method implemented
-    for common operations. See `DataCollatorForQuestionAnswering` for an example.
+    dataset reader, and returns an `InputBatchTensor`. It is responsible from
+    padding the required instances, collect them into a batch and convert it into a
+    `InputBatchTensor`. Moreover, it also creates the attention_mask and inserts
+    it into the batch if required. It is used as a callable just like the
+    `DataCollator` in the `transformers` library.
+
     Args:
         tokenizer ():
-        model_input_keys ():
+        model_forward_params ():
     """
 
-    label_pad_token_id = PAD_TOKEN_LABEL_ID
+    default_implementation = "default"
 
     def __init__(
         self,
         tokenizer: TransformerTokenizer,
-        model_input_keys: Tuple[str, ...],
+        model_forward_params: Tuple[str, ...],
     ):
         self._tokenizer = tokenizer
-        self._model_input_keys: Tuple[str, ...] = model_input_keys
+        self._model_forward_params: Tuple[str, ...] = model_forward_params
 
     def __call__(
         self,
@@ -57,11 +57,10 @@ def build_model_inputs(
         should_eliminate_model_incompatible_keys: bool = True,
     ) -> InputBatch:
         return_attention_mask = (
-            return_attention_mask or "attention_mask" in self._model_input_keys
+            return_attention_mask or "attention_mask" in self._model_forward_params
         )
         batch = self._create_empty_batch()
         for instance in instances:
-            instance = self._build_input_fields(instance)
             if should_eliminate_model_incompatible_keys:
                 self._eliminate_model_incompatible_keys(instance)
             if return_attention_mask:
@@ -71,30 +70,17 @@ def build_model_inputs(
         self._eliminate_empty_inputs(batch)
         return batch
 
-    @abstractmethod
-    def _build_input_fields(self, raw_instance: IndexedInstance) -> IndexedInstance:
-        """
-        Takes a raw `IndexedInstance`, performs some processing on it,
-        and returns an `IndexedInstance` again. Look at
-        `DataCollatorForQuestionAnswering._build_input_fields` for an example.
-        Args:
-            raw_instance ():
-        """
-        raise NotImplementedError
+    @staticmethod
+    def _create_empty_batch() -> InputBatch:
+        return defaultdict(list)
 
     def _eliminate_model_incompatible_keys(self, instance: IndexedInstance):
         incompatible_keys = [
-            key for key in instance.keys() if key not in self._model_input_keys
+            key for key in instance.keys() if key not in self._model_forward_params
         ]
         for key in incompatible_keys:
             del instance[key]
 
-    @staticmethod
-    def _eliminate_empty_inputs(batch: InputBatch):
-        incompatible_keys = [key for key, val in batch.items() if len(val) == 0]
-        for key in incompatible_keys:
-            del batch[key]
-
     @staticmethod
     def _add_attention_mask(
         instance: IndexedInstance,
@@ -103,12 +89,15 @@ def _add_attention_mask(
             instance["attention_mask"] = [1] * len(instance["input_ids"])
 
     @staticmethod
-    def _extend_token_ids(
-        instance: IndexedInstance, token_type_id: int, input_ids: List[int]
-    ):
-        instance["input_ids"].extend(input_ids)
-        token_type_ids = [token_type_id] * len(input_ids)
-        instance["token_type_ids"].extend(token_type_ids)
+    def _add_instance(batch: InputBatch, instance: IndexedInstance):
+        for field_name, encodings in instance.items():
+            batch[field_name].append(encodings)
+
+    @staticmethod
+    def _eliminate_empty_inputs(batch: InputBatch):
+        incompatible_keys = [key for key, val in batch.items() if len(val) == 0]
+        for key in incompatible_keys:
+            del batch[key]
 
     def pad(
         self,
@@ -128,6 +117,21 @@ def pad(
                     padding_side=padding_side,
                 )
 
+    def _pad_id(self, padded_field: str) -> int:
+        if padded_field == "input_ids":
+            pad_id = self._tokenizer.pad_token_id
+        elif padded_field == "token_type_ids":
+            pad_id = self._tokenizer.pad_token_type_id
+        elif padded_field == "labels":
+            pad_id = PAD_TOKEN_LABEL_ID
+        elif padded_field == "attention_mask":
+            pad_id = 0
+        elif padded_field == "special_tokens_mask":
+            pad_id = 1
+        else:
+            raise ValueError(f"{padded_field} is not a valid field for padding")
+        return pad_id
+
     @staticmethod
     def _pad_encodings(
         encodings: List[List[int]],
@@ -145,33 +149,12 @@ def _pad_encodings(
             elif padding_side == "left":
                 encodings[i] = pad_values + encoded_inputs
 
-    def _pad_id(self, padded_field: str) -> int:
-        if padded_field == "input_ids":
-            pad_id = self._tokenizer.pad_token_id
-        elif padded_field == "token_type_ids":
-            pad_id = self._tokenizer.pad_token_type_id
-        elif padded_field == "labels":
-            pad_id = self.__class__.label_pad_token_id
-        elif padded_field == "attention_mask":
-            pad_id = 0
-        elif padded_field == "special_tokens_mask":
-            pad_id = 1
-        else:
-            raise ValueError(f"{padded_field} is not a valid field for padding")
-        return pad_id
-
-    @staticmethod
-    def _create_empty_batch() -> InputBatch:
-        return defaultdict(list)
-
-    @staticmethod
-    def _add_instance(batch: InputBatch, instance: IndexedInstance):
-        for field_name, encodings in instance.items():
-            batch[field_name].append(encodings)
-
     @staticmethod
     def _convert_to_tensor(batch: InputBatch) -> InputBatchTensor:
         return {
             input_key: torch.tensor(encodings)
             for input_key, encodings in batch.items()
         }
+
+
+DataCollator.register("default")(DataCollator)
diff --git a/trapper/data/data_collators/__init__.py b/trapper/data/data_collators/__init__.py
deleted file mode 100644
index e0366fe..0000000
--- a/trapper/data/data_collators/__init__.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from trapper.data.data_collators.data_collator import (
-    InputBatch,
-    InputBatchTensor,
-    TransformerDataCollator,
-)
-from trapper.data.data_collators.squad import DataCollatorForQuestionAnswering
diff --git a/trapper/data/data_collators/squad/__init__.py b/trapper/data/data_collators/squad/__init__.py
deleted file mode 100644
index 02be6a8..0000000
--- a/trapper/data/data_collators/squad/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from trapper.data.data_collators.squad.question_answering_collator import (
-    DataCollatorForQuestionAnswering,
-)
diff --git a/trapper/data/data_processors/__init__.py b/trapper/data/data_processors/__init__.py
index ea53335..1973b42 100644
--- a/trapper/data/data_processors/__init__.py
+++ b/trapper/data/data_processors/__init__.py
@@ -1,7 +1,6 @@
 from trapper.data.data_processors.data_processor import (
-    IndexedDataset,
+    DataProcessor,
     IndexedInstance,
-    TransformerDataProcessor,
 )
 from trapper.data.data_processors.squad import (
     SquadDataProcessor,
diff --git a/trapper/data/data_processors/data_processor.py b/trapper/data/data_processors/data_processor.py
index 42fbfd6..0269a97 100644
--- a/trapper/data/data_processors/data_processor.py
+++ b/trapper/data/data_processors/data_processor.py
@@ -1,9 +1,6 @@
-import inspect
 import logging
 from abc import ABCMeta, abstractmethod
-from typing import Any, Dict, Iterator, List, Optional, Union
-
-from torch.utils.data import Dataset
+from typing import Any, Dict, List, Optional, Union
 
 from trapper.common import Registrable
 from trapper.common.constants import PositionTuple
@@ -22,31 +19,19 @@ class ImproperDataInstanceError(Exception):
      artifacts"""
 
 
-class IndexedDataset(Dataset):
-    def __init__(self, instances: List[IndexedInstance]):
-        self.instances = instances
-
-    def __getitem__(self, index) -> IndexedInstance:
-        return self.instances[index]
-
-    def __len__(self):
-        return len(self.instances)
-
-    def __iter__(self) -> Iterator[IndexedInstance]:
-        yield from self.instances
-
-
-class TransformerDataProcessor(Registrable, metaclass=ABCMeta):
+class DataProcessor(Registrable, metaclass=ABCMeta):
     """
-    This class is used for converting a raw instance dict from `datasets.Dataset`
-    to `IndexedInstance`. The abstract `text_to_instance` and `process` must be
-    implemented in the subclasses. Typically, the `process` method calls
-    `text_to_instance` with raw data as input to generate an `IndexedInstance`.
-    Some methods that are commonly used are implemented here for convenience.
-
-    Child classes have to set the following class variables:
+    A callable class used for converting a raw instance dict from `datasets.Dataset`
+    to `IndexedInstance`. Typically, used as the first processing step after the
+    raw data is read to extract the task-related fields from the raw data.
+    The abstract `text_to_instance` and `__call__` must be implemented in the
+    subclasses. Typically, the `__call__` method calls `text_to_instance` with
+    raw data as input to generate an `IndexedInstance`. Some methods that are
+    commonly used are implemented here for convenience.
+
+    Child classes may need to set the following class variables:
         - NUM_EXTRA_SPECIAL_TOKENS_IN_SEQUENCE : The total number of extra special
-            tokens (not unique) in the input ids.
+            tokens (do not have to be unique) in the input ids.
 
     Args:
         tokenizer ():
@@ -65,7 +50,7 @@ def tokenizer(self):
     def text_to_instance(self, *inputs) -> IndexedInstance:
         """
         Takes unpacked, raw input and converts them to an `IndexedInstance`.
-        Typically, called by the `process` method. An important suggestion while
+        Typically, invoked by the `__call__` method. An important suggestion while
         implementing this method in your custom subclass is to put the label input
         at the end as an optional parameter whose default value is `None`. By this
         way, you can reuse this method while reading a single instance during the
@@ -78,11 +63,14 @@ def text_to_instance(self, *inputs) -> IndexedInstance:
         raise NotImplementedError
 
     @abstractmethod
-    def process(self, instance_dict: Dict[str, Any]) -> Optional[IndexedInstance]:
+    def __call__(self, instance_dict: Dict[str, Any]) -> IndexedInstance:
         """
-        Processes an instance dict taken from a `datasets.Dataset`. Returns an
-        `IndexedInstance` if the input is successfully tokenized, indexed and
-        arranged. Otherwise, returns None.
+        Processes an instance dict taken from a `datasets.Dataset`. Typically,
+        extracts the task-related fields and pass them to `text_to_instance` method.
+        Returns an`IndexedInstance` with proper keys if the input is successfully
+        tokenized, indexed and arranged, otherwise returns a dummy
+        `IndexedInstance` with "filter_out"=True and the remaining keys are
+        associated with empty values suitable with the expected types.
 
         Args:
             instance_dict ():
@@ -91,6 +79,13 @@ def process(self, instance_dict: Dict[str, Any]) -> Optional[IndexedInstance]:
 
     @classmethod
     def _total_seq_len(cls, *sequences):
+        """
+        Computes the total number of tokens in an iterable of sequences by taking
+        the special tokens in the combined sequence into account as well.
+
+        Args:
+            *sequences ():
+        """
         total_seq_len = sum(len(seq) for seq in sequences)
         total_seq_len += cls.NUM_EXTRA_SPECIAL_TOKENS_IN_SEQUENCE
         return total_seq_len
diff --git a/trapper/data/data_processors/squad/question_answering_processor.py b/trapper/data/data_processors/squad/question_answering_processor.py
index 1f52dd8..1430d26 100644
--- a/trapper/data/data_processors/squad/question_answering_processor.py
+++ b/trapper/data/data_processors/squad/question_answering_processor.py
@@ -1,9 +1,9 @@
 import logging
 from typing import Any, Dict, Optional
 
-from trapper.common.constants import SpanTuple
+from trapper.common.constants import PositionTuple, SpanTuple
 from trapper.common.utils import convert_span_dict_to_tuple
-from trapper.data.data_processors import TransformerDataProcessor
+from trapper.data.data_processors import DataProcessor
 from trapper.data.data_processors.data_processor import (
     ImproperDataInstanceError,
     IndexedInstance,
@@ -13,11 +13,46 @@
 logger = logging.getLogger(__file__)
 
 
-@TransformerDataProcessor.register("squad-question-answering")
+@DataProcessor.register("squad-question-answering")
 class SquadQuestionAnsweringDataProcessor(SquadDataProcessor):
     NUM_EXTRA_SPECIAL_TOKENS_IN_SEQUENCE = 3  # <bos> context <eos> question <eos>
     MAX_SEQUENCE_LEN = 512
 
+    def __call__(self, instance_dict: Dict[str, Any]) -> Optional[IndexedInstance]:
+        paragraph_ind = instance_dict["id"]
+        context = instance_dict["context"]
+        question = {"text": instance_dict["question"], "start": None}
+        question = convert_span_dict_to_tuple(question)
+        if self._is_input_too_long(context, question):
+            return self.filtered_instance()
+        # Rename SQuAD answer_start as start for trapper tuple conversion.
+        answers = instance_dict["answers"]
+        first_answer = {
+            "start": answers["answer_start"][0],
+            "text": answers["text"][0],
+        }
+        first_answer = convert_span_dict_to_tuple(first_answer)
+        try:
+            return self.text_to_instance(
+                context=context,
+                question=question,
+                paragraph_ind=paragraph_ind,
+                answer=first_answer,
+            )
+        except ImproperDataInstanceError:
+            return self.filtered_instance()
+
+    @staticmethod
+    def filtered_instance() -> IndexedInstance:
+        return {
+            "answer": [],
+            "answer_position_tokenized": PositionTuple(start=-1, end=-1),
+            "context": [],
+            "context_index": "",
+            "question": [],
+            "discard_sample": True,
+        }
+
     def text_to_instance(
         self,
         context: str,
@@ -34,10 +69,6 @@ def text_to_instance(
             "context": self._tokenizer.convert_tokens_to_ids(tokenized_context),
             "question": self._tokenizer.convert_tokens_to_ids(tokenized_question),
         }
-        if question.start is not None:
-            self._store_tokenized_field_position(
-                instance, context, question.start, type_="question"
-            )
 
         if answer is not None:
             answer = self._join_whitespace_prefix(context, answer)
@@ -53,30 +84,6 @@ def text_to_instance(
         instance["context_index"] = paragraph_ind
         return instance
 
-    def process(self, instance_dict: Dict[str, Any]) -> Optional[IndexedInstance]:
-        paragraph_ind = instance_dict["paragraph_ind"]
-        context = instance_dict["context"]
-        question = {"text": instance_dict["question"], "start": None}
-        question = convert_span_dict_to_tuple(question)
-        if self._is_input_too_long(context, question):
-            return None
-        # Rename SQuAD answer_start as start for trapper tuple conversion.
-        answers = instance_dict["answers"]
-        first_answer = {
-            "start": answers["answer_start"][0],
-            "text": answers["text"][0],
-        }
-        first_answer = convert_span_dict_to_tuple(first_answer)
-        try:
-            return self.text_to_instance(
-                context=context,
-                question=question,
-                paragraph_ind=paragraph_ind,
-                answer=first_answer,
-            )
-        except ImproperDataInstanceError:
-            return None
-
     def _is_input_too_long(self, context: str, question: SpanTuple) -> bool:
         context_tokens = self.tokenizer.tokenize(context)
         question_tokens = self.tokenizer.tokenize(question.text)
diff --git a/trapper/data/data_processors/squad/squad_processor.py b/trapper/data/data_processors/squad/squad_processor.py
index 246027e..f23de4a 100644
--- a/trapper/data/data_processors/squad/squad_processor.py
+++ b/trapper/data/data_processors/squad/squad_processor.py
@@ -2,10 +2,16 @@
 from typing import Dict, List, Optional, Union
 
 from trapper.common.constants import PositionTuple, SpanTuple
-from trapper.data.data_processors.data_processor import TransformerDataProcessor
+from trapper.data.data_processors.data_processor import DataProcessor
 
 
-class SquadDataProcessor(TransformerDataProcessor, metaclass=ABCMeta):
+class SquadDataProcessor(DataProcessor, metaclass=ABCMeta):
+    """
+    Provides utility methods that can be used in SQuAD style tasks involving a
+    context and information fields inside the context e.g. answers in the case of
+    question answering.
+    """
+
     @staticmethod
     def _join_whitespace_prefix(context: str, field: SpanTuple) -> SpanTuple:
         """
diff --git a/trapper/data/dataset_loader.py b/trapper/data/dataset_loader.py
index 1b51130..2f3526f 100644
--- a/trapper/data/dataset_loader.py
+++ b/trapper/data/dataset_loader.py
@@ -1,140 +1,75 @@
-import inspect
 import logging
 from pathlib import Path
-from typing import Iterable, Mapping, Optional, Sequence, Union
+from typing import Union
 
 import datasets
-from allennlp.common.util import ensure_list
-from datasets import (
-    DownloadConfig,
-    Features,
-    GenerateMode,
-    Split,
-    Version,
-    load_dataset,
-)
-from datasets.tasks import TaskTemplate
 
 from trapper.common import Registrable
-from trapper.data.data_processors import IndexedDataset, IndexedInstance
-from trapper.data.data_processors.data_processor import TransformerDataProcessor
+from trapper.data import DataAdapter
+from trapper.data.data_processors.data_processor import DataProcessor
+from trapper.data.dataset_reader import DatasetReader
 
 logger = logging.getLogger(__file__)
 
 
 class DatasetLoader(Registrable):
     """
-    This class is used to read a dataset from the `datasets` library and return
-    a collection of `IndexedDataset`s. Its first argument is a
-    `TransformerDataProcessor` which is required for tokenization and handling
-    the special tokens during the pre-processing. The remaining arguments are
-    directly transferred to the `datasets.load_dataset` function and use for
-    loading the raw dataset from the `datasets` library. See
-    :py:func:`datasets.load_dataset` for the details of these parameters and how
-    the dataset is loaded.
+    This class is responsible for reading and pre-processing a dataset. This
+    involves reading the raw data instances, extracting the task-related fields,
+    tokenizing the instances, converting the fields into a format accepted by the
+    transformer models as well as taking care of the special tokens. All these
+    tasks are performed sequentially by three components in a pipelined manner.
 
     Args:
-        data_processor ():
-        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 ():
+        dataset_reader (): Reads the raw dataset.
+        data_processor (): Handles pre-processing i.e. tokenization and adding
+            the special tokens.
+        data_adapter (): Converts the instance into a `IndexedInstance` suitable
+            for directly feeding to the models.
     """
 
     default_implementation = "default"
 
     def __init__(
         self,
-        data_processor: TransformerDataProcessor,
-        path: str,
-        name: Optional[str] = None,
-        data_dir: Optional[str] = None,
-        data_files: Optional[
-            Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]
-        ] = None,
-        split: Optional[Union[str, Split]] = None,
-        cache_dir: Optional[str] = None,
-        features: Optional[Features] = None,
-        download_config: Optional[DownloadConfig] = None,
-        download_mode: Optional[GenerateMode] = None,
-        ignore_verifications: bool = False,
-        keep_in_memory: Optional[bool] = None,
-        save_infos: bool = False,
-        script_version: Optional[Union[str, Version]] = None,
-        use_auth_token: Optional[Union[bool, str]] = None,
-        task: Optional[Union[str, TaskTemplate]] = None,
-        streaming: bool = False,
-        **config_kwargs,
+        dataset_reader: DatasetReader,
+        data_processor: DataProcessor,
+        data_adapter: DataAdapter,
     ):
+        self._dataset_reader = dataset_reader
         self._data_processor = data_processor
-        locals_ = locals()
-        self._dataset = load_dataset(
-            *[locals_[arg] for arg in inspect.getfullargspec(load_dataset).args],
-            **config_kwargs,
-        )
+        self._data_adapter = data_adapter
+
+    @property
+    def dataset_reader(self):
+        return self._dataset_reader
 
     @property
     def data_processor(self):
         return self._data_processor
 
     @property
-    def split_names(self):
-        return tuple(
-            s for s in ("train", "validation", "test") if s in self._dataset
-        )
+    def data_adapter(self):
+        return self._data_adapter
 
-    def load(self, split_name: Union[Path, str]) -> IndexedDataset:
+    def load(self, split_name: Union[Path, str]) -> datasets.Dataset:
         """
-        Reads the dataset for the specified split.
+        Reads the specified split of the dataset, process the instances and
+        covert them into a format suitable for feeding to a model.
 
         Args:
             split_name (): one of "train", "validation" or "test"
 
         Returns:
-            an `IndexedDataset` that can be passed to `TransformerTrainer`
-        """
-        # TODO: Check if `ensure_list` is really necessary?
-        instances = ensure_list(self._load(self._get_raw_data(split_name)))
-        return IndexedDataset(instances)
-
-    def _load(self, split: datasets.Dataset) -> Iterable[IndexedInstance]:
-        """
-        Returns an `Iterable` of `IndexedInstance`s from a dataset split.
-
-        Args:
-            split (): The train, validation or test split of a dataset
-
-        Returns:
-            an iterable of `IndexedInstance`s
+            a processed split from the dataset, which can be passed to
+                `TransformerTrainer`
         """
-        for instance_dict in split:
-            indexed_instance = self.data_processor.process(instance_dict)
-            if indexed_instance is not None:
-                yield indexed_instance
-
-    def _get_raw_data(self, split_name: Union[Path, str]):
-        self._check_split_name(split_name)
-        return self._dataset[split_name]
-
-    def _check_split_name(self, split_name):
-        if split_name not in self.split_names:
-            raise ValueError(
-                f"split: {split_name} not in available splits for the "
-                f"dataset ({self.split_names})"
-            )
+        raw_data = self.dataset_reader.get_dataset(split_name)
+        processed_data = raw_data.map(self.data_processor)
+        processed_data = processed_data.filter(
+            lambda x: not x.get("discard_sample", False)
+        )
+        return processed_data.map(self.data_adapter)
 
 
 DatasetLoader.register("default")(DatasetLoader)
diff --git a/trapper/data/dataset_reader.py b/trapper/data/dataset_reader.py
new file mode 100644
index 0000000..41bdf52
--- /dev/null
+++ b/trapper/data/dataset_reader.py
@@ -0,0 +1,106 @@
+import inspect
+import logging
+from pathlib import Path
+from typing import Mapping, Optional, Sequence, Union
+
+from datasets import (
+    DownloadConfig,
+    Features,
+    GenerateMode,
+    Split,
+    Version,
+    load_dataset,
+)
+from datasets.tasks import TaskTemplate
+
+from trapper.common import Registrable
+
+logger = logging.getLogger(__file__)
+
+
+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.
+
+    .. _datasets.load_dataset: https://github.com/huggingface/datasets/blob/b057846bcc1ccfe0fda4d8d42f190d146a70ee64/src/datasets/load.py#L996
+
+    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 ():
+    """
+
+    default_implementation = "default"
+
+    def __init__(
+        self,
+        path: str,
+        name: Optional[str] = None,
+        data_dir: Optional[str] = None,
+        data_files: Optional[
+            Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]
+        ] = None,
+        split: Optional[Union[str, Split]] = None,
+        cache_dir: Optional[str] = None,
+        features: Optional[Features] = None,
+        download_config: Optional[DownloadConfig] = None,
+        download_mode: Optional[GenerateMode] = None,
+        ignore_verifications: bool = False,
+        keep_in_memory: Optional[bool] = None,
+        save_infos: bool = False,
+        script_version: Optional[Union[str, Version]] = None,
+        use_auth_token: Optional[Union[bool, str]] = None,
+        task: Optional[Union[str, TaskTemplate]] = None,
+        streaming: bool = False,
+        **config_kwargs,
+    ):
+        locals_ = locals()
+        self._dataset = load_dataset(
+            *[locals_[arg] for arg in inspect.getfullargspec(load_dataset).args],
+            **config_kwargs,
+        )
+
+    @property
+    def split_names(self):
+        return tuple(
+            s for s in ("train", "validation", "test") if s in self._dataset
+        )
+
+    def get_dataset(self, split_name: Union[Path, str]):
+        """
+        Returns the specified split of the dataset.
+
+        Args:
+            split_name (): one of "train", "validation" or "test"
+        """
+        self._check_split_name(split_name)
+        return self._dataset[split_name]
+
+    def _check_split_name(self, split_name):
+        if split_name not in self.split_names:
+            raise ValueError(
+                f"split: {split_name} not in available splits for the "
+                f"dataset ({self.split_names})"
+            )
+
+
+DatasetReader.register("default")(DatasetReader)
diff --git a/trapper/models/model.py b/trapper/models/model.py
index 5154016..825911a 100644
--- a/trapper/models/model.py
+++ b/trapper/models/model.py
@@ -33,7 +33,7 @@ class TransformerModel(PreTrainedModel, Registrable):
 
     Then, the legitimate parameters are stored as `forward_params` attribute of
     the returned model. Thanks to that, `trapper.training.TransformerTrainer`
-    can access that attribute to use as argument for the `model_input_keys`
+    can access that attribute to use as argument for the `model_forward_params`
     parameter while creating the dataset collator.
 
     Below are the explanations of the important class variables.
diff --git a/trapper/pipelines/pipeline.py b/trapper/pipelines/pipeline.py
index 0207741..604b308 100644
--- a/trapper/pipelines/pipeline.py
+++ b/trapper/pipelines/pipeline.py
@@ -5,11 +5,8 @@
 from transformers.pipelines import pipeline
 
 from trapper.common.params import Params
-from trapper.data import (
-    DatasetLoader,
-    TransformerDataCollator,
-    TransformerTokenizer,
-)
+from trapper.data import DataAdapter, DataProcessor, TransformerTokenizer
+from trapper.data.data_collator import DataCollator
 from trapper.models import TransformerModel
 
 
@@ -28,8 +25,9 @@ def create_pipeline_from_checkpoint(
 def _create_pipeline(checkpoint_path, params, task: str, **kwargs):
     model = _create_model(checkpoint_path, params)
     tokenizer = _create_tokenizer(checkpoint_path, params)
-    dataset_reader = _create_dataset_reader(params, tokenizer)
-    data_collator = _create_data_collator(model, params, tokenizer)
+    data_processor = _create_data_processor(params, tokenizer)
+    data_adapter = _create_data_adapter(params, tokenizer)
+    data_collator = _create_data_collator(model, tokenizer)
     config = AutoConfig.from_pretrained(checkpoint_path)
     pipeline_ = pipeline(
         task=task,
@@ -37,16 +35,18 @@ def _create_pipeline(checkpoint_path, params, task: str, **kwargs):
         tokenizer=tokenizer,
         config=config,
         framework="pt",
-        dataset_reader=dataset_reader,
+        data_processor=data_processor,
+        data_adapter=data_adapter,
         data_collator=data_collator,
         **kwargs
     )
     return pipeline_
 
 
-def _create_data_collator(model, params, tokenizer):
-    sub_cls = TransformerDataCollator.by_name(params["data_collator"]["type"])
-    return sub_cls(tokenizer=tokenizer, model_input_keys=model.forward_params)
+def _create_data_collator(model, tokenizer):
+    return DataCollator(
+        tokenizer=tokenizer, model_forward_params=model.forward_params
+    )
 
 
 def _validate_checkpoint_dir(path: Union[str, Path]):
@@ -77,6 +77,13 @@ def _create_tokenizer(checkpoint_path, params):
     )
 
 
-def _create_dataset_reader(params, tokenizer):
-    sub_cls = DatasetLoader.by_name(params["dataset_reader"]["type"])
-    return sub_cls(tokenizer)
+def _create_data_processor(params, tokenizer):
+    return DataProcessor.by_name(
+        params["dataset_loader"]["data_processor"]["type"]
+    )(tokenizer)
+
+
+def _create_data_adapter(params, tokenizer):
+    return DataAdapter.by_name(params["dataset_loader"]["data_adapter"]["type"])(
+        tokenizer
+    )
diff --git a/trapper/pipelines/question_answering_pipeline.py b/trapper/pipelines/question_answering_pipeline.py
index 5e5b087..c5d90a5 100644
--- a/trapper/pipelines/question_answering_pipeline.py
+++ b/trapper/pipelines/question_answering_pipeline.py
@@ -45,7 +45,10 @@
     convert_span_tuple_to_dict,
 )
 from trapper.data import IndexedInstance, SquadQuestionAnsweringDataProcessor
-from trapper.data.data_collators import DataCollatorForQuestionAnswering
+from trapper.data.data_adapters.question_answering_adapter import (
+    DataAdapterForQuestionAnswering,
+)
+from trapper.data.data_collator import DataCollator
 from trapper.models import TransformerModel
 
 
@@ -160,7 +163,8 @@ def __init__(
         model: PreTrainedModel,
         tokenizer: PreTrainedTokenizer,
         data_processor: SquadQuestionAnsweringDataProcessor,
-        data_collator: DataCollatorForQuestionAnswering,
+        data_adapter: DataAdapterForQuestionAnswering,
+        data_collator: DataCollator,
         modelcard: Optional[ModelCard] = None,
         framework: Optional[str] = None,
         device: int = -1,
@@ -180,6 +184,7 @@ def __init__(
         self._args_parser = QuestionAnsweringArgumentHandler()
         self.check_model_type(MODEL_FOR_QUESTION_ANSWERING_MAPPING)
         self._data_processor = data_processor
+        self._data_adapter = data_adapter
         self._data_collator = data_collator
 
     def __call__(self, *args, **kwargs):
@@ -226,6 +231,7 @@ def __call__(self, *args, **kwargs):
         all_answers = []
         for example in tqdm(examples, disable=kwargs["disable_tqdm"]):
             indexed_instance = self._data_processor.text_to_instance(**example)
+            indexed_instance = self._data_adapter(indexed_instance)
             # Manage tensor allocation on correct device
             with self.device_placement():
                 with torch.no_grad():
diff --git a/trapper/training/callbacks.py b/trapper/training/callbacks.py
index 0d6a741..81cd98d 100644
--- a/trapper/training/callbacks.py
+++ b/trapper/training/callbacks.py
@@ -1,9 +1,9 @@
-from transformers import TrainerCallback as _TransformersTrainerCallback
+from transformers import TrainerCallback as _TrainerCallback
 
 from trapper.common import Registrable
 
 
-class TrainerCallback(_TransformersTrainerCallback, Registrable):
+class TrainerCallback(_TrainerCallback, Registrable):
     """
     The base class that implements `Registrable` for
     `transformers.TrainerCallback`. To add a custom callback, just subclass
diff --git a/trapper/training/optimizers.py b/trapper/training/optimizers.py
index 022301b..dd07746 100644
--- a/trapper/training/optimizers.py
+++ b/trapper/training/optimizers.py
@@ -7,7 +7,7 @@
 @append_parent_docstr(parent_id=1)
 class Optimizer(Registrable, _allennlp_optimizers.Optimizer):
     """
-    The base registrable optimizer class that replaces the one from the
+    The base `Registrable` optimizer class that replaces the one from the
     `allennlp` library.
     """
 
diff --git a/trapper/training/train.py b/trapper/training/train.py
index 108427c..6d02574 100644
--- a/trapper/training/train.py
+++ b/trapper/training/train.py
@@ -7,7 +7,7 @@
 import transformers
 from transformers.trainer_utils import is_main_process
 
-from trapper.common.params import Params
+from trapper.common import Params
 from trapper.training import TransformerTrainer
 
 logger = logging.getLogger(__name__)
diff --git a/trapper/training/trainers/trainer.py b/trapper/training/trainers/trainer.py
index 493c484..178b96f 100644
--- a/trapper/training/trainers/trainer.py
+++ b/trapper/training/trainers/trainer.py
@@ -1,6 +1,7 @@
 import re
 from typing import Callable, Dict, List, Optional, Tuple
 
+import datasets
 import torch
 from torch.optim.lr_scheduler import LambdaLR
 from transformers import PreTrainedModel
@@ -10,13 +11,8 @@
 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 (
-    DatasetLoader,
-    IndexedDataset,
-    TransformerDataCollator,
-    TransformerDataProcessor,
-    TransformerTokenizer,
-)
+from trapper.data import DatasetLoader, TransformerTokenizer
+from trapper.data.data_collator import DataCollator
 from trapper.models import TransformerModel
 from trapper.training.callbacks import TrainerCallback
 from trapper.training.metrics import TransformerMetric
@@ -37,9 +33,9 @@ def __init__(
         self,
         model: PreTrainedModel = None,
         args: TransformerTrainingArguments = None,
-        data_collator: Optional[TransformerDataCollator] = None,
-        train_dataset: Optional[IndexedDataset] = None,
-        eval_dataset: Optional[IndexedDataset] = None,
+        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,
         compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
@@ -67,9 +63,8 @@ def from_partial_objects(
         dev_split_name: str,
         model: Lazy[TransformerModel],
         tokenizer: Lazy[TransformerTokenizer],
-        data_processor: Lazy[TransformerDataProcessor],
         dataset_loader: Lazy[DatasetLoader],
-        data_collator: Lazy[TransformerDataCollator],
+        data_collator: Lazy[DataCollator],
         optimizer: Lazy[Optimizer],
         compute_metrics: Optional[Lazy[TransformerMetric]] = None,
         no_grad: List[str] = None,
@@ -92,18 +87,17 @@ def from_partial_objects(
             [n, p] for n, p in model_.named_parameters() if p.requires_grad
         ]
         optimizer_ = optimizer.construct(model_parameters=params_with_grad)
-
+        dataset_loader_ = dataset_loader.construct(
+            tokenizer=tokenizer_, 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_input_keys=model_.forward_params
+            tokenizer=tokenizer_, model_forward_params=model_.forward_params
         )
         compute_metrics_ = cls._create_compute_metrics(
             compute_metrics, data_collator_
         )
-        data_processor_ = data_processor.construct(tokenizer=tokenizer_)
-        dataset_loader_ = dataset_loader.construct(data_processor=data_processor_)
-        train_dataset_ = dataset_loader_.load(train_split_name)
-        eval_dataset_ = dataset_loader_.load(dev_split_name)
-
         return cls(
             model=model_,
             args=args,
@@ -120,7 +114,7 @@ def from_partial_objects(
     def _create_compute_metrics(
         cls,
         compute_metrics: Optional[Lazy[TransformerMetric]],
-        data_collator: TransformerDataCollator,
+        data_collator: DataCollator,
     ) -> Optional[TransformerMetric]:
         if compute_metrics is None:
             return None


Test patch:
diff --git a/tests/trapper/data/data_collators/squad/test_question_answering_collator.py b/tests/trapper/data/data_collators/squad/test_question_answering_collator.py
index f376002..a501355 100644
--- a/tests/trapper/data/data_collators/squad/test_question_answering_collator.py
+++ b/tests/trapper/data/data_collators/squad/test_question_answering_collator.py
@@ -6,8 +6,11 @@
 from torch.utils.data.sampler import SequentialSampler
 from transformers.trainer_pt_utils import SequentialDistributedSampler
 
-from trapper.data import SquadQuestionAnsweringDataProcessor
-from trapper.data.data_collators import DataCollatorForQuestionAnswering, InputBatch
+from trapper.data import DatasetReader, SquadQuestionAnsweringDataProcessor
+from trapper.data.data_adapters.question_answering_adapter import (
+    DataAdapterForQuestionAnswering,
+)
+from trapper.data.data_collator import DataCollator, InputBatch
 from trapper.data.dataset_loader import DatasetLoader
 from trapper.data.tokenizers import QuestionAnsweringTokenizer
 from trapper.models.auto_wrappers import _TASK_TO_INPUT_FIELDS
@@ -48,24 +51,34 @@ def data_processor(tokenizer):
 
 
 @pytest.fixture
-def dataset_loader(data_processor):
+def dataset_loader(tokenizer):
+    dataset_reader = DatasetReader(path="squad_qa_test_fixture")
+    data_processor = SquadQuestionAnsweringDataProcessor(tokenizer)
+    data_adapter = DataAdapterForQuestionAnswering(tokenizer)
     return DatasetLoader(
-        data_processor=data_processor, path="squad_qa_test_fixture"
+        dataset_reader=dataset_reader,
+        data_processor=data_processor,
+        data_adapter=data_adapter,
     )
 
 
+def adapt_instances(instances, adapter):
+    return [adapter(i) for i in instances]
+
+
 @pytest.fixture
-def fixtures_dir(fixtures_root):
-    return fixtures_root / "data/question_answering"
+def raw_dev_dataset(dataset_loader, data_processor):
+    raw_dataset = dataset_loader.dataset_reader.get_dataset("validation")
+    return [dataset_loader.data_processor(i) for i in raw_dataset]
 
 
 @pytest.fixture
-def dev_dataset(dataset_loader, fixtures_dir):
+def adapted_dev_dataset(dataset_loader):
     return dataset_loader.load("validation")
 
 
 @pytest.fixture
-def train_dataset(dataset_loader, fixtures_dir):
+def adapted_train_dataset(dataset_loader):
     return dataset_loader.load("train")
 
 
@@ -76,7 +89,7 @@ def model_forward_params(args):
 
 @pytest.fixture
 def dataset_collator(tokenizer, model_forward_params):
-    return DataCollatorForQuestionAnswering(tokenizer, model_forward_params)
+    return DataCollator(tokenizer, model_forward_params)
 
 
 def get_sequential_sampler(distributed_training, dataset):
@@ -86,10 +99,10 @@ def get_sequential_sampler(distributed_training, dataset):
 
 
 @pytest.fixture
-def dev_data_items(dataset_collator, dev_dataset, args):
-    sampler = get_sequential_sampler(args.distributed, dev_dataset)
+def dev_data_items(dataset_collator, adapted_dev_dataset, args):
+    sampler = get_sequential_sampler(args.distributed, adapted_dev_dataset)
     loader = DataLoader(
-        dev_dataset,
+        adapted_dev_dataset,
         batch_size=args.dev_batch_size,
         sampler=sampler,
         collate_fn=dataset_collator,
@@ -98,10 +111,10 @@ def dev_data_items(dataset_collator, dev_dataset, args):
 
 
 @pytest.fixture
-def train_data_items(dataset_collator, train_dataset, args):
-    sampler = get_sequential_sampler(args.distributed, train_dataset)
+def train_data_items(dataset_collator, adapted_train_dataset, args):
+    sampler = get_sequential_sampler(args.distributed, adapted_train_dataset)
     loader = DataLoader(
-        train_dataset,
+        adapted_train_dataset,
         batch_size=args.train_batch_size,
         sampler=sampler,
         collate_fn=dataset_collator,
@@ -119,8 +132,8 @@ def test_data_sizes(dev_data_items, train_data_items):
 
 
 @pytest.fixture
-def collated_batch(dataset_collator, dev_dataset):
-    return dataset_collator.build_model_inputs(dev_dataset)
+def collated_batch(dataset_collator, adapted_dev_dataset):
+    return dataset_collator.build_model_inputs(adapted_dev_dataset)
 
 
 @pytest.mark.parametrize(
@@ -132,7 +145,13 @@ def collated_batch(dataset_collator, dev_dataset):
     ],
 )
 def test_batch_content(
-    args, tokenizer, dev_dataset, index, question, collated_batch
+    args,
+    tokenizer,
+    raw_dev_dataset,
+    adapted_dev_dataset,
+    index,
+    question,
+    collated_batch,
 ):
     if args.uncased:
         question = question.lower()
@@ -140,7 +159,7 @@ def test_batch_content(
         question, index, tokenizer, collated_batch
     )
 
-    raw_instance = dev_dataset[index]
+    raw_instance = raw_dev_dataset[index]
     token_type_ids = collated_batch["token_type_ids"][index]
     validate_token_type_ids(token_type_ids, raw_instance)
     validate_attention_mask(collated_batch)
diff --git a/tests/trapper/data/data_processors/squad/test_question_answering_processor.py b/tests/trapper/data/data_processors/squad/test_question_answering_processor.py
index af54cf7..06e2adc 100644
--- a/tests/trapper/data/data_processors/squad/test_question_answering_processor.py
+++ b/tests/trapper/data/data_processors/squad/test_question_answering_processor.py
@@ -3,17 +3,13 @@
 
 import pytest
 
-from trapper.data import DatasetLoader, SquadQuestionAnsweringDataProcessor
+from trapper.data import DatasetReader, SquadQuestionAnsweringDataProcessor
 from trapper.data.tokenizers import QuestionAnsweringTokenizer
 
 
 @dataclass
 class Arguments:
-    distributed: bool = False
-    shuffle: bool = False
     model_type: str = "roberta-base"
-    train_batch_size: int = 2
-    dev_batch_size: int = 1
 
 
 @pytest.fixture(scope="module")
@@ -26,31 +22,16 @@ def args():
     return args
 
 
-@pytest.fixture(scope="module")
-def tokenizer(args):
-    return QuestionAnsweringTokenizer.from_pretrained(args.model_type)
-
-
 @pytest.fixture
-def data_processor(tokenizer):
-    return SquadQuestionAnsweringDataProcessor(tokenizer)
+def dev_dataset():
+    dataset_reader = DatasetReader(path="squad_qa_test_fixture")
+    return dataset_reader.get_dataset("validation")
 
 
 @pytest.fixture
-def dataset_loader(data_processor):
-    return DatasetLoader(
-        data_processor=data_processor, path="squad_qa_test_fixture"
-    )
-
-
-@pytest.fixture
-def fixtures_dir(fixtures_root):
-    return fixtures_root / "data/question_answering"
-
-
-@pytest.fixture
-def dev_dataset(dataset_loader, fixtures_dir):
-    return dataset_loader.load("validation")
+def data_processor(args):
+    tokenizer = QuestionAnsweringTokenizer.from_pretrained(args.model_type)
+    return SquadQuestionAnsweringDataProcessor(tokenizer)
 
 
 @pytest.mark.parametrize(
@@ -61,12 +42,13 @@ def dev_dataset(dataset_loader, fixtures_dir):
         (2, "Where did Super Bowl 50 take place?"),
     ],
 )
-def test_data_processor(tokenizer, dev_dataset, args, index, question):
+def test_data_processor(data_processor, dev_dataset, args, index, question):
     if args.uncased:
         question = question.lower()
-    assert get_question(dev_dataset[index], tokenizer) == question
+    processed_instance = data_processor(dev_dataset[index])
+    assert decode_question(processed_instance, data_processor.tokenizer) == question
 
 
-def get_question(instance: Dict, tokenizer: QuestionAnsweringTokenizer) -> str:
+def decode_question(instance: Dict, tokenizer: QuestionAnsweringTokenizer) -> str:
     field_value = tokenizer.decode(instance["question"])
     return field_value.lstrip()
diff --git a/tests/trapper/data/test_dataset_reader.py b/tests/trapper/data/test_dataset_reader.py
new file mode 100644
index 0000000..04353fd
--- /dev/null
+++ b/tests/trapper/data/test_dataset_reader.py
@@ -0,0 +1,18 @@
+import pytest
+
+from trapper.data import DatasetReader
+
+
+@pytest.fixture
+def dataset_reader():
+    return DatasetReader(path="squad_qa_test_fixture")
+
+
+def test_split_names(dataset_reader):
+    assert "train" in dataset_reader.split_names
+    assert "validation" in dataset_reader.split_names
+
+
+def test_get_dataset(dataset_reader):
+    assert len(dataset_reader.get_dataset("train")) == 5
+    assert len(dataset_reader.get_dataset("validation")) == 6
diff --git a/tests/trapper/training/test_trainer_for_question_answering.py b/tests/trapper/training/test_trainer_for_question_answering.py
index b855578..e188682 100644
--- a/tests/trapper/training/test_trainer_for_question_answering.py
+++ b/tests/trapper/training/test_trainer_for_question_answering.py
@@ -1,8 +1,11 @@
+import datasets
 import pytest
+from transformers import DistilBertForQuestionAnswering, DistilBertTokenizerFast
 
-from trapper.common.params import Params
-from trapper.data import IndexedDataset
-from trapper.training import TransformerTrainer
+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
 
 
@@ -28,13 +31,12 @@ def trainer_params(temp_output_dir, temp_result_dir, temp_cache_dir):
         "train_split_name": "train",
         "dev_split_name": "validation",
         "tokenizer": {"type": "question-answering"},
-        "data_collator": {"type": "question-answering"},
-        "data_processor": {
-            "type": "squad-question-answering",
-        },
         "dataset_loader": {
-            "path": "squad_qa_test_fixture",
+            "dataset_reader": {"path": "squad_qa_test_fixture"},
+            "data_processor": {"type": "squad-question-answering"},
+            "data_adapter": {"type": "question-answering"},
         },
+        "data_collator": {},
         "model": {"type": "question_answering"},
         "args": {
             "type": "default",
@@ -79,15 +81,14 @@ def trainer(trainer_params) -> TransformerTrainer:
 
 
 def test_trainer_fields(trainer):
-    assert type(trainer.tokenizer).__name__ == "DistilBertTokenizerFast"
-    assert (
-        type(trainer.data_collator).__name__ == "DataCollatorForQuestionAnswering"
-    )
-    assert isinstance(trainer.train_dataset, IndexedDataset)
-    assert isinstance(trainer.eval_dataset, IndexedDataset)
-    assert type(trainer).__name__ == "TransformerTrainer"
-    assert type(trainer.model).__name__ == "DistilBertForQuestionAnswering"
-    assert type(trainer.optimizer).__name__ == "HuggingfaceAdamWOptimizer"
+    assert isinstance(trainer, TransformerTrainer)
+    assert isinstance(trainer.model, DistilBertForQuestionAnswering)
+    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):