test_modeling_tf_albert.py 11.4 KB
Newer Older
Lysandre's avatar
Lysandre 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

Lysandre's avatar
Lysandre committed
16

17
18
import unittest

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

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


if is_tf_available():
27
    import tensorflow as tf
28
29
    from transformers.modeling_tf_albert import (
        TFAlbertModel,
30
        TFAlbertForPreTraining,
31
        TFAlbertForMaskedLM,
32
        TFAlbertForMultipleChoice,
33
        TFAlbertForSequenceClassification,
34
        TFAlbertForQuestionAnswering,
Lysandre Debut's avatar
Lysandre Debut committed
35
        TFAlbertForTokenClassification,
36
        TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
37
    )
Lysandre's avatar
Lysandre committed
38
39


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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
class TFAlbertModelTester:
    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,
        embedding_size=16,
        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.embedding_size = 16
        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

    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)

        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)

        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)

        config = AlbertConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
Lysandre Debut's avatar
Lysandre Debut committed
113
            embedding_size=self.embedding_size,
114
115
116
117
118
119
120
121
122
            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
123
            return_dict=True,
124
125
126
127
128
129
130
131
132
133
134
135
136
        )

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

    def create_and_check_albert_model(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFAlbertModel(config=config)
        # inputs = {'input_ids': input_ids,
        #           'attention_mask': input_mask,
        #           'token_type_ids': token_type_ids}
        # sequence_output, pooled_output = model(**inputs)
        inputs = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
137
        result = model(inputs)
138
139

        inputs = [input_ids, input_mask]
Sylvain Gugger's avatar
Sylvain Gugger committed
140
        result = model(inputs)
141

Sylvain Gugger's avatar
Sylvain Gugger committed
142
        result = model(input_ids)
143
144

        self.parent.assertListEqual(
Sylvain Gugger's avatar
Sylvain Gugger committed
145
            list(result["last_hidden_state"].shape), [self.batch_size, self.seq_length, self.hidden_size]
146
        )
Sylvain Gugger's avatar
Sylvain Gugger committed
147
        self.parent.assertListEqual(list(result["pooler_output"].shape), [self.batch_size, self.hidden_size])
148
149
150
151
152
153
154

    def create_and_check_albert_for_pretraining(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        config.num_labels = self.num_labels
        model = TFAlbertForPreTraining(config=config)
        inputs = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
155
        result = model(inputs)
156
        self.parent.assertListEqual(
Sylvain Gugger's avatar
Sylvain Gugger committed
157
            list(result["prediction_logits"].shape), [self.batch_size, self.seq_length, self.vocab_size]
158
        )
Sylvain Gugger's avatar
Sylvain Gugger committed
159
        self.parent.assertListEqual(list(result["sop_logits"].shape), [self.batch_size, self.num_labels])
160
161
162
163
164
165

    def create_and_check_albert_for_masked_lm(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFAlbertForMaskedLM(config=config)
        inputs = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
166
167
        result = model(inputs)
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.seq_length, self.vocab_size])
168
169
170
171
172
173
174

    def create_and_check_albert_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 = TFAlbertForSequenceClassification(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
        result = model(inputs)
176
177
178
179
180
181
182
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.num_labels])

    def create_and_check_albert_for_question_answering(
        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    ):
        model = TFAlbertForQuestionAnswering(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
185
186
        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])

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
    def create_and_check_albert_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 = TFAlbertForMultipleChoice(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,
        }
        result = model(inputs)
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.num_choices])

Lysandre Debut's avatar
Lysandre Debut committed
203
204
205
206
207
208
209
210
211
212
213
214
215
    def create_and_check_albert_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 = TFAlbertForTokenClassification(config=config)
        inputs = {
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "token_type_ids": token_type_ids,
        }
        result = model(inputs)
        self.parent.assertListEqual(list(result["logits"].shape), [self.batch_size, self.seq_length, self.num_labels])

216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
    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


231
@require_tf
232
class TFAlbertModelTest(TFModelTesterMixin, unittest.TestCase):
Lysandre's avatar
Lysandre committed
233
234

    all_model_classes = (
235
236
237
238
239
240
        (
            TFAlbertModel,
            TFAlbertForPreTraining,
            TFAlbertForMaskedLM,
            TFAlbertForSequenceClassification,
            TFAlbertForQuestionAnswering,
Lysandre Debut's avatar
Lysandre Debut committed
241
242
            TFAlbertForTokenClassification,
            TFAlbertForMultipleChoice,
243
        )
244
245
        if is_tf_available()
        else ()
246
    )
Lysandre's avatar
Lysandre committed
247
248

    def setUp(self):
249
        self.model_tester = TFAlbertModelTester(self)
250
        self.config_tester = ConfigTester(self, config_class=AlbertConfig, hidden_size=37)
Lysandre's avatar
Lysandre committed
251
252
253
254
255
256
257
258

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

    def test_albert_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_albert_model(*config_and_inputs)

259
260
261
262
    def test_for_pretraining(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_albert_for_pretraining(*config_and_inputs)

Lysandre's avatar
Lysandre committed
263
264
    def test_for_masked_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
265
        self.model_tester.create_and_check_albert_for_masked_lm(*config_and_inputs)
Lysandre's avatar
Lysandre committed
266

267
268
269
270
    def test_for_multiple_choice(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_albert_for_multiple_choice(*config_and_inputs)

Lysandre's avatar
Lysandre committed
271
272
    def test_for_sequence_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
273
        self.model_tester.create_and_check_albert_for_sequence_classification(*config_and_inputs)
Lysandre's avatar
Lysandre committed
274

275
276
277
278
    def test_for_question_answering(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_albert_for_question_answering(*config_and_inputs)

279
    @slow
Lysandre's avatar
Lysandre committed
280
    def test_model_from_pretrained(self):
281
        for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
282
            model = TFAlbertModel.from_pretrained(model_name)
Lysandre's avatar
Lysandre committed
283
            self.assertIsNotNone(model)