test_modeling_tf_bert.py 12.9 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# 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.
Aymeric Augustin's avatar
Aymeric Augustin committed
15

thomwolf's avatar
thomwolf committed
16

17
18
import unittest

Aymeric Augustin's avatar
Aymeric Augustin committed
19
from transformers import BertConfig, is_tf_available
20
from transformers.testing_utils import require_tf, slow
thomwolf's avatar
thomwolf committed
21

22
from .test_configuration_common import ConfigTester
23
from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
thomwolf's avatar
thomwolf committed
24

thomwolf's avatar
thomwolf committed
25

thomwolf's avatar
thomwolf committed
26
if is_tf_available():
thomwolf's avatar
thomwolf committed
27
    import tensorflow as tf
28
29
    from transformers.modeling_tf_bert import (
        TFBertModel,
30
        TFBertLMHeadModel,
31
32
33
34
35
36
37
38
        TFBertForMaskedLM,
        TFBertForNextSentencePrediction,
        TFBertForPreTraining,
        TFBertForSequenceClassification,
        TFBertForMultipleChoice,
        TFBertForTokenClassification,
        TFBertForQuestionAnswering,
    )
thomwolf's avatar
thomwolf committed
39

thomwolf's avatar
thomwolf committed
40

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class TFBertModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=7,
        is_training=True,
        use_input_mask=True,
        use_token_type_ids=True,
        use_labels=True,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=5,
        num_attention_heads=4,
        intermediate_size=37,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=512,
        type_vocab_size=16,
        type_sequence_label_size=2,
        initializer_range=0.02,
        num_labels=3,
        num_choices=4,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = 13
        self.seq_length = 7
        self.is_training = True
        self.use_input_mask = True
        self.use_token_type_ids = True
        self.use_labels = True
        self.vocab_size = 99
        self.hidden_size = 32
        self.num_hidden_layers = 5
        self.num_attention_heads = 4
        self.intermediate_size = 37
        self.hidden_act = "gelu"
        self.hidden_dropout_prob = 0.1
        self.attention_probs_dropout_prob = 0.1
        self.max_position_embeddings = 512
        self.type_vocab_size = 16
        self.type_sequence_label_size = 2
        self.initializer_range = 0.02
        self.num_labels = 3
        self.num_choices = 4
        self.scope = None
thomwolf's avatar
thomwolf committed
89

90
91
92
93
94
95
    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)

        input_mask = None
        if self.use_input_mask:
            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)
thomwolf's avatar
thomwolf committed
96

97
98
99
        token_type_ids = None
        if self.use_token_type_ids:
            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
thomwolf's avatar
thomwolf committed
100

101
102
103
104
105
106
107
        sequence_labels = None
        token_labels = None
        choice_labels = None
        if self.use_labels:
            sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
            token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)
            choice_labels = ids_tensor([self.batch_size], self.num_choices)
thomwolf's avatar
thomwolf committed
108

109
110
111
112
113
114
115
116
117
118
119
120
        config = BertConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            intermediate_size=self.intermediate_size,
            hidden_act=self.hidden_act,
            hidden_dropout_prob=self.hidden_dropout_prob,
            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
            max_position_embeddings=self.max_position_embeddings,
            type_vocab_size=self.type_vocab_size,
            initializer_range=self.initializer_range,
Sylvain Gugger's avatar
Sylvain Gugger committed
121
            return_dict=True,
122
        )
thomwolf's avatar
thomwolf committed
123

124
        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
thomwolf's avatar
thomwolf committed
125

126
127
128
129
130
131
    def create_and_check_bert_model(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFBertModel(config=config)
        inputs = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        sequence_output, pooled_output = model(inputs)
thomwolf's avatar
thomwolf committed
132

133
        inputs = [input_ids, input_mask]
Sylvain Gugger's avatar
Sylvain Gugger committed
134
        result = model(inputs)
thomwolf's avatar
thomwolf committed
135

Sylvain Gugger's avatar
Sylvain Gugger committed
136
        result = model(input_ids)
thomwolf's avatar
thomwolf committed
137

138
        self.parent.assertListEqual(
Sylvain Gugger's avatar
Sylvain Gugger committed
139
            list(result["last_hidden_state"].shape), [self.batch_size, self.seq_length, self.hidden_size]
140
        )
Sylvain Gugger's avatar
Sylvain Gugger committed
141
        self.parent.assertListEqual(list(result["pooler_output"].shape), [self.batch_size, self.hidden_size])
thomwolf's avatar
thomwolf committed
142

143
144
145
146
147
148
149
150
151
152
    def create_and_check_bert_lm_head(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        config.is_decoder = True
        model = TFBertLMHeadModel(config=config)
        inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "token_type_ids": token_type_ids,
        }
Sylvain Gugger's avatar
Sylvain Gugger committed
153
        prediction_scores = model(inputs)["logits"]
154
155
156
157
        self.parent.assertListEqual(
            list(prediction_scores.numpy().shape), [self.batch_size, self.seq_length, self.vocab_size]
        )

158
159
160
161
    def create_and_check_bert_for_masked_lm(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFBertForMaskedLM(config=config)
162
163
164
165
166
        inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "token_type_ids": token_type_ids,
        }
Sylvain Gugger's avatar
Sylvain Gugger committed
167
168
        result = model(inputs)
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.seq_length, self.vocab_size])
thomwolf's avatar
thomwolf committed
169

