Added DataCollator for dynamic operations for each batch.

CHANGELOG.md

@@ -6,6 +6,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
+### Added
+- Added `DataCollator` for dynamic operations for each batch.
 
 ### Changed
 

allennlp/data/data_loaders/__init__.py

@@ -1,4 +1,5 @@
-from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict, allennlp_collate
+from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict
 from allennlp.data.data_loaders.multiprocess_data_loader import MultiProcessDataLoader, WorkerError
 from allennlp.data.data_loaders.multitask_data_loader import MultiTaskDataLoader
 from allennlp.data.data_loaders.simple_data_loader import SimpleDataLoader
+from allennlp.data.data_loaders.data_collator import allennlp_collate

allennlp/data/data_loaders/data_collator.py

@@ -0,0 +1,68 @@
+from typing import List, Tuple
+
+import torch
+from transformers.data.data_collator import DataCollatorForLanguageModeling
+
+from allennlp.common import Registrable
+from allennlp.data.batch import Batch
+from allennlp.data.data_loaders.data_loader import TensorDict
+from allennlp.data.instance import Instance
+
+
+def allennlp_collate(instances: List[Instance]) -> TensorDict:
+    """
+    This is the default function used to turn a list of `Instance`s into a `TensorDict`
+    batch.
+    """
+    batch = Batch(instances)
+    return batch.as_tensor_dict()
+
+
+class DataCollator(Registrable):
+    """
+    This class is similar with `DataCollator` in [Transformers]
+    (https://github.com/huggingface/transformers/blob/master/src/transformers/data/data_collator.py)
+    Allow to do some dynamic operations for tensor in different batches
+    Cause this method run before each epoch to convert `List[Instance]` to `TensorDict`
+    """
+
+    def __call__(self, instances: List[Instance]) -> TensorDict:
+        raise NotImplementedError
+
+
+class DefaultDataCollator(DataCollator):
+    def __call__(self, instances: List[Instance]) -> TensorDict:
+        return allennlp_collate(instances)
+
+
+@DataCollator.register("language-model")
+class LanguageModelingDataCollator(DataCollator):
+    """
+    Register as an `DataCollator` with name `LanguageModelingDataCollator`
+    Used for language modeling.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        mlm: bool = True,
+        mlm_probability: float = 0.15,
+        namespace: str = "tokens",
+    ):
+        from allennlp.common import cached_transformers
+
+        tokenizer = cached_transformers.get_tokenizer(model_name)
+        self._namespace = namespace
+        self._collator = DataCollatorForLanguageModeling(tokenizer, mlm, mlm_probability)
+
+    def __call__(self, instances: List[Instance]) -> TensorDict:
+        tensor_dicts = allennlp_collate(instances)
+        tensor_dicts = self.process_tokens(tensor_dicts)
+        return tensor_dicts
+
+    def process_tokens(self, tensor_dicts: TensorDict) -> Tuple[torch.Tensor, torch.Tensor]:
+        inputs = tensor_dicts[self._namespace]["token_ids"]["token_ids"]
+        inputs, labels = self._collator.mask_tokens(inputs)
+        tensor_dicts[self._namespace]["tokens"]["token_ids"] = inputs
+        tensor_dicts[self._namespace]["tokens"]["labels"] = labels
+        return tensor_dicts

allennlp/data/data_loaders/data_loader.py

@@ -1,10 +1,9 @@
-from typing import List, Dict, Union, Iterator
+from typing import Dict, Union, Iterator
 
 import torch
 
 from allennlp.common.registrable import Registrable
 from allennlp.data.instance import Instance
-from allennlp.data.batch import Batch
 from allennlp.data.vocabulary import Vocabulary
 
 
@@ -14,15 +13,6 @@
 """
 
 
-def allennlp_collate(instances: List[Instance]) -> TensorDict:
-    """
-    This is the default function used to turn a list of `Instance`s into a `TensorDict`
-    batch.
-    """
-    batch = Batch(instances)
-    return batch.as_tensor_dict()
-
-
 class DataLoader(Registrable):
     """
     A `DataLoader` is responsible for generating batches of instances from a

allennlp/data/data_loaders/multiprocess_data_loader.py

@@ -12,7 +12,8 @@
 from allennlp.common.util import lazy_groups_of, shuffle_iterable
 from allennlp.common.tqdm import Tqdm
 from allennlp.data.instance import Instance
-from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict, allennlp_collate
+from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict
+from allennlp.data.data_loaders.data_collator import DataCollator, DefaultDataCollator
 from allennlp.data.dataset_readers import DatasetReader, WorkerInfo, DatasetReaderInput
 from allennlp.data.fields import TextField
 from allennlp.data.samplers import BatchSampler
@@ -124,6 +125,8 @@ class MultiProcessDataLoader(DataLoader):
     quiet : `bool`, optional (default = `False`)
         If `True`, tqdm progress bars will be disabled.
 
+    collate_fn : `DataCollator`, optional ( default = `DefaultDataCollator`)
+
     # Best practices
 
     - **Large datasets**
@@ -207,6 +210,7 @@ def __init__(
         start_method: str = "fork",
         cuda_device: Optional[Union[int, str, torch.device]] = None,
         quiet: bool = False,
+        collate_fn: DataCollator = DefaultDataCollator(),
     ) -> None:
         # Do some parameter validation.
         if num_workers is not None and num_workers < 0:
@@ -244,7 +248,7 @@ def __init__(
         self.batch_sampler = batch_sampler
         self.batches_per_epoch = batches_per_epoch
         self.num_workers = num_workers
-        self.collate_fn = allennlp_collate
+        self.collate_fn = collate_fn
         self.max_instances_in_memory = max_instances_in_memory
         self.start_method = start_method
         self.quiet = quiet

allennlp/data/data_loaders/multitask_data_loader.py

@@ -258,4 +258,4 @@ def _make_data_loader(self, key: str) -> MultiProcessDataLoader:
             kwargs["max_instances_in_memory"] = self._max_instances_in_memory[key]
         if key in self._start_method:
             kwargs["start_method"] = self._start_method[key]
-        return MultiProcessDataLoader(**kwargs)
+        return MultiProcessDataLoader(**kwargs)

allennlp/data/data_loaders/simple_data_loader.py

@@ -7,7 +7,8 @@
 
 from allennlp.common.util import lazy_groups_of
 from allennlp.common.tqdm import Tqdm
-from allennlp.data.data_loaders.data_loader import DataLoader, allennlp_collate, TensorDict
+from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict
+from allennlp.data.data_loaders.data_collator import DefaultDataCollator
 from allennlp.data.dataset_readers import DatasetReader
 from allennlp.data.instance import Instance
 from allennlp.data.vocabulary import Vocabulary
@@ -36,6 +37,7 @@ def __init__(
         self.vocab = vocab
         self.cuda_device: Optional[torch.device] = None
         self._batch_generator: Optional[Iterator[TensorDict]] = None
+        self.collate_fn = DefaultDataCollator()
 
     def __len__(self) -> int:
         if self.batches_per_epoch is not None:
@@ -60,7 +62,7 @@ def _iter_batches(self) -> Iterator[TensorDict]:
         if self.shuffle:
             random.shuffle(self.instances)
         for batch in lazy_groups_of(self.iter_instances(), self.batch_size):
-            tensor_dict = allennlp_collate(batch)
+            tensor_dict = self.collate_fn(batch)
             if self.cuda_device is not None:
                 tensor_dict = nn_util.move_to_device(tensor_dict, self.cuda_device)
             yield tensor_dict