test_modeling_distilbert.py 10.4 KB
Newer Older
LysandreJik's avatar
LysandreJik committed
1
# coding=utf-8
Sylvain Gugger's avatar
Sylvain Gugger committed
2
# Copyright 2020 The HuggingFace Team. All rights reserved.
LysandreJik's avatar
LysandreJik committed
3
4
5
6
7
8
9
10
11
12
13
14
#
# 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
20
from transformers.testing_utils import require_torch, slow, torch_device
thomwolf's avatar
thomwolf committed
21

22
from .test_configuration_common import ConfigTester
23
from .test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
Aymeric Augustin's avatar
Aymeric Augustin committed
24
25


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

thomwolf's avatar
thomwolf committed
38
    class DistilBertModelTester(object):
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
        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
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
            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:
92
                input_mask = random_attention_mask([self.batch_size, self.seq_length])
LysandreJik's avatar
LysandreJik committed
93
94
95
96
97
98
99
100
101

            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
102
            config = DistilBertConfig(
thomwolf's avatar
thomwolf committed
103
                vocab_size=self.vocab_size,
LysandreJik's avatar
LysandreJik committed
104
105
106
107
108
109
110
111
                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,
112
113
                initializer_range=self.initializer_range,
            )
LysandreJik's avatar
LysandreJik committed
114
115
116

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

117
118
119
        def create_and_check_distilbert_model(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
thomwolf's avatar
thomwolf committed
120
            model = DistilBertModel(config=config)
121
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
122
            model.eval()
Sylvain Gugger's avatar
Sylvain Gugger committed
123
124
            result = model(input_ids, input_mask)
            result = model(input_ids)
Stas Bekman's avatar
Stas Bekman committed
125
126
            self.parent.assertEqual(
                result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)
127
            )
LysandreJik's avatar
LysandreJik committed
128

129
130
131
        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
132
            model = DistilBertForMaskedLM(config=config)
133
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
134
            model.eval()
Sylvain Gugger's avatar
Sylvain Gugger committed
135
            result = model(input_ids, attention_mask=input_mask, labels=token_labels)
Stas Bekman's avatar
Stas Bekman committed
136
            self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
LysandreJik's avatar
LysandreJik committed
137

138
139
140
        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
141
            model = DistilBertForQuestionAnswering(config=config)
142
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
143
            model.eval()
Sylvain Gugger's avatar
Sylvain Gugger committed
144
            result = model(
145
146
                input_ids, attention_mask=input_mask, start_positions=sequence_labels, end_positions=sequence_labels
            )
Stas Bekman's avatar
Stas Bekman committed
147
148
            self.parent.assertEqual(result.start_logits.shape, (self.batch_size, self.seq_length))
            self.parent.assertEqual(result.end_logits.shape, (self.batch_size, self.seq_length))
LysandreJik's avatar
LysandreJik committed
149

150
151
152
        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
153
            config.num_labels = self.num_labels
thomwolf's avatar
thomwolf committed
154
            model = DistilBertForSequenceClassification(config)
155
            model.to(torch_device)
LysandreJik's avatar
LysandreJik committed
156
            model.eval()
Sylvain Gugger's avatar
Sylvain Gugger committed
157
            result = model(input_ids, attention_mask=input_mask, labels=sequence_labels)
Stas Bekman's avatar
Stas Bekman committed
158
            self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))
LysandreJik's avatar
LysandreJik committed
159

160
161
162
        def create_and_check_distilbert_for_token_classification(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
163
164
            config.num_labels = self.num_labels
            model = DistilBertForTokenClassification(config=config)
165
            model.to(torch_device)
166
167
            model.eval()

Sylvain Gugger's avatar
Sylvain Gugger committed
168
            result = model(input_ids, attention_mask=input_mask, labels=token_labels)
Stas Bekman's avatar
Stas Bekman committed
169
            self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))
170

171
172
173
174
175
176
177
178
179
        def create_and_check_distilbert_for_multiple_choice(
            self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
        ):
            config.num_choices = self.num_choices
            model = DistilBertForMultipleChoice(config=config)
            model.to(torch_device)
            model.eval()
            multiple_choice_inputs_ids = input_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
            multiple_choice_input_mask = input_mask.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
Sylvain Gugger's avatar
Sylvain Gugger committed
180
            result = model(
Lysandre's avatar
Lysandre committed
181
182
183
                multiple_choice_inputs_ids,
                attention_mask=multiple_choice_input_mask,
                labels=choice_labels,
184
            )
Stas Bekman's avatar
Stas Bekman committed
185
            self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_choices))
186

LysandreJik's avatar
LysandreJik committed
187
188
189
        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
190
            inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
LysandreJik's avatar
LysandreJik committed
191
192
            return config, inputs_dict

193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213

@require_torch
class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):

    all_model_classes = (
        (
            DistilBertModel,
            DistilBertForMaskedLM,
            DistilBertForMultipleChoice,
            DistilBertForQuestionAnswering,
            DistilBertForSequenceClassification,
            DistilBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    test_pruning = True
    test_torchscript = True
    test_resize_embeddings = True
    test_head_masking = True

LysandreJik's avatar
LysandreJik committed
214
    def setUp(self):
215
        self.model_tester = DistilBertModelTester(self)
thomwolf's avatar
thomwolf committed
216
        self.config_tester = ConfigTester(self, config_class=DistilBertConfig, dim=37)
LysandreJik's avatar
LysandreJik committed
217
218
219
220

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

thomwolf's avatar
thomwolf committed
221
    def test_distilbert_model(self):
LysandreJik's avatar
LysandreJik committed
222
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
223
        self.model_tester.create_and_check_distilbert_model(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
224
225
226

    def test_for_masked_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
227
        self.model_tester.create_and_check_distilbert_for_masked_lm(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
228
229
230

    def test_for_question_answering(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
thomwolf's avatar
thomwolf committed
231
        self.model_tester.create_and_check_distilbert_for_question_answering(*config_and_inputs)
LysandreJik's avatar
LysandreJik committed
232
233
234

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

237
238
239
240
    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)

241
242
243
244
    def test_for_multiple_choice(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_multiple_choice(*config_and_inputs)

245
246
247
248
249
    @slow
    def test_model_from_pretrained(self):
        for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = DistilBertModel.from_pretrained(model_name)
            self.assertIsNotNone(model)