test_modeling_tf_xlnet.py 22.5 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
Sylvain Gugger's avatar
Sylvain Gugger committed
2
# Copyright 2020 The HuggingFace Team. All rights reserved.
thomwolf's avatar
thomwolf 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

thomwolf's avatar
thomwolf committed
16

Matt's avatar
Matt committed
17
18
from __future__ import annotations

19
import inspect
thomwolf's avatar
thomwolf committed
20
import random
21
import unittest
thomwolf's avatar
thomwolf committed
22

23
from transformers import XLNetConfig, is_tf_available
24
from transformers.testing_utils import require_tf, slow
thomwolf's avatar
thomwolf committed
25

Yih-Dar's avatar
Yih-Dar committed
26
27
from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
28
from ...test_pipeline_mixin import PipelineTesterMixin
Aymeric Augustin's avatar
Aymeric Augustin committed
29
30


thomwolf's avatar
thomwolf committed
31
32
33
if is_tf_available():
    import tensorflow as tf

Sylvain Gugger's avatar
Sylvain Gugger committed
34
    from transformers.models.xlnet.modeling_tf_xlnet import (
35
36
37
        TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLNetForMultipleChoice,
        TFXLNetForQuestionAnsweringSimple,
38
39
        TFXLNetForSequenceClassification,
        TFXLNetForTokenClassification,
40
41
        TFXLNetLMHeadModel,
        TFXLNetModel,
42
43
    )

44

45
46
class TFXLNetModelTester:
    def __init__(
Lysandre's avatar
Lysandre committed
47
48
        self,
        parent,
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
    ):
        self.parent = parent
        self.batch_size = 13
        self.seq_length = 7
        self.mem_len = 10
        # self.key_len = seq_length + mem_len
        self.clamp_len = -1
        self.reuse_len = 15
        self.is_training = True
        self.use_labels = True
        self.vocab_size = 99
        self.cutoffs = [10, 50, 80]
        self.hidden_size = 32
        self.num_attention_heads = 4
        self.d_inner = 128
64
        self.num_hidden_layers = 2
65
66
67
68
69
70
71
72
73
74
        self.type_sequence_label_size = 2
        self.untie_r = True
        self.bi_data = False
        self.same_length = False
        self.initializer_range = 0.05
        self.seed = 1
        self.type_vocab_size = 2
        self.bos_token_id = 1
        self.eos_token_id = 2
        self.pad_token_id = 5
75
        self.num_choices = 4
76
77
78
79
80

    def prepare_config_and_inputs(self):
        input_ids_1 = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
        input_ids_2 = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
        segment_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
81
        input_mask = random_attention_mask([self.batch_size, self.seq_length], dtype=tf.float32)
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

        input_ids_q = ids_tensor([self.batch_size, self.seq_length + 1], self.vocab_size)
        perm_mask = tf.zeros((self.batch_size, self.seq_length + 1, self.seq_length), dtype=tf.float32)
        perm_mask_last = tf.ones((self.batch_size, self.seq_length + 1, 1), dtype=tf.float32)
        perm_mask = tf.concat([perm_mask, perm_mask_last], axis=-1)
        # perm_mask[:, :, -1] = 1.0  # Previous tokens don't see last token
        target_mapping = tf.zeros((self.batch_size, 1, self.seq_length), dtype=tf.float32)
        target_mapping_last = tf.ones((self.batch_size, 1, 1), dtype=tf.float32)
        target_mapping = tf.concat([target_mapping, target_mapping_last], axis=-1)
        # target_mapping[:, 0, -1] = 1.0  # predict last token

        sequence_labels = None
        lm_labels = None
        is_impossible_labels = None
        if self.use_labels:
            lm_labels = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
            sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
            is_impossible_labels = ids_tensor([self.batch_size], 2, dtype=tf.float32)

        config = XLNetConfig(
            vocab_size=self.vocab_size,
            d_model=self.hidden_size,
            n_head=self.num_attention_heads,
            d_inner=self.d_inner,
            n_layer=self.num_hidden_layers,
            untie_r=self.untie_r,
            mem_len=self.mem_len,
            clamp_len=self.clamp_len,
            same_length=self.same_length,
            reuse_len=self.reuse_len,
            bi_data=self.bi_data,
            initializer_range=self.initializer_range,
            num_labels=self.type_sequence_label_size,
            bos_token_id=self.bos_token_id,
            pad_token_id=self.pad_token_id,
            eos_token_id=self.eos_token_id,
118
        )