170
171
172
173
174
    def create_and_check_bert_for_next_sequence_prediction(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFBertForNextSentencePrediction(config=config)
        inputs = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
175
176
        result = model(inputs)
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, 2])
thomwolf's avatar
thomwolf committed
177

178
179
180
181
182
    def create_and_check_bert_for_pretraining(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFBertForPreTraining(config=config)
        inputs = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
183
        result = model(inputs)
184
        self.parent.assertListEqual(
Sylvain Gugger's avatar
Sylvain Gugger committed
185
            list(result["prediction_logits"].shape), [self.batch_size, self.seq_length, self.vocab_size]
186
        )
Sylvain Gugger's avatar
Sylvain Gugger committed
187
        self.parent.assertListEqual(list(result["seq_relationship_logits"].shape), [self.batch_size, 2])
thomwolf's avatar
thomwolf committed
188

189
190
191
192
193
    def create_and_check_bert_for_sequence_classification(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        config.num_labels = self.num_labels
        model = TFBertForSequenceClassification(config=config)
194
195
196
197
        inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "token_type_ids": token_type_ids,
198
        }
199

Sylvain Gugger's avatar
Sylvain Gugger committed
200
        result = model(inputs)
201
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.num_labels])
thomwolf's avatar
thomwolf committed
202

203
204
205
206
207
208
209
210
211
212
213
214
215
    def create_and_check_bert_for_multiple_choice(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        config.num_choices = self.num_choices
        model = TFBertForMultipleChoice(config=config)
        multiple_choice_inputs_ids = tf.tile(tf.expand_dims(input_ids, 1), (1, self.num_choices, 1))
        multiple_choice_input_mask = tf.tile(tf.expand_dims(input_mask, 1), (1, self.num_choices, 1))
        multiple_choice_token_type_ids = tf.tile(tf.expand_dims(token_type_ids, 1), (1, self.num_choices, 1))
        inputs = {
            "input_ids": multiple_choice_inputs_ids,
            "attention_mask": multiple_choice_input_mask,
            "token_type_ids": multiple_choice_token_type_ids,
        }
Sylvain Gugger's avatar
Sylvain Gugger committed
216
        result = model(inputs)
217
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.num_choices])
thomwolf's avatar
thomwolf committed
218

219
220
221
222
223
    def create_and_check_bert_for_token_classification(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        config.num_labels = self.num_labels
        model = TFBertForTokenClassification(config=config)
224
225
226
227
228
        inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "token_type_ids": token_type_ids,
        }
Sylvain Gugger's avatar
Sylvain Gugger committed
229
        result = model(inputs)
230
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.seq_length, self.num_labels])
thomwolf's avatar
thomwolf committed
231

232
233
234
235
    def create_and_check_bert_for_question_answering(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFBertForQuestionAnswering(config=config)
236
237
238
239
        inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "token_type_ids": token_type_ids,
240
        }
241

Sylvain Gugger's avatar
Sylvain Gugger committed
242
        result = model(inputs)
243
244
        self.parent.assertListEqual(list(result["start_logits"].shape), [self.batch_size, self.seq_length])
        self.parent.assertListEqual(list(result["end_logits"].shape), [self.batch_size, self.seq_length])
thomwolf's avatar
thomwolf committed
245

246
247
248
249
250
251
252
253
254
255
256
257
258
    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
        ) = config_and_inputs
        inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict
thomwolf's avatar
thomwolf committed
259
260


261
262
@require_tf
class TFBertModelTest(TFModelTesterMixin, unittest.TestCase):
thomwolf's avatar
thomwolf committed
263

264
265
266
267
    all_model_classes = (
        (
            TFBertModel,
            TFBertForMaskedLM,
Lysandre Debut's avatar
Lysandre Debut committed
268
            TFBertLMHeadModel,
269
270
271
272
273
274
275
276
277
278
            TFBertForNextSentencePrediction,
            TFBertForPreTraining,
            TFBertForQuestionAnswering,
            TFBertForSequenceClassification,
            TFBertForTokenClassification,
            TFBertForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
thomwolf's avatar
thomwolf committed
279
280

    def setUp(self):
281
        self.model_tester = TFBertModelTester(self)
thomwolf's avatar
thomwolf committed
282
283
284
285
286
287
288
289
290
291
292
293
294
        self.config_tester = ConfigTester(self, config_class=BertConfig, hidden_size=37)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_bert_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_model(*config_and_inputs)

    def test_for_masked_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_for_masked_lm(*config_and_inputs)

295
296
297
298
    def test_for_causal_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_lm_head(*config_and_inputs)

thomwolf's avatar
thomwolf committed
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
    def test_for_multiple_choice(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_for_multiple_choice(*config_and_inputs)

    def test_for_next_sequence_prediction(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_for_next_sequence_prediction(*config_and_inputs)

    def test_for_pretraining(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_for_pretraining(*config_and_inputs)

    def test_for_question_answering(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_for_question_answering(*config_and_inputs)

    def test_for_sequence_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_bert_for_sequence_classification(*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_bert_for_token_classification(*config_and_inputs)

323
    @slow
thomwolf's avatar
thomwolf committed
324
    def test_model_from_pretrained(self):
325
        # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
326
        for model_name in ["bert-base-uncased"]:
327
            model = TFBertModel.from_pretrained(model_name)
thomwolf's avatar
thomwolf committed
328
            self.assertIsNotNone(model)