test_modeling_distilbert.py 10.3 KB
Newer Older
LysandreJik's avatar
LysandreJik 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

LysandreJik's avatar
LysandreJik committed
16

17
18
import unittest

19
from transformers import is_torch_available
thomwolf's avatar
thomwolf committed
20

21
from .test_configuration_common import ConfigTester
22
from .test_modeling_common import ModelTesterMixin, ids_tensor
23
from .utils import require_torch, torch_device
Aymeric Augustin's avatar
Aymeric Augustin committed
24
25


26
if is_torch_available():
27
28
29
30
31
32
33
34
35
    from transformers import (
        DistilBertConfig,
        DistilBertModel,
        DistilBertForMaskedLM,
        DistilBertForTokenClassification,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
    )

LysandreJik's avatar
LysandreJik committed
36

37
@require_torch
38
class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):
LysandreJik's avatar
LysandreJik committed
39

40
    all_model_classes = (
41
42
43
44
45
46
47
        (
            DistilBertModel,
            DistilBertForMaskedLM,
            DistilBertForQuestionAnswering,
            DistilBertForSequenceClassification,
            DistilBertForTokenClassification,
        )
48
49
50
        if is_torch_available()
        else None
    )
51
52
53
54
    test_pruning = True
    test_torchscript = True
    test_resize_embeddings = True
    test_head_masking = True
LysandreJik's avatar
LysandreJik committed
55

thomwolf's avatar
thomwolf committed
56
    class DistilBertModelTester(object):
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
        def __init__(
            self,
            parent,
            batch_size=13,
            seq_length=7,
            is_training=True,
            use_input_mask=True,
            use_token_type_ids=False,
            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,
        ):
LysandreJik's avatar
LysandreJik committed
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
113
114
115
116
117
118
119
            self.parent = parent
            self.batch_size = batch_size
            self.seq_length = seq_length
            self.is_training = is_training
            self.use_input_mask = use_input_mask
            self.use_token_type_ids = use_token_type_ids
            self.use_labels = use_labels
            self.vocab_size = vocab_size
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.num_attention_heads = num_attention_heads
            self.intermediate_size = intermediate_size
            self.hidden_act = hidden_act
            self.hidden_dropout_prob = hidden_dropout_prob
            self.attention_probs_dropout_prob = attention_probs_dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.type_sequence_label_size = type_sequence_label_size
            self.initializer_range = initializer_range
            self.num_labels = num_labels
            self.num_choices = num_choices
            self.scope = scope

        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)

            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
120
            config = DistilBertConfig(
thomwolf's avatar
thomwolf committed
121
                vocab_size=self.vocab_size,
LysandreJik's avatar
LysandreJik committed
122
123
124
125
126
127
128
129
                dim=self.hidden_size,
                n_layers=self.num_hidden_layers,
                n_heads=self.num_attention_heads,
                hidden_dim=self.intermediate_size,
                hidden_act=self.hidden_act,
                dropout=self.hidden_dropout_prob,
                attention_dropout=self.attention_probs_dropout_prob,
                max_position_embeddings=self.max_position_embeddings,
130
131
                initializer_range=self.initializer_range,
            )
LysandreJik's avatar
LysandreJik committed
132
133
134
135

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

        def check_loss_output(self, result):
136
            self.parent.assertListEqual(list(result["loss"].size()), [])
LysandreJik's avatar
LysandreJik committed
137

