From a88a0e4413d8de5ad235a211fb3b0326aadc5ce0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remilouf@gmail.com>
Date: Wed, 30 Oct 2019 16:06:29 +0100
Subject: [PATCH] add tests to encoder-decoder model

---
 transformers/tests/modeling_common_test.py    | 16 ++++++
 .../tests/modeling_encoder_decoder_test.py    | 52 +++++++++++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 transformers/tests/modeling_encoder_decoder_test.py

diff --git a/transformers/tests/modeling_common_test.py b/transformers/tests/modeling_common_test.py
index 2b66757c28..1c1794550c 100644
--- a/transformers/tests/modeling_common_test.py
+++ b/transformers/tests/modeling_common_test.py
@@ -704,6 +704,22 @@ def ids_tensor(shape, vocab_size, rng=None, name=None):
     return torch.tensor(data=values, dtype=torch.long).view(shape).contiguous()
 
 
+def floats_tensor(shape, scale=1.0, rng=None, name=None):
+    """Creates a random float32 tensor of the shape within the vocab size."""
+    if rng is None:
+        rng = global_rng
+
+    total_dims = 1
+    for dim in shape:
+        total_dims *= dim
+
+    values = []
+    for _ in range(total_dims):
+        values.append(rng.random() * scale)
+
+    return torch.tensor(data=values, dtype=torch.float).view(shape).contiguous()
+
+
 class ModelUtilsTest(unittest.TestCase):
     def test_model_from_pretrained(self):
         logging.basicConfig(level=logging.INFO)
diff --git a/transformers/tests/modeling_encoder_decoder_test.py b/transformers/tests/modeling_encoder_decoder_test.py
new file mode 100644
index 0000000000..1ffd0ebc4c
--- /dev/null
+++ b/transformers/tests/modeling_encoder_decoder_test.py
@@ -0,0 +1,52 @@
+# coding=utf-8
+# Copyright 2018 The Hugging Face Inc. Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import unittest
+import pytest
+
+from transformers import is_torch_available
+
+if is_torch_available():
+    from transformers import BertModel, BertForMaskedLM, Model2Model
+    from transformers.modeling_bert import BERT_PRETRAINED_MODEL_ARCHIVE_MAP
+else:
+    pytestmark = pytest.mark.skip("Require Torch")
+
+
+class EncoderDecoderModelTest(unittest.TestCase):
+    def test_model2model_from_pretrained(self):
+        logging.basicConfig(level=logging.INFO)
+        for model_name in list(BERT_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]:
+            model = Model2Model.from_pretrained(model_name)
+            self.assertIsInstance(model.encoder, BertModel)
+            self.assertIsInstance(model.decoder, BertForMaskedLM)
+            self.assertEqual(model.decoder.config.is_decoder, True)
+            self.assertEqual(model.encoder.config.is_decoder, False)
+
+    def test_model2model_from_pretrained_not_bert(self):
+        logging.basicConfig(level=logging.INFO)
+        with self.assertRaises(ValueError):
+            _ = Model2Model.from_pretrained('roberta')
+
+        with self.assertRaises(ValueError):
+            _ = Model2Model.from_pretrained('distilbert')
+
+        with self.assertRaises(ValueError):
+            _ = Model2Model.from_pretrained('does-not-exist')
+
+
+if __name__ == "__main__":
+    unittest.main()