thomwolf's avatar
thomwolf committed
119

120
        return (
121
122
123
124
125
126
127
128
129
130
131
            config,
            input_ids_1,
            input_ids_2,
            input_ids_q,
            perm_mask,
            input_mask,
            target_mapping,
            segment_ids,
            lm_labels,
            sequence_labels,
            is_impossible_labels,
132
        )
thomwolf's avatar
thomwolf committed
133

134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
    def set_seed(self):
        random.seed(self.seed)
        tf.random.set_seed(self.seed)

    def create_and_check_xlnet_base_model(
        self,
        config,
        input_ids_1,
        input_ids_2,
        input_ids_q,
        perm_mask,
        input_mask,
        target_mapping,
        segment_ids,
        lm_labels,
        sequence_labels,
        is_impossible_labels,
    ):
        model = TFXLNetModel(config)

        inputs = {"input_ids": input_ids_1, "input_mask": input_mask, "token_type_ids": segment_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
155
        result = model(inputs)
156
157

        inputs = [input_ids_1, input_mask]
Sylvain Gugger's avatar
Sylvain Gugger committed
158
        result = model(inputs)
159

160
        config.use_mems_eval = False
161
162
163
164
        model = TFXLNetModel(config)
        no_mems_outputs = model(inputs)
        self.parent.assertEqual(len(no_mems_outputs), 1)

165
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
166
        self.parent.assertListEqual(
167
168
            [mem.shape for mem in result.mems],
            [(self.seq_length, self.batch_size, self.hidden_size)] * self.num_hidden_layers,
169
        )
thomwolf's avatar
thomwolf committed
170

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
    def create_and_check_xlnet_lm_head(
        self,
        config,
        input_ids_1,
        input_ids_2,
        input_ids_q,
        perm_mask,
        input_mask,
        target_mapping,
        segment_ids,
        lm_labels,
        sequence_labels,
        is_impossible_labels,
    ):
        model = TFXLNetLMHeadModel(config)
thomwolf's avatar
thomwolf committed
186

187
        inputs_1 = {"input_ids": input_ids_1, "token_type_ids": segment_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
188
        all_logits_1, mems_1 = model(inputs_1).to_tuple()
189

190
        inputs_2 = {"input_ids": input_ids_2, "mems": mems_1, "token_type_ids": segment_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
191
        all_logits_2, mems_2 = model(inputs_2).to_tuple()
192

193
        inputs_3 = {"input_ids": input_ids_q, "perm_mask": perm_mask, "target_mapping": target_mapping}
Sylvain Gugger's avatar
Sylvain Gugger committed
194
        logits, _ = model(inputs_3).to_tuple()
195

196
        self.parent.assertEqual(all_logits_1.shape, (self.batch_size, self.seq_length, self.vocab_size))
197
        self.parent.assertListEqual(
198
199
            [mem.shape for mem in mems_1],
            [(self.seq_length, self.batch_size, self.hidden_size)] * self.num_hidden_layers,
200
        )
201
        self.parent.assertEqual(all_logits_2.shape, (self.batch_size, self.seq_length, self.vocab_size))
202
        self.parent.assertListEqual(
203
204
            [mem.shape for mem in mems_2],
            [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers,
205
        )
206

207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
    def create_and_check_xlnet_qa(
        self,
        config,
        input_ids_1,
        input_ids_2,
        input_ids_q,
        perm_mask,
        input_mask,
        target_mapping,
        segment_ids,
        lm_labels,
        sequence_labels,
        is_impossible_labels,
    ):
        model = TFXLNetForQuestionAnsweringSimple(config)

        inputs = {"input_ids": input_ids_1, "attention_mask": input_mask, "token_type_ids": segment_ids}
Sylvain Gugger's avatar
Sylvain Gugger committed
224
        result = model(inputs)
225

226
227
        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))
228
        self.parent.assertListEqual(
229
230
            [mem.shape for mem in result.mems],
            [(self.seq_length, self.batch_size, self.hidden_size)] * self.num_hidden_layers,
231
        )
232

233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
    def create_and_check_xlnet_sequence_classif(
        self,
        config,
        input_ids_1,
        input_ids_2,
        input_ids_q,
        perm_mask,
        input_mask,
        target_mapping,
        segment_ids,
        lm_labels,
        sequence_labels,
        is_impossible_labels,
    ):
        model = TFXLNetForSequenceClassification(config)

Sylvain Gugger's avatar
Sylvain Gugger committed
249
        result = model(input_ids_1)
250

251
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.type_sequence_label_size))
252
        self.parent.assertListEqual(
253
254
            [mem.shape for mem in result.mems],
            [(self.seq_length, self.batch_size, self.hidden_size)] * self.num_hidden_layers,
255
        )
256

257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
    def create_and_check_xlnet_for_token_classification(
        self,
        config,
        input_ids_1,
        input_ids_2,
        input_ids_q,
        perm_mask,
        input_mask,
        target_mapping,
        segment_ids,
        lm_labels,
        sequence_labels,
        is_impossible_labels,
    ):
        config.num_labels = input_ids_1.shape[1]
        model = TFXLNetForTokenClassification(config)
        inputs = {
            "input_ids": input_ids_1,
            "attention_mask": input_mask,
            # 'token_type_ids': token_type_ids
        }
Sylvain Gugger's avatar
Sylvain Gugger committed
278
        result = model(inputs)
279
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, config.num_labels))
280
        self.parent.assertListEqual(
281
282
            [mem.shape for mem in result.mems],
            [(self.seq_length, self.batch_size, self.hidden_size)] * self.num_hidden_layers,
283
        )
284

285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
    def create_and_check_xlnet_for_multiple_choice(
        self,
        config,
        input_ids_1,
        input_ids_2,
        input_ids_q,
        perm_mask,
        input_mask,
        target_mapping,
        segment_ids,
        lm_labels,
        sequence_labels,
        is_impossible_labels,
    ):
        config.num_choices = self.num_choices
        model = TFXLNetForMultipleChoice(config=config)
        multiple_choice_inputs_ids = tf.tile(tf.expand_dims(input_ids_1, 1), (1, self.num_choices, 1))
        multiple_choice_input_mask = tf.tile(tf.expand_dims(input_mask, 1), (1, self.num_choices, 1))
        multiple_choice_token_type_ids = tf.tile(tf.expand_dims(segment_ids, 1), (1, self.num_choices, 1))
        inputs = {
            "input_ids": multiple_choice_inputs_ids,
            "attention_mask": multiple_choice_input_mask,
            "token_type_ids": multiple_choice_token_type_ids,
        }
Sylvain Gugger's avatar
Sylvain Gugger committed
309
        result = model(inputs)
Julien Plu's avatar
Julien Plu committed
310

311
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_choices))
Julien Plu's avatar
Julien Plu committed
312
        self.parent.assertListEqual(
313
314
            [mem.shape for mem in result.mems],
            [(self.seq_length, self.batch_size * self.num_choices, self.hidden_size)] * self.num_hidden_layers,
Julien Plu's avatar
Julien Plu committed
315
        )
316

317
318
319
    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
320
321
322
323
324
325
326
327
328
329
330
            config,
            input_ids_1,
            input_ids_2,
            input_ids_q,
            perm_mask,
            input_mask,
            target_mapping,
            segment_ids,
            lm_labels,
            sequence_labels,
            is_impossible_labels,
331
332
333
        ) = config_and_inputs
        inputs_dict = {"input_ids": input_ids_1}
        return config, inputs_dict
334
335


336
@require_tf
337
class TFXLNetModelTest(TFModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
338
339
340
341
342
343
344
    all_model_classes = (
        (
            TFXLNetModel,
            TFXLNetLMHeadModel,
            TFXLNetForSequenceClassification,
            TFXLNetForTokenClassification,
            TFXLNetForQuestionAnsweringSimple,
345
            TFXLNetForMultipleChoice,
346
347
348
349
350
351
352
        )
        if is_tf_available()
        else ()
    )
    all_generative_model_classes = (
        (TFXLNetLMHeadModel,) if is_tf_available() else ()
    )  # TODO (PVP): Check other models whether language generation is also applicable
353
354
355
356
357
358
359
360
361
362
363
364
    pipeline_model_mapping = (
        {
            "feature-extraction": TFXLNetModel,
            "question-answering": TFXLNetForQuestionAnsweringSimple,
            "text-classification": TFXLNetForSequenceClassification,
            "text-generation": TFXLNetLMHeadModel,
            "token-classification": TFXLNetForTokenClassification,
            "zero-shot": TFXLNetForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
365
    test_head_masking = False
366
    test_onnx = False
thomwolf's avatar
thomwolf committed
367

368
369
370
371
372
373
    # Note that `TFXLNetModelTest` is not a subclass of `GenerationTesterMixin`, so no contrastive generation tests
    # from there is run against `TFXLNetModel`.
    @unittest.skip("XLNet has special cache mechanism and is currently not working with contrastive generation")
    def test_xla_generate_contrastive(self):
        super().test_xla_generate_contrastive()

374
375
376
377
378
379
380
    # TODO: Fix the failed tests
    def is_pipeline_test_to_skip(
        self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name
    ):
        # Exception encountered when calling layer '...'
        return True

thomwolf's avatar
thomwolf committed
381
    def setUp(self):
382
        self.model_tester = TFXLNetModelTester(self)
thomwolf's avatar
thomwolf committed
383
384
385
386
387
388
389
390
391
392
393
394
395
        self.config_tester = ConfigTester(self, config_class=XLNetConfig, d_inner=37)

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

    def test_xlnet_base_model(self):
        self.model_tester.set_seed()
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlnet_base_model(*config_and_inputs)

    def test_xlnet_lm_head(self):
        self.model_tester.set_seed()
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
396
        self.model_tester.create_and_check_xlnet_lm_head(*config_and_inputs)
thomwolf's avatar
thomwolf committed
397
398
399
400
401
402

    def test_xlnet_sequence_classif(self):
        self.model_tester.set_seed()
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlnet_sequence_classif(*config_and_inputs)

403
404
405
406
    def test_xlnet_token_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlnet_for_token_classification(*config_and_inputs)

thomwolf's avatar
thomwolf committed
407
408
409
410
411
    def test_xlnet_qa(self):
        self.model_tester.set_seed()
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlnet_qa(*config_and_inputs)

412
413
414
415
    def test_xlnet_for_multiple_choice(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlnet_for_multiple_choice(*config_and_inputs)

416
    @slow
thomwolf's avatar
thomwolf committed
417
    def test_model_from_pretrained(self):
418
        for model_name in TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
419
            model = TFXLNetModel.from_pretrained(model_name)
thomwolf's avatar
thomwolf committed
420
            self.assertIsNotNone(model)
patrickvonplaten's avatar
patrickvonplaten committed
421

Matt's avatar
Matt committed
422
423
424
425
    @unittest.skip("Some of the XLNet models misbehave with flexible input shapes.")
    def test_compile_tf_model(self):
        pass

426
427
428
429
430
431
432
433
434
    # overwrite since `TFXLNetLMHeadModel` doesn't cut logits/labels
    def test_loss_computation(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            model = model_class(config)
            if getattr(model, "hf_compute_loss", None):
                # The number of elements in the loss should be the same as the number of elements in the label
                prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
                added_label = prepared_for_class[
435
                    sorted(prepared_for_class.keys() - inputs_dict.keys(), reverse=True)[0]
436
                ]
Matt's avatar
Matt committed
437
                expected_loss_size = added_label.shape.as_list()[:1]
438
439
440
441
442
443
444
445
446
447
448
449
450

                # `TFXLNetLMHeadModel` doesn't cut logits/labels
                # if model.__class__ in get_values(TF_MODEL_FOR_CAUSAL_LM_MAPPING):
                #     # if loss is causal lm loss, labels are shift, so that one label per batch
                #     # is cut
                #     loss_size = loss_size - self.model_tester.batch_size

                # Test that model correctly compute the loss with kwargs
                prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
                input_name = "input_ids" if "input_ids" in prepared_for_class else "pixel_values"
                input_ids = prepared_for_class.pop(input_name)

                loss = model(input_ids, **prepared_for_class)[0]
451
                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])
452
453
454
455

                # Test that model correctly compute the loss with a dict
                prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)
                loss = model(prepared_for_class)[0]
456
                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486

                # Test that model correctly compute the loss with a tuple
                prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True)

                # Get keys that were added with the _prepare_for_class function
                label_keys = prepared_for_class.keys() - inputs_dict.keys()
                signature = inspect.signature(model.call).parameters
                signature_names = list(signature.keys())

                # Create a dictionary holding the location of the tensors in the tuple
                tuple_index_mapping = {0: input_name}
                for label_key in label_keys:
                    label_key_index = signature_names.index(label_key)
                    tuple_index_mapping[label_key_index] = label_key
                sorted_tuple_index_mapping = sorted(tuple_index_mapping.items())
                # Initialize a list with their default values, update the values and convert to a tuple
                list_input = []

                for name in signature_names:
                    if name != "kwargs":
                        list_input.append(signature[name].default)

                for index, value in sorted_tuple_index_mapping:
                    list_input[index] = prepared_for_class[value]

                tuple_input = tuple(list_input)

                # Send to model
                loss = model(tuple_input[:-1])[0]

487
                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])
488

patrickvonplaten's avatar
patrickvonplaten committed
489

490
@require_tf
patrickvonplaten's avatar
patrickvonplaten committed
491
492
493
494
class TFXLNetModelLanguageGenerationTest(unittest.TestCase):
    @slow
    def test_lm_generate_xlnet_base_cased(self):
        model = TFXLNetLMHeadModel.from_pretrained("xlnet-base-cased")
495
        # fmt: off
patrickvonplaten's avatar
patrickvonplaten committed
496
497
498
        input_ids = tf.convert_to_tensor(
            [
                [
499
                    67, 2840, 19, 18, 1484, 20, 965, 29077, 8719, 1273, 21, 45, 273, 17, 10, 15048, 28, 27511, 21, 4185, 11, 41, 2444, 9, 32, 1025, 20, 8719, 26, 23, 673, 966, 19, 29077, 20643, 27511, 20822, 20643, 19, 17, 6616, 17511, 18, 8978, 20, 18, 777, 9, 19233, 1527, 17669, 19, 24, 673, 17, 28756, 150, 12943, 4354, 153, 27, 442, 37, 45, 668, 21, 24, 256, 20, 416, 22, 2771, 4901, 9, 12943, 4354, 153, 51, 24, 3004, 21, 28142, 23, 65, 20, 18, 416, 34, 24, 2958, 22947, 9, 1177, 45, 668, 3097, 13768, 23, 103, 28, 441, 148, 48, 20522, 19, 12943, 4354, 153, 12860, 34, 18, 326, 27, 17492, 684, 21, 6709, 9, 8585, 123, 266, 19, 12943, 4354, 153, 6872, 24, 3004, 20, 18, 9225, 2198, 19, 12717, 103, 22, 401, 24, 6348, 9, 12943, 4354, 153, 1068, 2768, 2286, 19, 33, 104, 19, 176, 24, 9313, 19, 20086, 28, 45, 10292, 9, 4, 3,
patrickvonplaten's avatar
patrickvonplaten committed
500
501
502
503
                ]
            ],
            dtype=tf.int32,
        )
504
505
        # fmt: on

patrickvonplaten's avatar
patrickvonplaten committed
506
507
508
509
510
511
512
513
514
515
516
        #  In 1991, the remains of Russian Tsar Nicholas II and his family
        #  (except for Alexei and Maria) are discovered.
        #  The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the
        #  remainder of the story. 1883 Western Siberia,
        #  a young Grigori Rasputin is asked by his father and a group of men to perform magic.
        #  Rasputin has a vision and denounces one of the men as a horse thief. Although his
        #  father initially slaps him for making such an accusation, Rasputin watches as the
        #  man is chased outside and beaten. Twenty years later, Rasputin sees a vision of
        #  the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous,
        #  with people, even a bishop, begging for his blessing. """

517
        # fmt: off
patrickvonplaten's avatar
patrickvonplaten committed
518
        expected_output_ids = [
519
            67, 2840, 19, 18, 1484, 20, 965, 29077, 8719, 1273, 21, 45, 273, 17, 10, 15048, 28, 27511, 21, 4185, 11, 41, 2444, 9, 32, 1025, 20, 8719, 26, 23, 673, 966, 19, 29077, 20643, 27511, 20822, 20643, 19, 17, 6616, 17511, 18, 8978, 20, 18, 777, 9, 19233, 1527, 17669, 19, 24, 673, 17, 28756, 150, 12943, 4354, 153, 27, 442, 37, 45, 668, 21, 24, 256, 20, 416, 22, 2771, 4901, 9, 12943, 4354, 153, 51, 24, 3004, 21, 28142, 23, 65, 20, 18, 416, 34, 24, 2958, 22947, 9, 1177, 45, 668, 3097, 13768, 23, 103, 28, 441, 148, 48, 20522, 19, 12943, 4354, 153, 12860, 34, 18, 326, 27, 17492, 684, 21, 6709, 9, 8585, 123, 266, 19, 12943, 4354, 153, 6872, 24, 3004, 20, 18, 9225, 2198, 19, 12717, 103, 22, 401, 24, 6348, 9, 12943, 4354, 153, 1068, 2768, 2286, 19, 33, 104, 19, 176, 24, 9313, 19, 20086, 28, 45, 10292, 9, 4, 3, 19, 12943, 4354, 153, 27, 442, 22, 2771, 4901, 9, 69, 27, 442, 22, 2771, 24, 11335, 20, 18, 9225, 2198, 9, 69, 27, 442, 22, 2771, 24, 11335, 20, 18, 9225, 2198, 9, 69, 27, 442, 22, 2771,
patrickvonplaten's avatar
patrickvonplaten committed
520
        ]
521
        # fmt: on
patrickvonplaten's avatar
patrickvonplaten committed
522
523
524
525
526
527
528
529
        #  In 1991, the remains of Russian Tsar Nicholas II and his family (except for Alexei and Maria)
        #  are discovered. The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich,
        #  narrates the remainder of the story. 1883 Western Siberia, a young Grigori Rasputin
        #  is asked by his father and a group of men to perform magic. Rasputin has a vision and
        #  denounces one of the men as a horse thief. Although his father initially slaps
        #  him for making such an accusation, Rasputin watches as the man is chased outside and beaten.
        #  Twenty years later, Rasputin sees a vision of the Virgin Mary, prompting him to become a priest.
        #  Rasputin quickly becomes famous, with people, even a bishop, begging for his blessing.
530
531
        #  <sep><cls>, Rasputin is asked to perform magic. He is asked to perform a ritual of the Virgin Mary.
        #  He is asked to perform a ritual of the Virgin Mary. He is asked to perform
patrickvonplaten's avatar
patrickvonplaten committed
532
533
534

        output_ids = model.generate(input_ids, max_length=200, do_sample=False)

535
        self.assertListEqual(output_ids[0].numpy().tolist(), expected_output_ids)