138
139
140
        def create_and_check_distilbert_model(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
thomwolf's avatar
thomwolf committed
141
            model = DistilBertModel(config=config)
142
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
143
            model.eval()
144
145
            (sequence_output,) = model(input_ids, input_mask)
            (sequence_output,) = model(input_ids)
LysandreJik's avatar
LysandreJik committed
146
147
148
149
150

            result = {
                "sequence_output": sequence_output,
            }
            self.parent.assertListEqual(
151
152
                list(result["sequence_output"].size()), [self.batch_size, self.seq_length, self.hidden_size]
            )
LysandreJik's avatar
LysandreJik committed
153

154
155
156
        def create_and_check_distilbert_for_masked_lm(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
thomwolf's avatar
thomwolf committed
157
            model = DistilBertForMaskedLM(config=config)
158
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
159
            model.eval()
160
            loss, prediction_scores = model(input_ids, attention_mask=input_mask, labels=token_labels)
LysandreJik's avatar
LysandreJik committed
161
162
163
164
165
            result = {
                "loss": loss,
                "prediction_scores": prediction_scores,
            }
            self.parent.assertListEqual(
166
167
                list(result["prediction_scores"].size()), [self.batch_size, self.seq_length, self.vocab_size]
            )
LysandreJik's avatar
LysandreJik committed
168
169
            self.check_loss_output(result)

170
171
172
        def create_and_check_distilbert_for_question_answering(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
thomwolf's avatar
thomwolf committed
173
            model = DistilBertForQuestionAnswering(config=config)
174
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
175
            model.eval()
176
177
178
            loss, start_logits, end_logits = model(
                input_ids, attention_mask=input_mask, start_positions=sequence_labels, end_positions=sequence_labels
            )
LysandreJik's avatar
LysandreJik committed
179
180
181
182
183
            result = {
                "loss": loss,
                "start_logits": start_logits,
                "end_logits": end_logits,
            }
184
185
            self.parent.assertListEqual(list(result["start_logits"].size()), [self.batch_size, self.seq_length])
            self.parent.assertListEqual(list(result["end_logits"].size()), [self.batch_size, self.seq_length])
LysandreJik's avatar
LysandreJik committed
186
187
            self.check_loss_output(result)

188
189
190
        def create_and_check_distilbert_for_sequence_classification(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
LysandreJik's avatar
LysandreJik committed
191
            config.num_labels = self.num_labels
thomwolf's avatar
thomwolf committed
192
            model = DistilBertForSequenceClassification(config)
193
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
194
            model.eval()
195
            loss, logits = model(input_ids, attention_mask=input_mask, labels=sequence_labels)
LysandreJik's avatar
LysandreJik committed
196
197
198
199
            result = {
                "loss": loss,
                "logits": logits,
            }
200
            self.parent.assertListEqual(list(result["logits"].size()), [self.batch_size, self.num_labels])
LysandreJik's avatar
LysandreJik committed
201
202
            self.check_loss_output(result)

203
204
205
        def create_and_check_distilbert_for_token_classification(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
206
207
            config.num_labels = self.num_labels
            model = DistilBertForTokenClassification(config=config)
208
            model.to(torch_device)
209
210
211
212
213
214
215
216
            model.eval()

            loss, logits = model(input_ids, attention_mask=input_mask, labels=token_labels)
            result = {
                "loss": loss,
                "logits": logits,
            }
            self.parent.assertListEqual(
217
218
                list(result["logits"].size()), [self.batch_size, self.seq_length, self.num_labels]
            )
219
220
            self.check_loss_output(result)

LysandreJik's avatar
LysandreJik committed
221
222
223
        def prepare_config_and_inputs_for_common(self):
            config_and_inputs = self.prepare_config_and_inputs()
            (config, input_ids, input_mask, sequence_labels, token_labels, choice_labels) = config_and_inputs
224
            inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
LysandreJik's avatar
LysandreJik committed
225
226
227
            return config, inputs_dict

    def setUp(self):
thomwolf's avatar
thomwolf committed
228
229
        self.model_tester = DistilBertModelTest.DistilBertModelTester(self)
        self.config_tester = ConfigTester(self, config_class=DistilBertConfig, dim=37)
LysandreJik's avatar
LysandreJik committed
230
231
232
233

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

thomwolf's avatar
thomwolf committed
234
    def test_distilbert_model(self):
LysandreJik's avatar
LysandreJik committed
235
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
236
        self.model_tester.create_and_check_distilbert_model(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
237
238
239

    def test_for_masked_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
240
        self.model_tester.create_and_check_distilbert_for_masked_lm(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
241
242
243

    def test_for_question_answering(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
244
        self.model_tester.create_and_check_distilbert_for_question_answering(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
245
246
247

    def test_for_sequence_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
248
        self.model_tester.create_and_check_distilbert_for_sequence_classification(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
249

250
251
252
253
    def test_for_token_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_token_classification(*config_and_inputs)

254
    # @slow
LysandreJik's avatar
LysandreJik committed
255
    # def test_model_from_pretrained(self):
256
    #     for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
257
    #         model = DistilBertModel.from_pretrained(model_name)
LysandreJik's avatar
LysandreJik committed
258
    #         self.assertIsNotNone(model)