From 2184f87003c18ad8a172ecab9a821626522cf8e7 Mon Sep 17 00:00:00 2001 From: Lysandre Date: Mon, 3 Feb 2020 18:39:54 -0500 Subject: [PATCH] RoBERTa TensorFlow Tests --- tests/test_modeling_tf_roberta.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_modeling_tf_roberta.py b/tests/test_modeling_tf_roberta.py index 21b0ffee0e..5ee598b2dc 100644 --- a/tests/test_modeling_tf_roberta.py +++ b/tests/test_modeling_tf_roberta.py @@ -200,6 +200,10 @@ class TFRobertaModelTest(TFModelTesterMixin, unittest.TestCase): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_roberta_for_masked_lm(*config_and_inputs) + def test_for_token_classification(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_roberta_for_token_classification(*config_and_inputs) + @slow def test_model_from_pretrained(self): for model_name in list(TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]: