test_modeling_common.py 34.1 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# coding=utf-8
# Copyright 2019 HuggingFace Inc.
#
# 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

16
import copy
Aymeric Augustin's avatar
Aymeric Augustin committed
17
import logging
18
import os.path
Aymeric Augustin's avatar
Aymeric Augustin committed
19
import random
20
import tempfile
thomwolf's avatar
thomwolf committed
21
22
import unittest

23
from transformers import is_torch_available
24

25
from .utils import require_torch, slow, torch_device
26

Aymeric Augustin's avatar
Aymeric Augustin committed
27

28
if is_torch_available():
thomwolf's avatar
thomwolf committed
29
    import torch
30
    import numpy as np
thomwolf's avatar
thomwolf committed
31

32
33
34
35
36
37
38
    from transformers import (
        AdaptiveEmbedding,
        PretrainedConfig,
        PreTrainedModel,
        BertModel,
        BertConfig,
        BERT_PRETRAINED_MODEL_ARCHIVE_MAP,
39
        top_k_top_p_filtering,
40
    )
thomwolf's avatar
thomwolf committed
41

42

43
44
45
def _config_zero_init(config):
    configs_no_init = copy.deepcopy(config)
    for key in configs_no_init.__dict__.keys():
46
        if "_range" in key or "_std" in key or "initializer_factor" in key:
47
48
49
            setattr(configs_no_init, key, 0.0)
    return configs_no_init

thomwolf's avatar
thomwolf committed
50

51
52
53
54
55
@require_torch
class ModelTesterMixin:

    model_tester = None
    all_model_classes = ()
56
    all_generative_model_classes = ()
Patrick von Platen's avatar
Patrick von Platen committed
57
58
59
60
    test_torchscript = True
    test_pruning = True
    test_resize_embeddings = True
    test_head_masking = True
61
    test_missing_keys = True
62
63
    is_encoder_decoder = False

Patrick von Platen's avatar
Patrick von Platen committed
64
    def test_save_load(self):
65
66
67
68
69
70
71
72
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**inputs_dict)
73
            out_2 = outputs[0].cpu().numpy()
74
            out_2[np.isnan(out_2)] = 0
75

76
            with tempfile.TemporaryDirectory() as tmpdirname:
77
78
                model.save_pretrained(tmpdirname)
                model = model_class.from_pretrained(tmpdirname)
79
                model.to(torch_device)
80
                with torch.no_grad():
81
                    after_outputs = model(**inputs_dict)
thomwolf's avatar
thomwolf committed
82

83
84
85
                # Make sure we don't have nans
                out_1 = after_outputs[0].cpu().numpy()
                out_1[np.isnan(out_1)] = 0
thomwolf's avatar
thomwolf committed
86
87
                max_diff = np.amax(np.abs(out_1 - out_2))
                self.assertLessEqual(max_diff, 1e-5)
88

Patrick von Platen's avatar
Patrick von Platen committed
89
    def test_initialization(self):
90
91
92
93
94
95
96
97
98
99
100
101
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                if param.requires_grad:
                    self.assertIn(
                        param.data.mean().item(),
                        [0.0, 1.0],
                        msg="Parameter {} of model {} seems not properly initialized".format(name, model_class),
                    )
thomwolf's avatar
thomwolf committed
102

Patrick von Platen's avatar
Patrick von Platen committed
103
    def test_determinism(self):
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                first = model(**inputs_dict)[0]
                second = model(**inputs_dict)[0]
            out_1 = first.cpu().numpy()
            out_2 = second.cpu().numpy()
            out_1 = out_1[~np.isnan(out_1)]
            out_2 = out_2[~np.isnan(out_2)]
            max_diff = np.amax(np.abs(out_1 - out_2))
            self.assertLessEqual(max_diff, 1e-5)

Patrick von Platen's avatar
Patrick von Platen committed
120
    def test_attention_outputs(self):
121
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
sshleifer's avatar
sshleifer committed
122
        seq_len = getattr(self.model_tester, "seq_length", None)
sshleifer's avatar
sshleifer committed
123
124
        decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len)
        encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", seq_len)
125
126
        decoder_key_length = getattr(self.model_tester, "key_length", decoder_seq_length)
        encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

        for model_class in self.all_model_classes:
            config.output_attentions = True
            config.output_hidden_states = False
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**inputs_dict)
            attentions = outputs[-1]
            self.assertEqual(model.config.output_attentions, True)
            self.assertEqual(model.config.output_hidden_states, False)
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],
143
            )
144
            out_len = len(outputs)
thomwolf's avatar
thomwolf committed
145

146
            if self.is_encoder_decoder:
Sam Shleifer's avatar
Sam Shleifer committed
147
148
149
150
                correct_outlen = (
                    4  # decoder_features_or_logits, decoder_attentions, encoder_features, encoder_attentions
                )
                decoder_attention_idx = 1
151
                if "lm_labels" in inputs_dict:  # loss will come first
Sam Shleifer's avatar
Sam Shleifer committed
152
153
154
155
156
157
                    correct_outlen += 1  # compute loss
                    decoder_attention_idx += 1
                self.assertEqual(out_len, correct_outlen)

                decoder_attentions = outputs[decoder_attention_idx]
                self.assertIsInstance(decoder_attentions, (list, tuple))
158
                self.assertEqual(len(decoder_attentions), self.model_tester.num_hidden_layers)
thomwolf's avatar
thomwolf committed
159
                self.assertListEqual(
160
161
                    list(decoder_attentions[0].shape[-3:]),
                    [self.model_tester.num_attention_heads, decoder_seq_length, decoder_key_length],
162
                )
thomwolf's avatar
thomwolf committed
163

164
            # Check attention is always last and order is fine
thomwolf's avatar
thomwolf committed
165
166
            config.output_attentions = True
            config.output_hidden_states = True
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**inputs_dict)
            self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1), len(outputs))
            self.assertEqual(model.config.output_attentions, True)
            self.assertEqual(model.config.output_hidden_states, True)

            self_attentions = outputs[-1]
            self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],
            )
thomwolf's avatar
thomwolf committed
182

Patrick von Platen's avatar
Patrick von Platen committed
183
    def test_torchscript(self):
184
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
thomwolf's avatar
thomwolf committed
185

186
        self._create_and_check_torchscript(config, inputs_dict)
thomwolf's avatar
thomwolf committed
187

Patrick von Platen's avatar
Patrick von Platen committed
188
    def test_torchscript_output_attentions(self):
189
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
thomwolf's avatar
thomwolf committed
190

191
192
        config.output_attentions = True
        self._create_and_check_torchscript(config, inputs_dict)
thomwolf's avatar
thomwolf committed
193

Patrick von Platen's avatar
Patrick von Platen committed
194
    def test_torchscript_output_hidden_state(self):
195
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
196

197
198
        config.output_hidden_states = True
        self._create_and_check_torchscript(config, inputs_dict)
thomwolf's avatar
thomwolf committed
199

200
    def _create_and_check_torchscript(self, config, inputs_dict):
Patrick von Platen's avatar
Patrick von Platen committed
201
        if not self.test_torchscript:
202
            return
203

204
205
206
207
208
209
210
        configs_no_init = _config_zero_init(config)  # To be sure we have no Nan
        configs_no_init.torchscript = True
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            model.to(torch_device)
            model.eval()
            inputs = inputs_dict["input_ids"]  # Let's keep only input_ids
thomwolf's avatar
thomwolf committed
211

212
213
214
215
            try:
                traced_gpt2 = torch.jit.trace(model, inputs)
            except RuntimeError:
                self.fail("Couldn't trace module.")
thomwolf's avatar
thomwolf committed
216

217
            with tempfile.TemporaryDirectory() as tmp_dir_name:
218
                pt_file_name = os.path.join(tmp_dir_name, "traced_model.pt")
thomwolf's avatar
thomwolf committed
219

220
221
222
223
                try:
                    torch.jit.save(traced_gpt2, pt_file_name)
                except Exception:
                    self.fail("Couldn't save module.")
thomwolf's avatar
thomwolf committed
224

225
226
227
228
                try:
                    loaded_model = torch.jit.load(pt_file_name)
                except Exception:
                    self.fail("Couldn't load module.")
LysandreJik's avatar
LysandreJik committed
229

230
231
            model.to(torch_device)
            model.eval()
thomwolf's avatar
thomwolf committed
232

233
234
            loaded_model.to(torch_device)
            loaded_model.eval()
thomwolf's avatar
thomwolf committed
235

236
237
238
239
            model_state_dict = model.state_dict()
            loaded_model_state_dict = loaded_model.state_dict()

            self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys()))
thomwolf's avatar
thomwolf committed
240

241
            models_equal = True
242
243
            for layer_name, p1 in model_state_dict.items():
                p2 = loaded_model_state_dict[layer_name]
244
245
                if p1.data.ne(p2.data).sum() > 0:
                    models_equal = False
thomwolf's avatar
thomwolf committed
246

247
            self.assertTrue(models_equal)
thomwolf's avatar
thomwolf committed
248

Patrick von Platen's avatar
Patrick von Platen committed
249
250
    def test_headmasking(self):
        if not self.test_head_masking:
251
            return
252

253
254
255
        global_rng.seed(42)
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        global_rng.seed()
LysandreJik's avatar
LysandreJik committed
256

257
258
259
260
261
262
263
        config.output_attentions = True
        config.output_hidden_states = True
        configs_no_init = _config_zero_init(config)  # To be sure we have no Nan
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            model.to(torch_device)
            model.eval()
LysandreJik's avatar
LysandreJik committed
264

265
266
267
            # Prepare head_mask
            # Set require_grad after having prepared the tensor to avoid error (leaf variable has been moved into the graph interior)
            head_mask = torch.ones(
268
                self.model_tester.num_hidden_layers, self.model_tester.num_attention_heads, device=torch_device,
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
            )
            head_mask[0, 0] = 0
            head_mask[-1, :-1] = 0
            head_mask.requires_grad_(requires_grad=True)
            inputs = inputs_dict.copy()
            inputs["head_mask"] = head_mask

            outputs = model(**inputs)

            # Test that we can get a gradient back for importance score computation
            output = sum(t.sum() for t in outputs[0])
            output = output.sum()
            output.backward()
            multihead_outputs = head_mask.grad

            attentions = outputs[-1]

            # Remove Nan
            for t in attentions:
                self.assertLess(
                    torch.sum(torch.isnan(t)), t.numel() / 4
                )  # Check we don't have more than 25% nans (arbitrary)
            attentions = [
                t.masked_fill(torch.isnan(t), 0.0) for t in attentions
            ]  # remove them (the test is less complete)

            self.assertIsNotNone(multihead_outputs)
            self.assertEqual(len(multihead_outputs), self.model_tester.num_hidden_layers)
            self.assertAlmostEqual(attentions[0][..., 0, :, :].flatten().sum().item(), 0.0)
            self.assertNotEqual(attentions[0][..., -1, :, :].flatten().sum().item(), 0.0)
            self.assertNotEqual(attentions[1][..., 0, :, :].flatten().sum().item(), 0.0)
            self.assertAlmostEqual(attentions[-1][..., -2, :, :].flatten().sum().item(), 0.0)
            self.assertNotEqual(attentions[-1][..., -1, :, :].flatten().sum().item(), 0.0)

Patrick von Platen's avatar
Patrick von Platen committed
303
304
    def test_head_pruning(self):
        if not self.test_pruning:
305
306
307
            return

        for model_class in self.all_model_classes:
308
            (config, inputs_dict,) = self.model_tester.prepare_config_and_inputs_for_common()
309

310
311
            if "head_mask" in inputs_dict:
                del inputs_dict["head_mask"]
312

313
314
315
316
317
            config.output_attentions = True
            config.output_hidden_states = False
            model = model_class(config=config)
            model.to(torch_device)
            model.eval()
318
319
320
321
            heads_to_prune = {
                0: list(range(1, self.model_tester.num_attention_heads)),
                -1: [0],
            }
322
323
324
            model.prune_heads(heads_to_prune)
            with torch.no_grad():
                outputs = model(**inputs_dict)
325

326
            attentions = outputs[-1]
327

328
329
330
            self.assertEqual(attentions[0].shape[-3], 1)
            self.assertEqual(attentions[1].shape[-3], self.model_tester.num_attention_heads)
            self.assertEqual(attentions[-1].shape[-3], self.model_tester.num_attention_heads - 1)
LysandreJik's avatar
LysandreJik committed
331

Patrick von Platen's avatar
Patrick von Platen committed
332
333
    def test_head_pruning_save_load_from_pretrained(self):
        if not self.test_pruning:
334
            return
LysandreJik's avatar
LysandreJik committed
335

336
        for model_class in self.all_model_classes:
337
            (config, inputs_dict,) = self.model_tester.prepare_config_and_inputs_for_common()
338
339
340

            if "head_mask" in inputs_dict:
                del inputs_dict["head_mask"]
341

342
343
344
345
346
            config.output_attentions = True
            config.output_hidden_states = False
            model = model_class(config=config)
            model.to(torch_device)
            model.eval()
347
348
349
350
            heads_to_prune = {
                0: list(range(1, self.model_tester.num_attention_heads)),
                -1: [0],
            }
351
            model.prune_heads(heads_to_prune)
352

353
            with tempfile.TemporaryDirectory() as temp_dir_name:
354
355
                model.save_pretrained(temp_dir_name)
                model = model_class.from_pretrained(temp_dir_name)
356
                model.to(torch_device)
357

358
359
360
361
362
363
            with torch.no_grad():
                outputs = model(**inputs_dict)
            attentions = outputs[-1]
            self.assertEqual(attentions[0].shape[-3], 1)
            self.assertEqual(attentions[1].shape[-3], self.model_tester.num_attention_heads)
            self.assertEqual(attentions[-1].shape[-3], self.model_tester.num_attention_heads - 1)
364

Patrick von Platen's avatar
Patrick von Platen committed
365
366
    def test_head_pruning_save_load_from_config_init(self):
        if not self.test_pruning:
367
            return
368

369
        for model_class in self.all_model_classes:
370
            (config, inputs_dict,) = self.model_tester.prepare_config_and_inputs_for_common()
371

372
373
            if "head_mask" in inputs_dict:
                del inputs_dict["head_mask"]
374

375
376
            config.output_attentions = True
            config.output_hidden_states = False
377

378
379
380
381
            heads_to_prune = {
                0: list(range(1, self.model_tester.num_attention_heads)),
                -1: [0],
            }
382
            config.pruned_heads = heads_to_prune
383

384
385
386
            model = model_class(config=config)
            model.to(torch_device)
            model.eval()
387

388
389
390
            with torch.no_grad():
                outputs = model(**inputs_dict)
            attentions = outputs[-1]
391

392
393
394
            self.assertEqual(attentions[0].shape[-3], 1)
            self.assertEqual(attentions[1].shape[-3], self.model_tester.num_attention_heads)
            self.assertEqual(attentions[-1].shape[-3], self.model_tester.num_attention_heads - 1)
395

Patrick von Platen's avatar
Patrick von Platen committed
396
397
    def test_head_pruning_integration(self):
        if not self.test_pruning:
398
            return
399

400
        for model_class in self.all_model_classes:
401
            (config, inputs_dict,) = self.model_tester.prepare_config_and_inputs_for_common()
402

403
404
            if "head_mask" in inputs_dict:
                del inputs_dict["head_mask"]
405

406
407
            config.output_attentions = True
            config.output_hidden_states = False
408

409
410
            heads_to_prune = {0: [0], 1: [1, 2]}
            config.pruned_heads = heads_to_prune
411

412
413
414
            model = model_class(config=config)
            model.to(torch_device)
            model.eval()
415

416
417
418
            with torch.no_grad():
                outputs = model(**inputs_dict)
            attentions = outputs[-1]
419

420
421
422
423
            self.assertEqual(attentions[0].shape[-3], self.model_tester.num_attention_heads - 1)
            self.assertEqual(attentions[1].shape[-3], self.model_tester.num_attention_heads - 2)
            self.assertEqual(attentions[2].shape[-3], self.model_tester.num_attention_heads)
            self.assertEqual(attentions[3].shape[-3], self.model_tester.num_attention_heads)
thomwolf's avatar
thomwolf committed
424

425
            with tempfile.TemporaryDirectory() as temp_dir_name:
426
427
                model.save_pretrained(temp_dir_name)
                model = model_class.from_pretrained(temp_dir_name)
428
                model.to(torch_device)
thomwolf's avatar
thomwolf committed
429

430
431
432
            with torch.no_grad():
                outputs = model(**inputs_dict)
            attentions = outputs[-1]
LysandreJik's avatar
LysandreJik committed
433

434
435
436
437
            self.assertEqual(attentions[0].shape[-3], self.model_tester.num_attention_heads - 1)
            self.assertEqual(attentions[1].shape[-3], self.model_tester.num_attention_heads - 2)
            self.assertEqual(attentions[2].shape[-3], self.model_tester.num_attention_heads)
            self.assertEqual(attentions[3].shape[-3], self.model_tester.num_attention_heads)
