fix bug with interleaving dataset reader (#5122)

* fix bug with interleaving dataset reader

* more tests

* Update allennlp/data/dataset_readers/interleaving_dataset_reader.py

* Update allennlp/data/dataset_readers/interleaving_dataset_reader.py

Author: epwalsh

CHANGELOG.md

@@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed a stall when using distributed training and gradient accumulation at the same time
 - Fixed an issue where using the `from_pretrained_transformer` `Vocabulary` constructor in distributed training via the `allennlp train` command
   would result in the data being iterated through unnecessarily.
+- Fixed a bug regarding token indexers with the `InterleavingDatasetReader` when used with multi-process data loading.
 - Fixed a warning from `transformers` when using `max_length` in the `PretrainedTransformerTokenizer`.
 
 ### Removed

allennlp/data/dataset_readers/interleaving_dataset_reader.py

@@ -1,8 +1,15 @@
-from typing import Dict, Mapping, Iterable, Union
+from typing import Dict, Mapping, Iterable, Union, Optional
 import json
 
+from overrides import overrides
+
 from allennlp.common.checks import ConfigurationError
-from allennlp.data.dataset_readers.dataset_reader import DatasetReader, PathOrStr
+from allennlp.data.dataset_readers.dataset_reader import (
+    DatasetReader,
+    PathOrStr,
+    WorkerInfo,
+    DistributedInfo,
+)
 from allennlp.data.fields import MetadataField
 from allennlp.data.instance import Instance
 
@@ -52,6 +59,18 @@ def __init__(
             raise ConfigurationError(f"invalid scheme: {scheme}")
         self._scheme = scheme
 
+    @overrides
+    def _set_worker_info(self, info: Optional[WorkerInfo]) -> None:
+        super()._set_worker_info(info)
+        for reader in self._readers.values():
+            reader._set_worker_info(info)
+
+    @overrides
+    def _set_distributed_info(self, info: Optional[DistributedInfo]) -> None:
+        super()._set_distributed_info(info)
+        for reader in self._readers.values():
+            reader._set_distributed_info(info)
+
     def _read_round_robin(self, datasets: Mapping[str, Iterable[Instance]]) -> Iterable[Instance]:
         remaining = set(datasets)
         dataset_iterators = {key: iter(dataset) for key, dataset in datasets.items()}
@@ -72,6 +91,7 @@ def _read_all_at_once(self, datasets: Mapping[str, Iterable[Instance]]) -> Itera
                 instance.fields[self._dataset_field_name] = MetadataField(key)
                 yield instance
 
+    @overrides
     def _read(self, file_path: Union[str, Dict[str, PathOrStr]]) -> Iterable[Instance]:
         if isinstance(file_path, str):
             try:
@@ -97,6 +117,11 @@ def _read(self, file_path: Union[str, Dict[str, PathOrStr]]) -> Iterable[Instanc
         else:
             raise RuntimeError("impossible to get here")
 
-    def text_to_instance(self) -> Instance:  # type: ignore
+    @overrides
+    def text_to_instance(self, dataset_key: str, *args, **kwargs) -> Instance:  # type: ignore
+        return self._readers[dataset_key].text_to_instance(*args, **kwargs)  # type: ignore[call-arg]
 
-        raise RuntimeError("text_to_instance doesn't make sense here")
+    @overrides
+    def apply_token_indexers(self, instance: Instance) -> None:
+        dataset = instance.fields[self._dataset_field_name].metadata  # type: ignore[attr-defined]
+        self._readers[dataset].apply_token_indexers(instance)

tests/data/dataset_readers/interleaving_dataset_reader_test.py

@@ -1,11 +1,15 @@
 from typing import Iterable
 
+import pytest
+
 from allennlp.common.testing import AllenNlpTestCase
 from allennlp.data.dataset_readers import DatasetReader, InterleavingDatasetReader
+from allennlp.data.data_loaders import MultiProcessDataLoader
 from allennlp.data.fields import TextField
 from allennlp.data.instance import Instance
 from allennlp.data.token_indexers import SingleIdTokenIndexer
 from allennlp.data.tokenizers import SpacyTokenizer
+from allennlp.data.vocabulary import Vocabulary
 
 
 class PlainTextReader(DatasetReader):
@@ -20,9 +24,11 @@ def _read(self, file_path: str) -> Iterable[Instance]:
                 yield self.text_to_instance(line)
 
     def text_to_instance(self, line: str) -> Instance:  # type: ignore
-
         tokens = self._tokenizer.tokenize(line)
-        return Instance({"line": TextField(tokens, self._token_indexers)})
+        return Instance({"line": TextField(tokens)})
+
+    def apply_token_indexers(self, instance):
+        instance.fields["line"].token_indexers = self._token_indexers
 
 
 class TestInterleavingDatasetReader(AllenNlpTestCase):
@@ -72,3 +78,26 @@ def test_all_at_once(self):
 
         # should be in 3 buckets
         assert len(buckets) == 3
+
+    @pytest.mark.parametrize("lazy", (True, False))
+    def test_with_multi_process_loading(self, lazy):
+        readers = {"a": PlainTextReader(), "b": PlainTextReader(), "c": PlainTextReader()}
+        reader = InterleavingDatasetReader(readers)
+        data_dir = self.FIXTURES_ROOT / "data"
+        file_path = {
+            "a": data_dir / "babi.txt",
+            "b": data_dir / "conll2003.txt",
+            "c": data_dir / "conll2003.txt",
+        }
+        vocab = Vocabulary.from_instances(reader.read(file_path))
+        loader = MultiProcessDataLoader(
+            reader,
+            file_path,
+            num_workers=1,
+            batch_size=1,
+            max_instances_in_memory=2 if lazy else None,
+        )
+        loader.index_with(vocab)
+
+        list(loader.iter_instances())
+        list(loader)