Added a TokenEmbedder for use with pytorch-transformers

allennlp/modules/token_embedders/__init__.py

@@ -16,3 +16,4 @@
         LanguageModelTokenEmbedder
 from allennlp.modules.token_embedders.bag_of_word_counts_token_embedder import BagOfWordCountsTokenEmbedder
 from allennlp.modules.token_embedders.pass_through_token_embedder import PassThroughTokenEmbedder
+from allennlp.modules.token_embedders.pretrained_transformer_embedder import PretrainedTransformerEmbedder

allennlp/modules/token_embedders/pretrained_transformer_embedder.py

@@ -0,0 +1,17 @@
+from pytorch_transformers.modeling_auto import AutoModel
+import torch
+
+from allennlp.modules.token_embedders.token_embedder import TokenEmbedder
+
+
+@TokenEmbedder.register("pretrained_transformer")
+class PretrainedTransformerEmbedder(TokenEmbedder):
+    """
+    Uses a pretrained model from ``pytorch-transformers`` as a ``TokenEmbedder``.
+    """
+    def __init__(self, model_name: str) -> None:
+        super().__init__()
+        self.transformer_model = AutoModel.from_pretrained(model_name)
+
+    def forward(self, token_ids: torch.LongTensor) -> torch.Tensor:
+        return self.transformer_model(token_ids)[0]

allennlp/tests/modules/token_embedders/pretrained_transformer_embedder_test.py

@@ -0,0 +1,16 @@
+# pylint: disable=no-self-use,invalid-name
+import torch
+
+from allennlp.common import Params
+from allennlp.modules.token_embedders import PretrainedTransformerEmbedder
+from allennlp.common.testing import AllenNlpTestCase
+
+class TestPretrainedTransformerEmbedder(AllenNlpTestCase):
+    def test_forward_runs_when_initialized_from_params(self):
+        # This code just passes things off to pytorch-transformers, so we only have a very simple
+        # test.
+        params = Params({'model_name': 'bert-base-uncased'})
+        embedder = PretrainedTransformerEmbedder.from_params(params)
+        tensor = torch.randint(0, 100, (1, 4))
+        output = embedder(tensor)
+        assert tuple(output.size()) == (1, 4, 768)