From f2d5c04e1f4bcbf5e52e9a87555d9e87cd8f52f9 Mon Sep 17 00:00:00 2001 From: sandip Date: Wed, 3 Feb 2021 20:19:18 +0530 Subject: [PATCH] Added integration tests for TensorFlow implementation of the ALBERT model (#9976) * TF Albert integration test * TF Alber integration test added --- tests/test_modeling_tf_albert.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_modeling_tf_albert.py b/tests/test_modeling_tf_albert.py index e020cc5cf7..19a2c06d26 100644 --- a/tests/test_modeling_tf_albert.py +++ b/tests/test_modeling_tf_albert.py @@ -303,3 +303,26 @@ class TFAlbertModelTest(TFModelTesterMixin, unittest.TestCase): for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]: model = TFAlbertModel.from_pretrained(model_name) self.assertIsNotNone(model) + + +@require_tf +class TFAlbertModelIntegrationTest(unittest.TestCase): + @slow + def test_inference_masked_lm(self): + model = TFAlbertForPreTraining.from_pretrained("albert-base-v2") + input_ids = tf.constant([[0, 1, 2, 3, 4, 5]]) + output = model(input_ids)[0] + + expected_shape = [1, 6, 30000] + self.assertEqual(output.shape, expected_shape) + + expected_slice = tf.constant( + [ + [ + [4.595668, 0.74462754, -1.818147], + [4.5954347, 0.7454184, -1.8188258], + [4.5954905, 0.7448235, -1.8182316], + ] + ] + ) + tf.debugging.assert_near(output[:, :3, :3], expected_slice, atol=1e-4)