Move default predictors (#4154)

* Move DecomposableAttention test

* Move default predictors

* Fix test

* Productivity through formatting

Author: dirkgr

allennlp/models/basic_classifier.py

@@ -179,3 +179,5 @@ def make_output_human_readable(
     def get_metrics(self, reset: bool = False) -> Dict[str, float]:
         metrics = {"accuracy": self._accuracy.get_metric(reset)}
         return metrics
+
+    default_predictor = "text_classifier"

allennlp/models/model.py

@@ -67,6 +67,7 @@ class Model(torch.nn.Module, Registrable):
     """
 
     _warn_for_unseparable_batches: Set[str] = set()
+    default_predictor: Optional[str] = None
 
     def __init__(self, vocab: Vocabulary, regularizer: RegularizerApplicator = None) -> None:
         super().__init__()

allennlp/models/simple_tagger.py

@@ -215,3 +215,5 @@ def get_metrics(self, reset: bool = False) -> Dict[str, float]:
             else:
                 metrics_to_return.update({x: y for x, y in f1_dict.items() if "overall" in x})
         return metrics_to_return
+
+    default_predictor = "sentence-tagger"

allennlp/predictors/__init__.py

@@ -7,6 +7,5 @@
 a `Predictor` that wraps it.
 """
 from allennlp.predictors.predictor import Predictor
-from allennlp.predictors.decomposable_attention import DecomposableAttentionPredictor
 from allennlp.predictors.sentence_tagger import SentenceTaggerPredictor
 from allennlp.predictors.text_classifier import TextClassifierPredictor

allennlp/predictors/decomposable_attention.py

@@ -15,24 +15,6 @@
 from allennlp.models.archival import Archive, load_archive
 from allennlp.nn import util
 
-# a mapping from model `type` to the default Predictor for that type
-DEFAULT_PREDICTORS = {
-    "atis_parser": "atis-parser",
-    "basic_classifier": "text_classifier",
-    "biaffine_parser": "biaffine-dependency-parser",
-    "bimpm": "textual-entailment",
-    "constituency_parser": "constituency-parser",
-    "coref": "coreference-resolution",
-    "crf_tagger": "sentence-tagger",
-    "decomposable_attention": "textual-entailment",
-    "event2mind": "event2mind",
-    "simple_tagger": "sentence-tagger",
-    "srl": "semantic-role-labeling",
-    "srl_bert": "semantic-role-labeling",
-    "quarel_parser": "quarel-parser",
-    "wikitables_mml_parser": "wikitables-parser",
-}
-
 
 class Predictor(Registrable):
     """
@@ -299,8 +281,8 @@ def from_archive(
 
         if not predictor_name:
             model_type = config.get("model").get("type")
-            if model_type in DEFAULT_PREDICTORS:
-                predictor_name = DEFAULT_PREDICTORS[model_type]
+            model_class, _ = Model.resolve_class_name(model_type)
+            predictor_name = model_class.default_predictor
         predictor_class: Type[Predictor] = Predictor.by_name(  # type: ignore
             predictor_name
         ) if predictor_name is not None else cls

allennlp/predictors/predictor.py

@@ -18,7 +18,6 @@
 from allennlp.data.dataset_readers import DatasetReader, TextClassificationJsonReader
 from allennlp.models.archival import load_archive
 from allennlp.predictors import Predictor, TextClassifierPredictor
-from allennlp.predictors.predictor import DEFAULT_PREDICTORS
 
 
 class TestPredict(AllenNlpTestCase):
@@ -224,31 +223,42 @@ def test_base_predictor(self):
         model_path = str(self.classifier_model_path)
         archive = load_archive(model_path)
         model_type = archive.config.get("model").get("type")
-        # Makes sure that we don't have a DEFAULT_PREDICTOR for it. Otherwise the base class
+        # Makes sure that we don't have a default_predictor for it. Otherwise the base class
         # implementation wouldn't be used
-        del DEFAULT_PREDICTORS["basic_classifier"]
-        assert model_type not in DEFAULT_PREDICTORS
+        from allennlp.models import Model
 
-        # Doesn't use a --predictor
-        sys.argv = [
-            "__main__.py",  # executable
-            "predict",  # command
-            model_path,
-            str(self.classifier_data_path),  # input_file
-            "--output-file",
-            str(self.outfile),
-            "--silent",
-            "--use-dataset-reader",
-        ]
-        main()
-        assert os.path.exists(self.outfile)
-        with open(self.outfile, "r") as f:
-            results = [json.loads(line) for line in f]
+        model_class, _ = Model.resolve_class_name(model_type)
+        saved_default_predictor = model_class.default_predictor
+        model_class.default_predictor = None
+        try:
+            # Doesn't use a --predictor
+            sys.argv = [
+                "__main__.py",  # executable
+                "predict",  # command
+                model_path,
+                str(self.classifier_data_path),  # input_file
+                "--output-file",
+                str(self.outfile),
+                "--silent",
+                "--use-dataset-reader",
+            ]
+            main()
+            assert os.path.exists(self.outfile)
+            with open(self.outfile, "r") as f:
+                results = [json.loads(line) for line in f]
 
-        assert len(results) == 3
-        for result in results:
-            assert set(result.keys()) == {"logits", "probs", "label", "loss", "tokens", "token_ids"}
-        DEFAULT_PREDICTORS["basic_classifier"] = "text_classifier"
+            assert len(results) == 3
+            for result in results:
+                assert set(result.keys()) == {
+                    "logits",
+                    "probs",
+                    "label",
+                    "loss",
+                    "tokens",
+                    "token_ids",
+                }
+        finally:
+            model_class.default_predictor = saved_default_predictor
 
     def test_batch_prediction_works_with_known_model(self):
         with open(self.infile, "w") as f: