test_modeling_ctrl.py 7.38 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
from transformers import is_torch_available
19
from transformers.testing_utils import require_torch, slow, torch_device
keskarnitish's avatar
keskarnitish committed
20

21
from .test_configuration_common import ConfigTester
22
from .test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
keskarnitish's avatar
keskarnitish committed
23
24


Aymeric Augustin's avatar
Aymeric Augustin committed
25
if is_torch_available():
26
    import torch
27
28

    from transformers import CTRL_PRETRAINED_MODEL_ARCHIVE_LIST, CTRLConfig, CTRLLMHeadModel, CTRLModel
Aymeric Augustin's avatar
Aymeric Augustin committed
29
30


31
32
33
34
35
36
37
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
class CTRLModelTester:
    def __init__(
        self, parent,
    ):
        self.parent = parent
        self.batch_size = 14
        self.seq_length = 7
        self.is_training = True
        self.use_token_type_ids = True
        self.use_input_mask = True
        self.use_labels = True
        self.use_mc_token_ids = 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

    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:
64
            input_mask = random_attention_mask([self.batch_size, self.seq_length])
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

        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(
            vocab_size=self.vocab_size,
            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,
Sylvain Gugger's avatar
Sylvain Gugger committed
92
            n_ctx=self.max_position_embeddings,
93
            # type_vocab_size=self.type_vocab_size,
Sylvain Gugger's avatar
Sylvain Gugger committed
94
95
            # initializer_range=self.initializer_range,
            return_dict=True,
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        )

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

        return (
            config,
            input_ids,
            input_mask,
            head_mask,
            token_type_ids,
            mc_token_ids,
            sequence_labels,
            token_labels,
            choice_labels,
        )

    def create_and_check_ctrl_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
        model = CTRLModel(config=config)
        model.to(torch_device)
        model.eval()

        model(input_ids, token_type_ids=token_type_ids, head_mask=head_mask)
        model(input_ids, token_type_ids=token_type_ids)
Sylvain Gugger's avatar
Sylvain Gugger committed
119
        result = model(input_ids)
Stas Bekman's avatar
Stas Bekman committed
120
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
121
        self.parent.assertEqual(len(result.past_key_values), config.n_layer)
122
123
124
125
126
127

    def create_and_check_lm_head_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
        model = CTRLLMHeadModel(config)
        model.to(torch_device)
        model.eval()

Sylvain Gugger's avatar
Sylvain Gugger committed
128
        result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
Stas Bekman's avatar
Stas Bekman committed
129
130
        self.parent.assertEqual(result.loss.shape, ())
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

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

        (
            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}

        return config, inputs_dict


152
@require_torch
153
class CTRLModelTest(ModelTesterMixin, unittest.TestCase):
keskarnitish's avatar
keskarnitish committed
154
155

    all_model_classes = (CTRLModel, CTRLLMHeadModel) if is_torch_available() else ()
156
    all_generative_model_classes = (CTRLLMHeadModel,) if is_torch_available() else ()
157
    test_pruning = True
keskarnitish's avatar
keskarnitish committed
158
159
160
161
162
    test_torchscript = False
    test_resize_embeddings = False
    test_head_masking = False

    def setUp(self):
163
        self.model_tester = CTRLModelTester(self)
keskarnitish's avatar
keskarnitish committed
164
165
166
167
168
169
170
171
172
173
174
175
176
        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)

177
    @slow
keskarnitish's avatar
keskarnitish committed
178
    def test_model_from_pretrained(self):
179
        for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
180
            model = CTRLModel.from_pretrained(model_name)
keskarnitish's avatar
keskarnitish committed
181
            self.assertIsNotNone(model)
182
183


184
@require_torch
185
186
187
188
class CTRLModelLanguageGenerationTest(unittest.TestCase):
    @slow
    def test_lm_generate_ctrl(self):
        model = CTRLLMHeadModel.from_pretrained("ctrl")
189
        model.to(torch_device)
patrickvonplaten's avatar
patrickvonplaten committed
190
        input_ids = torch.tensor(
Patrick von Platen's avatar
Patrick von Platen committed
191
192
            [[11859, 0, 1611, 8]], dtype=torch.long, device=torch_device
        )  # Legal the president is
193
194
        expected_output_ids = [
            11859,
Patrick von Platen's avatar
Patrick von Platen committed
195
196
            0,
            1611,
197
            8,
Patrick von Platen's avatar
Patrick von Platen committed
198
199
200
            5,
            150,
            26449,
201
            2,
Patrick von Platen's avatar
Patrick von Platen committed
202
203
204
            19,
            348,
            469,
205
            3,
Patrick von Platen's avatar
Patrick von Platen committed
206
207
208
209
210
211
212
213
214
            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
215

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