test_modeling_ctrl.py 9.14 KB
Newer Older
keskarnitish's avatar
keskarnitish committed
1
2
3
4
5
6
7
8
9
10
11
12
13
# coding=utf-8
# Copyright 2018 Salesforce and HuggingFace Inc. team.
# 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
14

keskarnitish's avatar
keskarnitish committed
15

16
17
import unittest

keskarnitish's avatar
keskarnitish committed
18
19
from transformers import is_torch_available

20
from .test_configuration_common import ConfigTester
21
from .test_modeling_common import ModelTesterMixin, ids_tensor
22
from .utils import require_torch, slow, torch_device
keskarnitish's avatar
keskarnitish committed
23
24


Aymeric Augustin's avatar
Aymeric Augustin committed
25
if is_torch_available():
26
    import torch
Aymeric Augustin's avatar
Aymeric Augustin committed
27
28
29
    from transformers import CTRLConfig, CTRLModel, CTRL_PRETRAINED_MODEL_ARCHIVE_MAP, CTRLLMHeadModel


30
@require_torch
31
class CTRLModelTest(ModelTesterMixin, unittest.TestCase):
keskarnitish's avatar
keskarnitish committed
32
33

    all_model_classes = (CTRLModel, CTRLLMHeadModel) if is_torch_available() else ()
34
    all_generative_model_classes = (CTRLLMHeadModel,) if is_torch_available() else ()
keskarnitish's avatar
keskarnitish committed
35
36
37
38
39
40
    test_pruning = False
    test_torchscript = False
    test_resize_embeddings = False
    test_head_masking = False

    class CTRLModelTester(object):
41
42
43
        def __init__(
            self,
            parent,
44
            batch_size=14,
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
            seq_length=7,
            is_training=True,
            use_token_type_ids=True,
            use_input_mask=True,
            use_labels=True,
            use_mc_token_ids=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,
        ):
keskarnitish's avatar
keskarnitish committed
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
113
114
            self.parent = parent
            self.batch_size = batch_size
            self.seq_length = seq_length
            self.is_training = is_training
            self.use_token_type_ids = use_token_type_ids
            self.use_input_mask = use_input_mask
            self.use_labels = use_labels
            self.use_mc_token_ids = use_mc_token_ids
            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)

            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)

            mc_token_ids = None
            if self.use_mc_token_ids:
                mc_token_ids = ids_tensor([self.batch_size, self.num_choices], self.seq_length)

            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 = CTRLConfig(
thomwolf's avatar
thomwolf committed
115
                vocab_size=self.vocab_size,
keskarnitish's avatar
keskarnitish committed
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
                n_embd=self.hidden_size,
                n_layer=self.num_hidden_layers,
                n_head=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,
                n_positions=self.max_position_embeddings,
                n_ctx=self.max_position_embeddings
                # type_vocab_size=self.type_vocab_size,
                # initializer_range=self.initializer_range
            )

            head_mask = ids_tensor([self.num_hidden_layers, self.num_attention_heads], 2)

131
132
133
134
135
136
137
138
139
140
141
            return (
                config,
                input_ids,
                input_mask,
                head_mask,
                token_type_ids,
                mc_token_ids,
                sequence_labels,
                token_labels,
                choice_labels,
            )
keskarnitish's avatar
keskarnitish committed
142
143

        def check_loss_output(self, result):
144
            self.parent.assertListEqual(list(result["loss"].size()), [])
keskarnitish's avatar
keskarnitish committed
145
146
147

        def create_and_check_ctrl_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
            model = CTRLModel(config=config)
148
            model.to(torch_device)
keskarnitish's avatar
keskarnitish committed
149
150
151
152
            model.eval()

            model(input_ids, token_type_ids=token_type_ids, head_mask=head_mask)
            model(input_ids, token_type_ids=token_type_ids)
thomwolf's avatar
thomwolf committed
153
            sequence_output, presents = model(input_ids)