thomwolf's avatar
thomwolf committed
438

439
440
            heads_to_prune = {0: [0], 2: [1, 2]}
            model.prune_heads(heads_to_prune)
441

442
443
444
            with torch.no_grad():
                outputs = model(**inputs_dict)
            attentions = outputs[-1]
445

446
447
448
449
            self.assertEqual(attentions[0].shape[-3], self.model_tester.num_attention_heads - 1)
            self.assertEqual(attentions[1].shape[-3], self.model_tester.num_attention_heads - 2)
            self.assertEqual(attentions[2].shape[-3], self.model_tester.num_attention_heads - 2)
            self.assertEqual(attentions[3].shape[-3], self.model_tester.num_attention_heads)
450

451
            self.assertDictEqual(model.config.pruned_heads, {0: [0], 1: [1, 2], 2: [1, 2]})
thomwolf's avatar
thomwolf committed
452

Patrick von Platen's avatar
Patrick von Platen committed
453
    def test_hidden_states_output(self):
454
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
thomwolf's avatar
thomwolf committed
455

456
457
458
459
        for model_class in self.all_model_classes:
            config.output_hidden_states = True
            config.output_attentions = False
            model = model_class(config)
460
            model.to(torch_device)
thomwolf's avatar
thomwolf committed
461
            model.eval()
thomwolf's avatar
thomwolf committed
462
            with torch.no_grad():
463
464
465
466
467
468
469
470
471
472
473
474
475
                outputs = model(**inputs_dict)
            hidden_states = outputs[-1]
            self.assertEqual(model.config.output_attentions, False)
            self.assertEqual(model.config.output_hidden_states, True)
            self.assertEqual(len(hidden_states), self.model_tester.num_hidden_layers + 1)
            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [
                    self.model_tester.encoder_seq_length
                    if hasattr(self.model_tester, "encoder_seq_length")
                    else self.model_tester.seq_length,
                    self.model_tester.hidden_size,
                ],
476
            )
thomwolf's avatar
thomwolf committed
477

Patrick von Platen's avatar
Patrick von Platen committed
478
    def test_resize_tokens_embeddings(self):
479
        (original_config, inputs_dict,) = self.model_tester.prepare_config_and_inputs_for_common()
Patrick von Platen's avatar
Patrick von Platen committed
480
        if not self.test_resize_embeddings:
481
482
483
484
485
            return

        for model_class in self.all_model_classes:
            config = copy.deepcopy(original_config)
            model = model_class(config)
486
            model.to(torch_device)
487
488
489
490
491
492
493
494
495
496
497

            model_vocab_size = config.vocab_size
            # Retrieve the embeddings and clone theme
            model_embed = model.resize_token_embeddings(model_vocab_size)
            cloned_embeddings = model_embed.weight.clone()

            # Check that resizing the token embeddings with a larger vocab size increases the model's vocab size
            model_embed = model.resize_token_embeddings(model_vocab_size + 10)
            self.assertEqual(model.config.vocab_size, model_vocab_size + 10)
            # Check that it actually resizes the embeddings matrix
            self.assertEqual(model_embed.weight.shape[0], cloned_embeddings.shape[0] + 10)
498
499
            # Check that the model can still do a forward pass successfully (every parameter should be resized)
            model(**inputs_dict)
500
501
502
503
504
505
506

            # Check that resizing the token embeddings with a smaller vocab size decreases the model's vocab size
            model_embed = model.resize_token_embeddings(model_vocab_size - 15)
            self.assertEqual(model.config.vocab_size, model_vocab_size - 15)
            # Check that it actually resizes the embeddings matrix
            self.assertEqual(model_embed.weight.shape[0], cloned_embeddings.shape[0] - 15)

507
508
509
510
511
            # Check that the model can still do a forward pass successfully (every parameter should be resized)
            # Input ids should be clamped to the maximum size of the vocabulary
            inputs_dict["input_ids"].clamp_(max=model_vocab_size - 15 - 1)
            model(**inputs_dict)

512
513
514
515
516
517
518
519
            # Check that adding and removing tokens has not modified the first part of the embedding matrix.
            models_equal = True
            for p1, p2 in zip(cloned_embeddings, model_embed.weight):
                if p1.data.ne(p2.data).sum() > 0:
                    models_equal = False

            self.assertTrue(models_equal)

Patrick von Platen's avatar
Patrick von Platen committed
520
    def test_model_common_attributes(self):
521
522
523
524
525
526
527
528
529
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertIsInstance(model.get_input_embeddings(), (torch.nn.Embedding, AdaptiveEmbedding))
            model.set_input_embeddings(torch.nn.Embedding(10, 10))
            x = model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(x, torch.nn.Linear))

530
    def test_correct_missing_keys(self):
531
532
        if not self.test_missing_keys:
            return
533
534
535
536
537
538
539
540
541
542
543
544
545
546
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            base_model_prefix = model.base_model_prefix

            if hasattr(model, base_model_prefix):
                with tempfile.TemporaryDirectory() as temp_dir_name:
                    model.base_model.save_pretrained(temp_dir_name)
                    model, loading_info = model_class.from_pretrained(temp_dir_name, output_loading_info=True)

                    with self.subTest(msg="Missing keys for {}".format(model.__class__.__name__)):
                        self.assertGreater(len(loading_info["missing_keys"]), 0)

547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
    def test_tie_model_weights(self):
        if not self.test_torchscript:
            return

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        def check_same_values(layer_1, layer_2):
            equal = True
            for p1, p2 in zip(layer_1.weight, layer_2.weight):
                if p1.data.ne(p2.data).sum() > 0:
                    equal = False
            return equal

        for model_class in self.all_model_classes:
            config.torchscript = True
            model_not_tied = model_class(config)
            if model_not_tied.get_output_embeddings() is None:
                continue

            params_not_tied = list(model_not_tied.parameters())

            config_tied = copy.deepcopy(config)
            config_tied.torchscript = False
            model_tied = model_class(config_tied)
            params_tied = list(model_tied.parameters())

            # Check that the embedding layer and decoding layer are the same in size and in value
            self.assertGreater(len(params_not_tied), len(params_tied))
            # self.assertTrue(check_same_values(embeddings, decoding))

            # # Check that after modification, they remain the same.
            # embeddings.weight.data.div_(2)
            # # Check that the embedding layer and decoding layer are the same in size and in value
            # self.assertTrue(embeddings.weight.shape, decoding.weight.shape)
            # self.assertTrue(check_same_values(embeddings, decoding))

            # # Check that after modification, they remain the same.
            # decoding.weight.data.div_(4)
            # # Check that the embedding layer and decoding layer are the same in size and in value
            # self.assertTrue(embeddings.weight.shape, decoding.weight.shape)
            # self.assertTrue(check_same_values(embeddings, decoding))

            # Check that after resize they remain tied.
            model_tied.resize_token_embeddings(config.vocab_size + 10)
            params_tied_2 = list(model_tied.parameters())
            self.assertGreater(len(params_not_tied), len(params_tied))
            self.assertEqual(len(params_tied_2), len(params_tied))

            # decoding.weight.data.mul_(20)
            # # Check that the embedding layer and decoding layer are the same in size and in value
            # self.assertTrue(model.transformer.wte.weight.shape, model.lm_head.weight.shape)
            # self.assertTrue(check_same_values(model.transformer.wte, model.lm_head))

Patrick von Platen's avatar
Patrick von Platen committed
600
    def test_inputs_embeds(self):
Sam Shleifer's avatar
Sam Shleifer committed
601

602
603
604
605
606
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        if not self.is_encoder_decoder:
            input_ids = inputs_dict["input_ids"]
            del inputs_dict["input_ids"]
        else:
607
            encoder_input_ids = inputs_dict["input_ids"]
Sam Shleifer's avatar
Sam Shleifer committed
608
            decoder_input_ids = inputs_dict.get("decoder_input_ids", encoder_input_ids)
609
            del inputs_dict["input_ids"]
Sam Shleifer's avatar
Sam Shleifer committed
610
            inputs_dict.pop("decoder_input_ids", None)
611
612
613

        for model_class in self.all_model_classes:
            model = model_class(config)
614
            model.to(torch_device)
thomwolf's avatar
thomwolf committed
615
            model.eval()
616
617
618
619
620

            wte = model.get_input_embeddings()
            if not self.is_encoder_decoder:
                inputs_dict["inputs_embeds"] = wte(input_ids)
            else:
621
                inputs_dict["inputs_embeds"] = wte(encoder_input_ids)
622
623
                inputs_dict["decoder_inputs_embeds"] = wte(decoder_input_ids)

thomwolf's avatar
thomwolf committed
624
            with torch.no_grad():
625
                model(**inputs_dict)
626

627
    def test_lm_head_model_random_generate(self):
628
629

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
630
        input_ids = inputs_dict.get("input_ids")
631

632
633
634
        if self.is_encoder_decoder:
            config.output_past = True  # needed for Bart TODO: might have to update for other encoder-decoder models

635
636
637
638
639
640
641
        for model_class in self.all_generative_model_classes:
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            if config.bos_token_id is None:
                with self.assertRaises(AssertionError):
642
                    model.generate(do_sample=True, max_length=5)
643
                # batch_size = 1
644
                self._check_generated_tokens(model.generate(input_ids, do_sample=True))
645
                # batch_size = 1, num_beams > 1
646
                self._check_generated_tokens(model.generate(input_ids, do_sample=True, num_beams=3))
647
648
            else:
                # batch_size = 1
649
                self._check_generated_tokens(model.generate(do_sample=True, max_length=5))
650
                # batch_size = 1, num_beams > 1
651
                self._check_generated_tokens(model.generate(do_sample=True, max_length=5, num_beams=3))
652

653
654
655
656
657
658
659
660
661
            with self.assertRaises(AssertionError):
                # generating multiple sequences when greedy no beam generation
                # is not allowed as it would always generate the same sequences
                model.generate(input_ids, do_sample=False, num_return_sequences=2)

            with self.assertRaises(AssertionError):
                # generating more sequences than having beams leads is not possible
                model.generate(input_ids, do_sample=False, num_return_sequences=3, num_beams=2)

662
            # batch_size > 1, sample
663
            self._check_generated_tokens(model.generate(input_ids, do_sample=True, num_return_sequences=3))
664
            # batch_size > 1, greedy
665
            self._check_generated_tokens(model.generate(input_ids, do_sample=False))
666

667
            # batch_size > 1, num_beams > 1, sample
668
669
670
            self._check_generated_tokens(
                model.generate(input_ids, do_sample=True, num_beams=3, num_return_sequences=3,)
            )
671
672
673
674
675
676
677
678
679
680
            # batch_size > 1, num_beams > 1, greedy
            self._check_generated_tokens(
                model.generate(input_ids, do_sample=False, num_beams=3, num_return_sequences=3)
            )

    def _check_generated_tokens(self, output_ids):
        for token_id in output_ids[0].tolist():
            self.assertGreaterEqual(token_id, 0)
            self.assertLess(token_id, self.model_tester.vocab_size)

681

682
global_rng = random.Random()
thomwolf's avatar
thomwolf committed
683
684


thomwolf's avatar
thomwolf committed
685
def ids_tensor(shape, vocab_size, rng=None, name=None):
686
    #  Creates a random int32 tensor of the shape within the vocab size
thomwolf's avatar
thomwolf committed
687
    if rng is None:
688
        rng = global_rng
thomwolf's avatar
thomwolf committed
689

thomwolf's avatar
thomwolf committed
690
691
692
    total_dims = 1
    for dim in shape:
        total_dims *= dim
thomwolf's avatar
thomwolf committed
693

thomwolf's avatar
thomwolf committed
694
695
696
    values = []
    for _ in range(total_dims):
        values.append(rng.randint(0, vocab_size - 1))
thomwolf's avatar
thomwolf committed
697

698
    return torch.tensor(data=values, dtype=torch.long, device=torch_device).view(shape).contiguous()
thomwolf's avatar
thomwolf committed
699
700


701
702
703
704
705
706
707
708
709
710
711
712
713
def floats_tensor(shape, scale=1.0, rng=None, name=None):
    """Creates a random float32 tensor of the shape within the vocab size."""
    if rng is None:
        rng = global_rng

    total_dims = 1
    for dim in shape:
        total_dims *= dim

    values = []
    for _ in range(total_dims):
        values.append(rng.random() * scale)

714
    return torch.tensor(data=values, dtype=torch.float, device=torch_device).view(shape).contiguous()
715
716


717
@require_torch
thomwolf's avatar
thomwolf committed
718
class ModelUtilsTest(unittest.TestCase):
719
    @slow
