allow multiple token embedders, but only use first with non-empty type (#129)

* allow multiple token embedders, but only use first with non-empty type

* blacken the modified code

* update CHANGELOG

* Throw an exception if we already know the model isn't going to work.

* Checks the config by parsing it

Co-authored-by: Anthony P Panza <anthony.p.panza@boeing.com>
Co-authored-by: Dirk Groeneveld <dirkg@allenai.org>

Author: 3 people

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Updated `LanguageModelTokenEmbedder` to allow allow multiple token embedders, but only use first with non-empty type
 - Fixed evaluation of metrics when using distributed setting.
 - Fixed a bug introduced in 1.0 where the SRL model did not reproduce the original result.
 

allennlp_models/lm/modules/token_embedders/language_model.py

@@ -2,10 +2,14 @@
 from typing import Dict, Tuple, TYPE_CHECKING
 
 import torch
+from allennlp.common import Params
 
 from allennlp.common.checks import ConfigurationError
 from allennlp.data import TokenIndexer, Token
+from allennlp.modules import TextFieldEmbedder
 from allennlp.modules.scalar_mix import ScalarMix
+from allennlp.modules.text_field_embedders import BasicTextFieldEmbedder
+from allennlp.modules.token_embedders import EmptyEmbedder
 from allennlp.modules.token_embedders.token_embedder import TokenEmbedder
 from allennlp.nn.util import (
     remove_sentence_boundaries,
@@ -76,15 +80,29 @@ def __init__(
 
         # Extract the name of the tokens that the LM was trained on.
         text_field_embedder = dict_config["model"]["text_field_embedder"]
-        token_names = list(text_field_embedder["token_embedders"].keys())
-        if len(token_names) != 1:
-            # We don't currently support embedding with language models trained with multiple
-            # embedded indices.
-            #
-            # Note: We only care about embedded indices. This does not include "tokens" which
-            # is just used to compute the loss in LanguageModel.
-            raise ConfigurationError(f"LM from {archive_file} trained with multiple embedders!")
-        self._token_name = token_names[0]
+        text_field_embedder = TextFieldEmbedder.from_params(Params(text_field_embedder))
+        if not isinstance(text_field_embedder, BasicTextFieldEmbedder):
+            raise ConfigurationError(
+                f"Language model from {archive_file} uses a non-standard TextFieldEmbedder!"
+            )
+        non_empty_embedders = [
+            name
+            for name, token_embedder in text_field_embedder._token_embedders.items()
+            if not isinstance(token_embedder, EmptyEmbedder)
+        ]
+
+        if len(non_empty_embedders) == 0:
+            # Only empty embedders were contained in the language model
+            # We need at least one non-empty embedder in the language model
+            raise ConfigurationError(
+                f"Language model from {archive_file} trained with only empty embedders!"
+            )
+        elif len(non_empty_embedders) > 1:
+            raise ConfigurationError(
+                f"Language model from {archive_file} trained with multiple non-empty embedders!"
+            )
+
+        self._token_name = non_empty_embedders[0]
 
         # TODO(brendanr): Find a way to remove this hack. The issue fundamentally is that the
         # BasicTextFieldEmbedder concatenates multiple embedded representations. When a