keskarnitish's avatar
keskarnitish committed
154
155
156

            result = {
                "sequence_output": sequence_output,
thomwolf's avatar
thomwolf committed
157
                "presents": presents,
keskarnitish's avatar
keskarnitish committed
158
159
            }
            self.parent.assertListEqual(
160
161
                list(result["sequence_output"].size()), [self.batch_size, self.seq_length, self.hidden_size]
            )
thomwolf's avatar
thomwolf committed
162
            self.parent.assertEqual(len(result["presents"]), config.n_layer)
keskarnitish's avatar
keskarnitish committed
163
164
165

        def create_and_check_lm_head_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
            model = CTRLLMHeadModel(config)
166
            model.to(torch_device)
keskarnitish's avatar
keskarnitish committed
167
168
169
170
            model.eval()

            loss, lm_logits, _ = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)

171
172
            result = {"loss": loss, "lm_logits": lm_logits}
            self.parent.assertListEqual(list(result["loss"].size()), [])
keskarnitish's avatar
keskarnitish committed
173
            self.parent.assertListEqual(
174
175
                list(result["lm_logits"].size()), [self.batch_size, self.seq_length, self.vocab_size]
            )
keskarnitish's avatar
keskarnitish committed
176
177
178
179

        def prepare_config_and_inputs_for_common(self):
            config_and_inputs = self.prepare_config_and_inputs()

180
181
182
183
184
185
186
187
188
189
190
191
192
            (
                config,
                input_ids,
                input_mask,
                head_mask,
                token_type_ids,
                mc_token_ids,
                sequence_labels,
                token_labels,
                choice_labels,
            ) = config_and_inputs

            inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "head_mask": head_mask}
keskarnitish's avatar
keskarnitish committed
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210

            return config, inputs_dict

    def setUp(self):
        self.model_tester = CTRLModelTest.CTRLModelTester(self)
        self.config_tester = ConfigTester(self, config_class=CTRLConfig, n_embd=37)

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

    def test_ctrl_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_model(*config_and_inputs)

    def test_ctrl_lm_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_lm_head_model(*config_and_inputs)

211
    @slow
keskarnitish's avatar
keskarnitish committed
212
213
    def test_model_from_pretrained(self):
        for model_name in list(CTRL_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]:
214
            model = CTRLModel.from_pretrained(model_name)
keskarnitish's avatar
keskarnitish committed
215
            self.assertIsNotNone(model)
216
217


218
@require_torch
219
220
221
222
class CTRLModelLanguageGenerationTest(unittest.TestCase):
    @slow
    def test_lm_generate_ctrl(self):
        model = CTRLLMHeadModel.from_pretrained("ctrl")
223
        model.to(torch_device)
patrickvonplaten's avatar
patrickvonplaten committed
224
        input_ids = torch.tensor(
Patrick von Platen's avatar
Patrick von Platen committed
225
226
            [[11859, 0, 1611, 8]], dtype=torch.long, device=torch_device
        )  # Legal the president is
227
228
        expected_output_ids = [
            11859,
Patrick von Platen's avatar
Patrick von Platen committed
229
230
            0,
            1611,
231
            8,
Patrick von Platen's avatar
Patrick von Platen committed
232
233
234
            5,
            150,
            26449,
235
            2,
Patrick von Platen's avatar
Patrick von Platen committed
236
237
238
            19,
            348,
            469,
239
            3,
Patrick von Platen's avatar
Patrick von Platen committed
240
241
242
243
244
245
246
247
248
            2595,
            48,
            20740,
            246533,
            246533,
            19,
            30,
            5,
        ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a
249

patrickvonplaten's avatar
patrickvonplaten committed
250
        output_ids = model.generate(input_ids, do_sample=False)
Patrick von Platen's avatar
Patrick von Platen committed
251
        self.assertListEqual(output_ids[0].tolist(), expected_output_ids)