Patrick von Platen's avatar
Patrick von Platen committed
720
    def test_model_from_pretrained(self):
thomwolf's avatar
thomwolf committed
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
        logging.basicConfig(level=logging.INFO)
        for model_name in list(BERT_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]:
            config = BertConfig.from_pretrained(model_name)
            self.assertIsNotNone(config)
            self.assertIsInstance(config, PretrainedConfig)

            model = BertModel.from_pretrained(model_name)
            model, loading_info = BertModel.from_pretrained(model_name, output_loading_info=True)
            self.assertIsNotNone(model)
            self.assertIsInstance(model, PreTrainedModel)
            for value in loading_info.values():
                self.assertEqual(len(value), 0)

            config = BertConfig.from_pretrained(model_name, output_attentions=True, output_hidden_states=True)
            model = BertModel.from_pretrained(model_name, output_attentions=True, output_hidden_states=True)
            self.assertEqual(model.config.output_attentions, True)
            self.assertEqual(model.config.output_hidden_states, True)
            self.assertEqual(model.config, config)
739
740
741
742
743
744


@require_torch
class UtilsFunctionsTest(unittest.TestCase):

    # tests whether the top_k_top_p function behaves as expected
Patrick von Platen's avatar
Patrick von Platen committed
745
    def test_top_k_top_p_filtering(self):
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
        logits = torch.tensor(
            [
                [
                    8.2220991,  # 3rd highest value; idx. 0
                    -0.5620044,
                    5.23229752,
                    4.0386393,
                    -6.8798378,
                    -0.54785802,
                    -3.2012153,
                    2.92777176,
                    1.88171953,
                    7.35341276,  # 5th highest value; idx. 9
                    8.43207833,  # 2nd highest value; idx. 10
                    -9.85711836,
                    -5.96209236,
                    -1.13039161,
                    -7.1115294,
                    -0.8369633,
                    -5.3186408,
                    7.06427407,
                    0.81369344,
                    -0.82023817,
                    -5.9179796,
                    0.58813443,
                    -6.99778438,
                    4.71551189,
                    -0.18771637,
                    7.44020759,  # 4th highest value; idx. 25
                    9.38450987,  # 1st highest value; idx. 26
                    2.12662941,
                    -9.32562038,
                    2.35652522,
                ],  # cummulative prob of 5 highest values <= 0.6
                [
                    0.58425518,
                    4.53139238,
                    -5.57510464,
                    -6.28030699,
                    -7.19529503,
                    -4.02122551,
                    1.39337037,
                    -6.06707057,
                    1.59480517,
                    -9.643119,
                    0.03907799,
                    0.67231762,
                    -8.88206726,
                    6.27115922,  # 4th highest value; idx. 13
                    2.28520723,
                    4.82767506,
                    4.30421368,
                    8.8275313,  # 2nd highest value; idx. 17
                    5.44029958,  # 5th highest value; idx. 18
                    -4.4735794,
                    7.38579536,  # 3rd highest value; idx. 20
                    -2.91051663,
                    2.61946077,
                    -2.5674762,
                    -9.48959302,
                    -4.02922645,
                    -1.35416918,
                    9.67702323,  # 1st highest value; idx. 27
                    -5.89478553,
                    1.85370467,
                ],  # cummulative prob of 5 highest values <= 0.6
            ],
            dtype=torch.float,
            device=torch_device,
        )

        non_inf_expected_idx = torch.tensor(
            [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]],
            dtype=torch.long,
            device=torch_device,
        )  # expected non filtered idx as noted above

        non_inf_expected_output = torch.tensor(
            [
                8.2221,
                7.3534,
                8.4321,
                7.4402,
                9.3845,
                6.2712,
                8.8275,
                5.4403,
                7.3858,
                9.6770,
            ],  # expected non filtered values as noted above
            dtype=torch.float,
            device=torch_device,
        )

        output = top_k_top_p_filtering(logits, top_k=10, top_p=0.6, min_tokens_to_keep=4)
        non_inf_output = output[output != -float("inf")].to(device=torch_device)
        non_inf_idx = (output != -float("inf")).nonzero().to(device=torch_device)

        self.assertTrue(torch.allclose(non_inf_expected_output, non_inf_output, atol=1e-12))
        self.assertTrue(torch.all(torch.eq(non_inf_expected_idx, non_inf_idx)))