Predicting with a dataset reader on a multitask model (#5115)

* Create a way to use allennlp predict with a dataset and a multitask model

* Fix type ignoration

* Changelog

* Fix to the predictor

Author: dirkgr

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Ported the following Huggingface `LambdaLR`-based schedulers: `ConstantLearningRateScheduler`, `ConstantWithWarmupLearningRateScheduler`, `CosineWithWarmupLearningRateScheduler`, `CosineHardRestartsWithWarmupLearningRateScheduler`.
 - Added a T5 implementation to `modules.transformers`.
 - Added new `sub_token_mode` parameter to `pretrained_transformer_mismatched_embedder` class to support first sub-token embedding
+- Added a way to run a multi task model with a dataset reader as part of `allennlp predict`.
 - Added new `eval_mode` in `PretrainedTransformerEmbedder`. If it is set to `True`, the transformer is _always_ run in evaluation mode, which, e.g., disables dropout and does not update batch normalization statistics.
 - Added additional parameters to the W&B callback: `entity`, `group`, `name`, `notes`, and `wandb_kwargs`.
 

allennlp/commands/predict.py

@@ -3,7 +3,6 @@
 or dataset to JSON predictions using a trained model and its
 [`Predictor`](../predictors/predictor.md#predictor) wrapper.
 """
-
 from typing import List, Iterator, Optional
 import argparse
 import sys
@@ -16,6 +15,7 @@
 from allennlp.common.checks import check_for_gpu, ConfigurationError
 from allennlp.common.file_utils import cached_path
 from allennlp.common.util import lazy_groups_of
+from allennlp.data.dataset_readers import MultiTaskDatasetReader
 from allennlp.models.archival import load_archive
 from allennlp.predictors.predictor import Predictor, JsonDict
 from allennlp.data import Instance
@@ -73,6 +73,14 @@ def add_subparser(self, parser: argparse._SubParsersAction) -> argparse.Argument
             "flag is set.",
         )
 
+        subparser.add_argument(
+            "--multitask-head",
+            type=str,
+            default=None,
+            help="If you are using a dataset reader to make predictions, and the model is a"
+            "multitask model, you have to specify the name of the model head to use here.",
+        )
+
         subparser.add_argument(
             "-o",
             "--overrides",
@@ -144,6 +152,7 @@ def __init__(
         batch_size: int,
         print_to_console: bool,
         has_dataset_reader: bool,
+        multitask_head: Optional[str] = None,
     ) -> None:
         self._predictor = predictor
         self._input_file = input_file
@@ -152,6 +161,24 @@ def __init__(
         self._print_to_console = print_to_console
         self._dataset_reader = None if not has_dataset_reader else predictor._dataset_reader
 
+        self._multitask_head = multitask_head
+        if self._multitask_head is not None:
+            if self._dataset_reader is None:
+                raise ConfigurationError(
+                    "You must use a dataset reader when using --multitask-head."
+                )
+            if not isinstance(self._dataset_reader, MultiTaskDatasetReader):
+                raise ConfigurationError(
+                    "--multitask-head only works with a multitask dataset reader."
+                )
+        if (
+            isinstance(self._dataset_reader, MultiTaskDatasetReader)
+            and self._multitask_head is None
+        ):
+            raise ConfigurationError(
+                "You must specify --multitask-head when using a multitask dataset reader."
+            )
+
     def _predict_json(self, batch_data: List[JsonDict]) -> Iterator[str]:
         if len(batch_data) == 1:
             results = [self._predictor.predict_json(batch_data[0])]
@@ -196,7 +223,15 @@ def _get_instance_data(self) -> Iterator[Instance]:
         elif self._dataset_reader is None:
             raise ConfigurationError("To generate instances directly, pass a DatasetReader.")
         else:
-            yield from self._dataset_reader.read(self._input_file)
+            if isinstance(self._dataset_reader, MultiTaskDatasetReader):
+                assert (
+                    self._multitask_head is not None
+                )  # This is properly checked by the constructor.
+                yield from self._dataset_reader.read(
+                    self._input_file, force_task=self._multitask_head
+                )
+            else:
+                yield from self._dataset_reader.read(self._input_file)
 
     def run(self) -> None:
         has_reader = self._dataset_reader is not None
@@ -235,5 +270,6 @@ def _predict(args: argparse.Namespace) -> None:
         args.batch_size,
         not args.silent,
         args.use_dataset_reader,
+        args.multitask_head,
     )
     manager.run()

allennlp/data/data_loaders/multitask_data_loader.py

@@ -6,11 +6,6 @@
 from overrides import overrides
 
 from allennlp.common import util
-from allennlp.data.dataset_readers.dataset_reader import (
-    DatasetReader,
-    DatasetReaderInput,
-    WorkerInfo,
-)
 from allennlp.data.batch import Batch
 from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict
 from allennlp.data.data_loaders.multiprocess_data_loader import MultiProcessDataLoader
@@ -251,7 +246,7 @@ def _get_instances_for_epoch(self) -> Dict[str, Iterable[Instance]]:
 
     def _make_data_loader(self, key: str) -> MultiProcessDataLoader:
         kwargs: Dict[str, Any] = {
-            "reader": _MultitaskDatasetReaderShim(self.readers[key], key),
+            "reader": self.readers[key],
             "data_path": self.data_paths[key],
             # We don't load batches from this data loader, only instances, but we have to set
             # something for the batch size, so we set 1.
@@ -264,39 +259,3 @@ def _make_data_loader(self, key: str) -> MultiProcessDataLoader:
         if key in self._start_method:
             kwargs["start_method"] = self._start_method[key]
         return MultiProcessDataLoader(**kwargs)
-
-
-@DatasetReader.register("multitask_shim")
-class _MultitaskDatasetReaderShim(DatasetReader):
-    """This dataset reader wraps another dataset reader and adds the name of the "task" into
-    each instance as a metadata field. This exists only to support `MultitaskDataLoader`. You
-    should not have to use this yourself."""
-
-    def __init__(self, inner: DatasetReader, head: str, **kwargs):
-        super().__init__(**kwargs)
-        self.inner = inner
-        self.head = head
-
-    def _set_worker_info(self, info: Optional[WorkerInfo]) -> None:
-        """
-        Should only be used internally.
-        """
-        super()._set_worker_info(info)
-        self.inner._set_worker_info(info)
-
-    def read(self, file_path: DatasetReaderInput) -> Iterator[Instance]:
-        from allennlp.data.fields import MetadataField
-
-        for instance in self.inner.read(file_path):
-            instance.add_field("task", MetadataField(self.head))
-            yield instance
-
-    def text_to_instance(self, *inputs) -> Instance:
-        from allennlp.data.fields import MetadataField
-
-        instance = self.inner.text_to_instance(*inputs)
-        instance.add_field("task", MetadataField(self.head))
-        return instance
-
-    def apply_token_indexers(self, instance: Instance) -> None:
-        self.inner.apply_token_indexers(instance)

allennlp/data/dataset_readers/multitask.py

@@ -1,8 +1,12 @@
 from os import PathLike
-from typing import Dict, Iterator, Union
+from typing import Dict, Iterator, Union, Optional
 
 from allennlp.data.instance import Instance
-from allennlp.data.dataset_readers.dataset_reader import DatasetReader
+from allennlp.data.dataset_readers.dataset_reader import (
+    DatasetReader,
+    WorkerInfo,
+    DatasetReaderInput,
+)
 
 
 @DatasetReader.register("multitask")
@@ -25,7 +29,51 @@ class MultiTaskDatasetReader(DatasetReader):
     """
 
     def __init__(self, readers: Dict[str, DatasetReader]) -> None:
-        self.readers = readers
+        self.readers = {
+            task: _MultitaskDatasetReaderShim(reader, task) for task, reader in readers.items()
+        }
 
-    def read(self, file_paths: Dict[str, Union[PathLike, str]]) -> Dict[str, Iterator[Instance]]:  # type: ignore
-        raise RuntimeError("This class is not designed to be called like this")
+    def read(  # type: ignore
+        self,
+        file_paths: Union[PathLike, str, Dict[str, Union[PathLike, str]]],
+        *,
+        force_task: Optional[str] = None
+    ) -> Union[Iterator[Instance], Dict[str, Iterator[Instance]]]:
+        if force_task is None:
+            raise RuntimeError("This class is not designed to be called like this.")
+        return self.readers[force_task].read(file_paths)
+
+
+@DatasetReader.register("multitask_shim")
+class _MultitaskDatasetReaderShim(DatasetReader):
+    """This dataset reader wraps another dataset reader and adds the name of the "task" into
+    each instance as a metadata field. You should not have to use this yourself."""
+
+    def __init__(self, inner: DatasetReader, head: str, **kwargs):
+        super().__init__(**kwargs)
+        self.inner = inner
+        self.head = head
+
+    def _set_worker_info(self, info: Optional[WorkerInfo]) -> None:
+        """
+        Should only be used internally.
+        """
+        super()._set_worker_info(info)
+        self.inner._set_worker_info(info)
+
+    def read(self, file_path: DatasetReaderInput) -> Iterator[Instance]:
+        from allennlp.data.fields import MetadataField
+
+        for instance in self.inner.read(file_path):
+            instance.add_field("task", MetadataField(self.head))
+            yield instance
+
+    def text_to_instance(self, *inputs) -> Instance:
+        from allennlp.data.fields import MetadataField
+
+        instance = self.inner.text_to_instance(*inputs)
+        instance.add_field("task", MetadataField(self.head))
+        return instance
+
+    def apply_token_indexers(self, instance: Instance) -> None:
+        self.inner.apply_token_indexers(instance)

allennlp/predictors/multitask.py

@@ -52,7 +52,7 @@ def __init__(self, model: MultiTaskModel, dataset_reader: MultiTaskDatasetReader
             predictor_class: Type[Predictor] = (
                 Predictor.by_name(predictor_name) if predictor_name is not None else Predictor  # type: ignore
             )
-            self.predictors[name] = predictor_class(model, dataset_reader.readers[name])
+            self.predictors[name] = predictor_class(model, dataset_reader.readers[name].inner)
 
     @overrides
     def predict_instance(self, instance: Instance) -> JsonDict: