Resize transformers word embeddings layer for additional_special_tokens (#4946)

* Adding a mechanism to resize the word embeddings layer from transformers models in case additional special tokens are provided in tokenizer_kwargs.

* Updating changelog

* Reformatting test file with black

* Fixing failed test for transformer model that don't implement get_input_embeddings()

* Adding message to warn user about the transformer model being unable to resize it's embeddings layer when additional tokens are provided

* Reformatting with black

Co-authored-by: Dirk Groeneveld <[email protected]>

Author: pvcastro

CHANGELOG.md

@@ -28,6 +28,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Added `tokenizer_kwargs` and `transformer_kwargs` arguments to `PretrainedTransformerBackbone`
+- Resize transformers word embeddings layer for `additional_special_tokens`
 
 ### Changed
 

allennlp/modules/token_embedders/pretrained_transformer_embedder.py

@@ -1,3 +1,4 @@
+import logging
 import math
 from typing import Optional, Tuple, Dict, Any
 
@@ -12,6 +13,8 @@
 from allennlp.modules.token_embedders.token_embedder import TokenEmbedder
 from allennlp.nn.util import batched_index_select
 
+logger = logging.getLogger(__name__)
+
 
 @TokenEmbedder.register("pretrained_transformer")
 class PretrainedTransformerEmbedder(TokenEmbedder):
@@ -101,6 +104,19 @@ def __init__(
             model_name,
             tokenizer_kwargs=tokenizer_kwargs,
         )
+
+        try:
+            if self.transformer_model.get_input_embeddings().num_embeddings != len(
+                tokenizer.tokenizer
+            ):
+                self.transformer_model.resize_token_embeddings(len(tokenizer.tokenizer))
+        except NotImplementedError:
+            # Can't resize for transformers models that don't implement base_model.get_input_embeddings()
+            logger.warning(
+                "Could not resize the token embedding matrix of the transformer model. "
+                "This model does not support resizing."
+            )
+
         self._num_added_start_tokens = len(tokenizer.single_sequence_start_tokens)
         self._num_added_end_tokens = len(tokenizer.single_sequence_end_tokens)
         self._num_added_tokens = self._num_added_start_tokens + self._num_added_end_tokens

allennlp/training/metric_tracker.py

@@ -128,6 +128,6 @@ def combined_score(self, metrics: Dict[str, float]) -> float:
             )
         except KeyError as e:
             raise ConfigurationError(
-                f"You configured the trainer to use the {e.args[0]}"
+                f"You configured the trainer to use the {e.args[0]} "
                 "metric for early stopping, but the model did not produce that metric."
             )

tests/modules/token_embedders/pretrained_transformer_embedder_test.py

@@ -315,3 +315,18 @@ def test_encoder_decoder_model(self):
         token_ids = torch.LongTensor([[1, 2, 3], [2, 3, 4]])
         mask = torch.ones_like(token_ids).bool()
         token_embedder(token_ids, mask)
+
+    def test_embeddings_resize(self):
+        regular_token_embedder = PretrainedTransformerEmbedder("bert-base-cased")
+        assert (
+            regular_token_embedder.transformer_model.embeddings.word_embeddings.num_embeddings
+            == 28996
+        )
+        tokenizer_kwargs = {"additional_special_tokens": ["<NEW_TOKEN>"]}
+        enhanced_token_embedder = PretrainedTransformerEmbedder(
+            "bert-base-cased", tokenizer_kwargs=tokenizer_kwargs
+        )
+        assert (
+            enhanced_token_embedder.transformer_model.embeddings.word_embeddings.num_embeddings
+            == 28997
+        )