test_utils.py 131 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright 2020 The HuggingFace Team 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 clone 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.


17
import copy
18
import inspect
19
import tempfile
20
import unittest
21
import warnings
22

23
import numpy as np
24
from parameterized import parameterized
25

26
from transformers import is_torch_available, pipeline, set_seed
27
from transformers.testing_utils import (
28
    is_flaky,
29
30
31
32
33
34
    require_accelerate,
    require_torch,
    require_torch_multi_accelerator,
    slow,
    torch_device,
)
35

36
from ..test_modeling_common import floats_tensor, ids_tensor
37
from .test_framework_agnostic import GenerationIntegrationTestsMixin
38

39
40
41
42

if is_torch_available():
    import torch

43
    from transformers import (
44
        AutoModelForCausalLM,
45
        AutoModelForSeq2SeqLM,
46
47
        AutoModelForSpeechSeq2Seq,
        AutoModelForVision2Seq,
48
        AutoTokenizer,
49
        BartForCausalLM,
50
51
52
53
        BartForConditionalGeneration,
        BartTokenizer,
        GPT2LMHeadModel,
        GPT2Tokenizer,
54
        ImageGPTForCausalImageModeling,
55
        SpeechEncoderDecoderModel,
56
    )
57
    from transformers.cache_utils import DynamicCache
58
59
60
61
62
63
    from transformers.generation import (
        BeamSampleDecoderOnlyOutput,
        BeamSampleEncoderDecoderOutput,
        BeamSearchDecoderOnlyOutput,
        BeamSearchEncoderDecoderOutput,
        DisjunctiveConstraint,
64
65
66
67
        GenerateBeamDecoderOnlyOutput,
        GenerateBeamEncoderDecoderOutput,
        GenerateDecoderOnlyOutput,
        GenerateEncoderDecoderOutput,
68
69
        GreedySearchDecoderOnlyOutput,
        GreedySearchEncoderDecoderOutput,
70
        LogitsProcessorList,
71
        MaxLengthCriteria,
72
        MinLengthLogitsProcessor,
73
74
75
76
77
        PhrasalConstraint,
        SampleDecoderOnlyOutput,
        SampleEncoderDecoderOutput,
        StoppingCriteria,
        StoppingCriteriaList,
78
    )
79
    from transformers.generation.utils import _speculative_sampling
80
81
82
83
84


class GenerationTesterMixin:
    model_tester = None
    all_generative_model_classes = ()
Suraj Patil's avatar
Suraj Patil committed
85
    input_name = "input_ids"
86
    max_new_tokens = 3
87

88
    def _get_input_ids_and_config(self, batch_size=2):
89
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
Suraj Patil's avatar
Suraj Patil committed
90
        input_ids = inputs_dict[self.input_name]
91

92
        input_ids = input_ids[:batch_size]
93
94
95

        if config.eos_token_id is not None and config.pad_token_id is None:
            # hack to allow generate for models such as GPT2 as is done in `generate()`
96
97
98
            if isinstance(config.eos_token_id, int):
                config.eos_token_id = [config.eos_token_id]
            config.pad_token_id = config.eos_token_id[0]
99
        attention_mask = torch.ones_like(input_ids, dtype=torch.long)
100

101
102
103
104
105
        # It is important set set the eos_token_id to None to ensure that no sequences
        # shorter than `max_length` can be generated
        config.eos_token_id = None
        config.forced_eos_token_id = None

106
        return config, input_ids, attention_mask
107
108

    @staticmethod
109
    def _get_logits_processor_and_warper_kwargs(
110
111
112
113
        input_length,
        forced_bos_token_id=None,
        forced_eos_token_id=None,
    ):
114
115
116
        process_kwargs = {
            "bad_words_ids": [[1, 0]],
            "repetition_penalty": 1.2,
117
            "remove_invalid_values": True,
118
        }
119
120
121
122
        # NoRepeatNGramLogitsProcessor + forced tokens may result in no valid continuations
        if forced_bos_token_id is None and forced_eos_token_id is None:
            process_kwargs["no_repeat_ngram_size"] = 2

123
        warp_kwargs = {"top_k": 10, "top_p": 0.7, "temperature": 0.7}
124
        return process_kwargs, warp_kwargs
125
126

    @staticmethod
127
    def _get_beam_kwargs(num_return_sequences=1):
128
129
130
131
132
133
        beam_kwargs = {
            "early_stopping": False,
            "length_penalty": 2.0,
            "num_beams": 2,
            "num_return_sequences": num_return_sequences,
        }
134
        return beam_kwargs
135

136
    @staticmethod
137
    def _get_diverse_beam_kwargs(num_return_sequences=1):
138
139
140
141
142
143
144
145
        beam_kwargs = {
            "early_stopping": False,
            "length_penalty": 2.0,
            "num_beams": 2,
            "num_return_sequences": num_return_sequences,
            "num_beam_groups": 2,  # one beam per group
            "diversity_penalty": 2.0,
        }
146
        return beam_kwargs
147

148
    @staticmethod
149
    def _get_constrained_beam_kwargs(num_return_sequences=1):
150
151
152
153
154
155
        beam_kwargs = {
            "early_stopping": False,
            "length_penalty": 2.0,
            "num_beams": num_return_sequences * 4,
            "num_return_sequences": num_return_sequences,
        }
156
        return beam_kwargs
157

158
    @staticmethod
159
160
161
    def _get_encoder_outputs(
        model, input_ids, attention_mask, output_attentions=None, output_hidden_states=None, num_interleave=1
    ):
162
        encoder = model.get_encoder()
163
164
165
166
167
168
        encoder_outputs = encoder(
            input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
169
170
171
        encoder_outputs["last_hidden_state"] = encoder_outputs.last_hidden_state.repeat_interleave(
            num_interleave, dim=0
        )
172
173
174
        generation_config = copy.deepcopy(model.generation_config)
        model._prepare_special_tokens(generation_config)
        input_ids = torch.zeros_like(input_ids[:, :1]) + generation_config.decoder_start_token_id
175
176
177
        attention_mask = None
        return encoder_outputs, input_ids, attention_mask

178
179
180
181
182
183
    def _greedy_generate(
        self,
        model,
        input_ids,
        attention_mask,
        output_scores=False,
184
        output_logits=False,
185
186
187
188
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
189
        logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
190
191
192
            input_ids.shape[-1],
            forced_bos_token_id=model.config.forced_bos_token_id,
            forced_eos_token_id=model.config.forced_eos_token_id,
193
194
        )

195
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
196
197
198
199
        output_generate = model.generate(
            input_ids,
            do_sample=False,
            num_beams=1,
200
            max_new_tokens=self.max_new_tokens,
201
202
203
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_scores=output_scores,
204
            output_logits=output_logits,
205
206
            return_dict_in_generate=return_dict_in_generate,
            **logits_process_kwargs,
207
            **model_kwargs,
208
209
        )

210
        return output_generate
211
212
213
214
215
216
217
218
219
220

    def _sample_generate(
        self,
        model,
        input_ids,
        attention_mask,
        num_return_sequences,
        logits_warper_kwargs,
        process_kwargs,
        output_scores=False,
221
        output_logits=False,
222
223
224
225
226
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        torch.manual_seed(0)
227
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
228
229
230
231
        output_generate = model.generate(
            input_ids,
            do_sample=True,
            num_beams=1,
232
            max_new_tokens=self.max_new_tokens,
233
234
            num_return_sequences=num_return_sequences,
            output_scores=output_scores,
235
            output_logits=output_logits,
236
237
238
239
240
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **logits_warper_kwargs,
            **process_kwargs,
241
            **model_kwargs,
242
243
        )

244
        return output_generate
245
246
247
248
249
250
251
252
253

    def _beam_search_generate(
        self,
        model,
        input_ids,
        attention_mask,
        beam_kwargs,
        logits_process_kwargs,
        output_scores=False,
254
        output_logits=False,
255
256
257
258
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
259
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
260
261
262
        output_generate = model.generate(
            input_ids,
            do_sample=False,
263
            max_new_tokens=self.max_new_tokens,
264
            output_scores=output_scores,
265
            output_logits=output_logits,
266
267
268
269
270
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **beam_kwargs,
            **logits_process_kwargs,
271
            **model_kwargs,
272
273
        )

274
        return output_generate
275
276
277
278
279
280
281
282
283

    def _beam_sample_generate(
        self,
        model,
        input_ids,
        attention_mask,
        beam_kwargs,
        logits_warper_kwargs,
        output_scores=False,
284
        output_logits=False,
285
286
287
288
289
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        torch.manual_seed(0)
290
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
291
292
293
        output_generate = model.generate(
            input_ids,
            do_sample=True,
294
            max_new_tokens=self.max_new_tokens,
295
            output_scores=output_scores,
296
            output_logits=output_logits,
297
298
299
300
301
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **beam_kwargs,
            **logits_warper_kwargs,
302
            **model_kwargs,
303
304
        )

305
        return output_generate
306
307
308
309
310
311
312
313
314

    def _group_beam_search_generate(
        self,
        model,
        input_ids,
        attention_mask,
        beam_kwargs,
        logits_process_kwargs,
        output_scores=False,
315
        output_logits=False,
316
317
318
319
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
320
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
321
322
323
        output_generate = model.generate(
            input_ids,
            do_sample=False,
324
            max_new_tokens=self.max_new_tokens,
325
            output_scores=output_scores,
326
            output_logits=output_logits,
327
328
329
330
331
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **beam_kwargs,
            **logits_process_kwargs,
332
            **model_kwargs,
333
334
        )

335
        return output_generate
336

337
338
339
340
341
342
343
344
345
    def _constrained_beam_search_generate(
        self,
        model,
        input_ids,
        attention_mask,
        constraints,
        beam_kwargs,
        logits_process_kwargs,
        output_scores=False,
346
        output_logits=False,
347
348
349
350
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
351
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
352
353
354
        output_generate = model.generate(
            input_ids,
            do_sample=False,
355
            max_new_tokens=self.max_new_tokens,
356
            output_scores=output_scores,
357
            output_logits=output_logits,
358
359
360
361
362
363
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            constraints=constraints,
            **beam_kwargs,
            **logits_process_kwargs,
364
            **model_kwargs,
365
366
        )

367
        return output_generate
368

369
370
371
372
373
374
    def _contrastive_generate(
        self,
        model,
        input_ids,
        attention_mask,
        output_scores=False,
375
        output_logits=False,
376
377
378
379
380
381
382
383
384
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        contrastive_search_kwargs = {
            "penalty_alpha": 0.6,
            "top_k": 5,
        }

385
        logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
386
387
388
389
390
391
392
393
394
395
            input_ids.shape[-1],
            forced_bos_token_id=model.config.forced_bos_token_id,
            forced_eos_token_id=model.config.forced_eos_token_id,
        )

        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
        output_generate = model.generate(
            input_ids,
            do_sample=False,
            num_beams=1,
396
            max_new_tokens=self.max_new_tokens,
397
398
399
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_scores=output_scores,
400
            output_logits=output_logits,
401
402
403
404
405
406
            return_dict_in_generate=return_dict_in_generate,
            **logits_process_kwargs,
            **model_kwargs,
            **contrastive_search_kwargs,
        )

407
        return output_generate
408

409
410
    def test_greedy_generate(self):
        for model_class in self.all_generative_model_classes:
411
            config, input_ids, attention_mask = self._get_input_ids_and_config()
412

413
            model = model_class(config).to(torch_device).eval()
414
            output_generate = self._greedy_generate(model=model, input_ids=input_ids, attention_mask=attention_mask)
415

416
417
418
419
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
420

421
422
    def test_greedy_generate_dict_outputs(self):
        for model_class in self.all_generative_model_classes:
423
            config, input_ids, attention_mask = self._get_input_ids_and_config()
424

425
426
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
427
            output_generate = self._greedy_generate(
428
429
430
431
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_scores=True,
432
                output_logits=True,
433
434
435
436
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )
437
438

            if model.config.is_encoder_decoder:
439
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
440
441
                self.assertIsInstance(output_generate, GenerateEncoderDecoderOutput)
                # Retrocompatibility check
442
443
                self.assertIsInstance(output_generate, GreedySearchEncoderDecoderOutput)
            else:
444
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
445
446
                self.assertIsInstance(output_generate, GenerateDecoderOnlyOutput)
                # Retrocompatibility check
447
                self.assertIsInstance(output_generate, GreedySearchDecoderOnlyOutput)
448

449
            self._check_outputs(output_generate, input_ids, model.config)
450
451
452

    def test_greedy_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
453
            config, input_ids, attention_mask = self._get_input_ids_and_config()
454
455

            if not hasattr(config, "use_cache"):
456
                self.skipTest("This model doesn't support caching")
457
458

            config.use_cache = True
459
            config.is_decoder = True
460
            model = model_class(config).to(torch_device).eval()
461
            output_generate = self._greedy_generate(
462
463
                model=model,
                input_ids=input_ids,
464
                attention_mask=attention_mask,
465
                output_scores=True,
466
                output_logits=True,
467
468
469
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
470
            )
471

472
473
474
475
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
476
            self._check_outputs(output_generate, input_ids, model.config, use_cache=True)
477
478
479

    def test_sample_generate(self):
        for model_class in self.all_generative_model_classes:
480
            config, input_ids, attention_mask = self._get_input_ids_and_config()
481

482
483
            model = model_class(config).to(torch_device).eval()
            process_kwargs, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(
484
485
486
487
488
                input_ids.shape[-1],
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
            )

489
            output_generate = self._sample_generate(
490
491
492
493
494
495
496
497
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                num_return_sequences=1,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
            )

498
499
500
501
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
502

503
504
    def test_sample_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
505
            config, input_ids, attention_mask = self._get_input_ids_and_config()
506

507
508
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
509

510
            process_kwargs, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(
511
512
513
                input_ids.shape[-1],
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
514
            )
515

516
            output_generate = self._sample_generate(
517
518
                model=model,
                input_ids=input_ids,
519
                attention_mask=attention_mask,
520
521
522
523
                num_return_sequences=2,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
                output_scores=True,
524
                output_logits=True,
525
526
527
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
528
529
530
            )

            if model.config.is_encoder_decoder:
531
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
532
533
                self.assertIsInstance(output_generate, GenerateEncoderDecoderOutput)
                # Retrocompatibility check
534
                self.assertIsInstance(output_generate, SampleEncoderDecoderOutput)
535
            else:
536
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
537
538
                self.assertIsInstance(output_generate, GenerateDecoderOnlyOutput)
                # Retrocompatibility check
539
540
                self.assertIsInstance(output_generate, SampleDecoderOnlyOutput)

541
            self._check_outputs(output_generate, input_ids, model.config, num_return_sequences=2)
542
543
544

    def test_beam_search_generate(self):
        for model_class in self.all_generative_model_classes:
545
            config, input_ids, attention_mask = self._get_input_ids_and_config()
546

547
            model = model_class(config).to(torch_device).eval()
548

549
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
550
551
552
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
553
            )
554
            beam_kwargs = self._get_beam_kwargs()
555

556
            output_generate = self._beam_search_generate(
557
558
                model=model,
                input_ids=input_ids,
559
                attention_mask=attention_mask,
560
561
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
562
            )
563

564
565
566
567
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
568
569
570

    def test_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
571
            config, input_ids, attention_mask = self._get_input_ids_and_config()
572
573

            # disable cache
574
            config.use_cache = False
575

576
            model = model_class(config).to(torch_device).eval()
577
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
578
579
580
581
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )
582
583
            beam_kwargs = self._get_beam_kwargs()
            output_generate = self._beam_search_generate(
584
585
                model=model,
                input_ids=input_ids,
586
                attention_mask=attention_mask,
587
588
589
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
590
                output_logits=True,
591
592
593
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
594
595
            )
            if model.config.is_encoder_decoder:
596
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
597
598
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
599
                self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput)
600
            else:
601
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
602
603
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
604
605
                self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput)

606
607
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
608
609
610
611
612
            )

    def test_beam_search_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
            # enable cache
613
            config, input_ids, attention_mask = self._get_input_ids_and_config()
614
615

            if not hasattr(config, "use_cache"):
616
                self.skipTest("This model doesn't support caching")
617
618

            model = model_class(config).to(torch_device).eval()
619
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
620
621
622
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
623
624
            )

625
            beam_kwargs = self._get_beam_kwargs()
626
627

            config.use_cache = True
628
            config.is_decoder = True
629
            model = model_class(config).to(torch_device).eval()
630
            output_generate = self._beam_search_generate(
631
632
633
634
635
636
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
637
                output_logits=True,
638
639
640
641
642
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

643
644
645
646
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
647
648
649
            self._check_outputs(
                output_generate, input_ids, model.config, use_cache=True, num_return_sequences=beam_kwargs["num_beams"]
            )
650

651
    @require_accelerate
652
    @require_torch_multi_accelerator
653
654
    def test_model_parallel_beam_search(self):
        for model_class in self.all_generative_model_classes:
655
656
657
            if "xpu" in torch_device:
                return unittest.skip("device_map='auto' does not work with XPU devices")

658
659
660
            if model_class._no_split_modules is None:
                continue

661
            config, input_ids, attention_mask = self._get_input_ids_and_config()
662
663
664
665
666
667
668
669
670

            model = model_class(config).eval()
            with tempfile.TemporaryDirectory() as tmp_dir:
                model.cpu().save_pretrained(tmp_dir)
                new_model = model_class.from_pretrained(tmp_dir, device_map="auto")

                new_model.generate(
                    input_ids,
                    attention_mask=attention_mask,
671
                    max_new_tokens=self.max_new_tokens,
672
673
674
                    num_beams=2,
                )

675
676
    def test_beam_sample_generate(self):
        for model_class in self.all_generative_model_classes:
677
            config, input_ids, attention_mask = self._get_input_ids_and_config()
678

679
            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])
680

681
            model = model_class(config).to(torch_device).eval()
682
            beam_kwargs = self._get_beam_kwargs()
683

684
            output_generate = self._beam_sample_generate(
685
686
                model=model,
                input_ids=input_ids,
687
                attention_mask=attention_mask,
688
689
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
690
            )
691

692
693
694
695
696
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])

697
698
699
700
701
702
703
704
705
706
707
708
            if "inputs_embeds" in set(inspect.signature(model.prepare_inputs_for_generation).parameters):
                input_embeds = model.get_input_embeddings()(input_ids)
                beam_kwargs.update({"inputs_embeds": input_embeds})
                output_generate2 = self._beam_sample_generate(
                    model=model,
                    input_ids=None,
                    attention_mask=attention_mask,
                    beam_kwargs=beam_kwargs,
                    logits_warper_kwargs=logits_warper_kwargs,
                )

                torch.testing.assert_close(output_generate[:, input_embeds.shape[1] :], output_generate2)
709
710
711

    def test_beam_sample_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
712
            config, input_ids, attention_mask = self._get_input_ids_and_config()
713
714

            # disable cache
715
            config.use_cache = False
716

717
            model = model_class(config).to(torch_device).eval()
718
719
            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])
            beam_kwargs = self._get_beam_kwargs()
720

721
            output_generate = self._beam_sample_generate(
722
723
724
725
726
727
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
                output_scores=True,
728
                output_logits=True,
729
730
731
732
733
734
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            if model.config.is_encoder_decoder:
735
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
736
737
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
738
                self.assertIsInstance(output_generate, BeamSampleEncoderDecoderOutput)
739
            else:
740
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
741
742
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
743
                self.assertIsInstance(output_generate, BeamSampleDecoderOnlyOutput)
744

745
746
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
747
            )
748

749
    def test_generate_without_input_ids(self):
750
        config, _, _ = self._get_input_ids_and_config()
751

752
753
754
        # if no bos token id => cannot generate from None
        if config.bos_token_id is None:
            return
755

756
757
758
759
        # hack in case they are equal, otherwise the attn mask will be [0]
        if config.bos_token_id == config.pad_token_id:
            config.pad_token_id = None

760
761
762
        for model_class in self.all_generative_model_classes:
            model = model_class(config).to(torch_device)
            model.eval()
763

764
765
766
            output_ids_generate = model.generate(
                do_sample=False, max_new_tokens=self.max_new_tokens, remove_invalid_values=True
            )
767
            self.assertIsNotNone(output_ids_generate)
768

769
770
    def test_group_beam_search_generate(self):
        for model_class in self.all_generative_model_classes:
771
            config, input_ids, attention_mask = self._get_input_ids_and_config()
772

773
            model = model_class(config).to(torch_device).eval()
774
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
775
776
777
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
778
779
780
            )

            # check `generate()` and `group_beam_search()` are equal
781
782
            beam_kwargs = self._get_diverse_beam_kwargs()
            output_generate = self._group_beam_search_generate(
783
784
                model=model,
                input_ids=input_ids,
785
                attention_mask=attention_mask,
786
787
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
788
            )
789
790
791
792
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
793

794
            # check `group_beam_search` for higher than 1 `num_return_sequences`
795
            num_return_sequences = 2
796
797
            beam_kwargs = self._get_diverse_beam_kwargs(num_return_sequences=num_return_sequences)
            output_generate = self._group_beam_search_generate(
798
799
800
801
802
803
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
804
805
806
807
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
808

809
810
    def test_group_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
811
            config, input_ids, attention_mask = self._get_input_ids_and_config()
812
            config.use_cache = False
813

814
            model = model_class(config).to(torch_device).eval()
815
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
816
817
818
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
819
820
            )

821
822
            beam_kwargs = self._get_diverse_beam_kwargs()
            output_generate = self._group_beam_search_generate(
823
824
                model=model,
                input_ids=input_ids,
825
                attention_mask=attention_mask,
826
827
828
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
829
                output_logits=True,
830
831
832
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
833
834
            )
            if model.config.is_encoder_decoder:
835
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
836
837
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
838
                self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput)
839
            else:
840
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
841
842
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
843
844
                self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput)

845
846
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
847
848
            )

849
850
    # TODO: @gante
    @is_flaky()
851
852
    def test_constrained_beam_search_generate(self):
        for model_class in self.all_generative_model_classes:
853
            config, input_ids, attention_mask = self._get_input_ids_and_config()
854
855
856

            model = model_class(config).to(torch_device).eval()

857
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
858
859
860
861
862
863
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )

            # Sample constraints
864
865
            min_id = 3
            max_id = config.vocab_size
866

867
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
868
869
870
871
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

872
873
            beam_kwargs = self._get_constrained_beam_kwargs()
            output_generate = self._constrained_beam_search_generate(
874
875
876
877
878
879
880
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                constraints=constraints,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
881
882
883
884
885
886

            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])

887
888
889
            for generation_output in output_generate:
                self._check_sequence_inside_sequence(force_tokens, generation_output)

890
            # check`constrained_beam_search` for higher than 1 `num_return_sequences`
891
            # Sample constraints
892
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
893
894
895
896
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

897
            beam_kwargs = self._get_constrained_beam_kwargs(num_return_sequences=2)
898

899
            output_generate = self._constrained_beam_search_generate(
900
901
902
903
904
905
906
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                constraints=constraints,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
907
908
909
910
911

            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
912
913
914
915
916
917

            for generation_output in output_generate:
                self._check_sequence_inside_sequence(force_tokens, generation_output)

    def test_constrained_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
918
            config, input_ids, attention_mask = self._get_input_ids_and_config()
919
920
921
922
923

            # disable cache
            config.use_cache = False

            model = model_class(config).to(torch_device).eval()
924
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
925
926
927
928
929
930
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )

            # Sample constraints
931
932
            min_id = 3
            max_id = model.config.vocab_size
933
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
934
935
936
937
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

938
939
            beam_kwargs = self._get_constrained_beam_kwargs()
            output_generate = self._constrained_beam_search_generate(
940
941
942
943
944
945
946
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                constraints=constraints,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
947
                output_logits=True,
948
949
950
951
952
953
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

            if model.config.is_encoder_decoder:
954
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
955
956
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
957
958
                self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput)
            else:
959
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
960
961
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
962
963
                self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput)

964
965
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
966
967
            )

968
969
970
    def test_contrastive_generate(self):
        for model_class in self.all_generative_model_classes:
            # won't fix: FSMT and Reformer have a different cache variable type (and format).
971
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
972
                self.skipTest("Won't fix: old model with different cache format")
973

974
            config, input_ids, attention_mask = self._get_input_ids_and_config()
975
976
977

            # NOTE: contrastive search only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
978
                self.skipTest("This model doesn't support caching")
979
980
981
982
983
            config.use_cache = True
            config.is_decoder = True

            # test old generation output for backwards compatibility
            model = model_class(config).to(torch_device).eval()
984
            output_generate = self._contrastive_generate(
985
                model=model, input_ids=input_ids, attention_mask=attention_mask
986
            )
987
988
989
990
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
991
992
993
994

    def test_contrastive_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
            # won't fix: FSMT and Reformer have a different cache variable type (and format).
995
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
996
                self.skipTest("Won't fix: old model with different cache format")
997

998
            config, input_ids, attention_mask = self._get_input_ids_and_config()
999
1000
1001

            # NOTE: contrastive search only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
1002
                self.skipTest("This model doesn't support caching")
1003
1004
1005
1006
            config.use_cache = True
            config.is_decoder = True

            model = model_class(config).to(torch_device).eval()
1007
            output_generate = self._contrastive_generate(
1008
1009
1010
1011
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_scores=True,
1012
                output_logits=True,
1013
1014
1015
1016
1017
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

1018
1019
1020
1021
            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
1022
            self._check_outputs(output_generate, input_ids, model.config, use_cache=True)
1023

1024
1025
1026
    def test_contrastive_generate_low_memory(self):
        # Check that choosing 'low_memory' does not change the model output
        for model_class in self.all_generative_model_classes:
1027
1028
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer", "speech2text"]):
                self.skipTest("Won't fix: old model with different cache format")
tomeras91's avatar
tomeras91 committed
1029
            if any(model_name in model_class.__name__.lower() for model_name in ["gptbigcode", "jamba"]):
1030
                self.skipTest("TODO: fix me")
1031

1032
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1033
1034
1035

            # NOTE: contrastive search only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
1036
                self.skipTest("This model doesn't support caching")
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048

            config.use_cache = True
            config.is_decoder = True

            # test output equality of low versus high memory
            model = model_class(config).to(torch_device).eval()

            low_output = model.generate(
                input_ids,
                top_k=4,
                penalty_alpha=0.6,
                low_memory=True,
1049
                max_new_tokens=self.max_new_tokens,
1050
1051
1052
1053
1054
1055
1056
1057
                attention_mask=attention_mask,
            )

            high_output = model.generate(
                input_ids,
                top_k=4,
                penalty_alpha=0.6,
                low_memory=False,
1058
                max_new_tokens=self.max_new_tokens,
1059
1060
1061
1062
                attention_mask=attention_mask,
            )
            self.assertListEqual(low_output.tolist(), high_output.tolist())

1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
    def test_beam_search_low_memory(self):
        # Check that choosing 'low_memory' does not change the model output
        for model_class in self.all_generative_model_classes:
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
                self.skipTest("Won't fix: old model with different cache format")
            if any(
                model_name in model_class.__name__.lower()
                for model_name in [
                    "bloom",
                    "ctrl",
                    "gptbigcode",
                    "transo_xl",
                    "xlnet",
                    "cpm",
tomeras91's avatar
tomeras91 committed
1077
                    "jamba",
1078
1079
1080
                ]
            ):
                self.skipTest("May fix in the future: need model-specific fixes")
1081
            config, input_ids, _ = self._get_input_ids_and_config(batch_size=2)
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
            # batch_size=1 is ok, but batch_size>1 will cause non-identical output

            config.use_cache = True
            config.is_decoder = True

            # test output equality of low versus high memory
            model = model_class(config).to(torch_device).eval()

            low_output = model.generate(input_ids, max_new_tokens=8, num_beams=5, early_stopping=True, low_memory=True)

            high_output = model.generate(
                input_ids, max_new_tokens=8, num_beams=5, early_stopping=True, low_memory=False
            )
            self.assertListEqual(low_output.tolist(), high_output.tolist())

1097
    @parameterized.expand([("random",), ("same",)])
1098
    @is_flaky()  # Read NOTE (1) below. If there are API issues, all attempts will fail.
1099
    def test_assisted_decoding_matches_greedy_search(self, assistant_type):
1100
        # This test ensures that the assisted generation does not introduce output changes over greedy search.
1101
1102
1103
1104
1105
        # NOTE (1): The sentence above is true most of the time, there is a tiny difference in the logits due to matmul
        # shape differences -- and it may result in a different output. The input shape difference happens in the
        # main model, that runs the forward pass with several candidates at once (as opposed to generating one token at
        # a time). See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535 for more info.
        # NOTE (2): It breaks the pattern in the tests above, for multiple reasons:
1106
        # - assisted_decoding, contrarily to the other methods, can't be called on its own (e.g. needs to
1107
        # prepare the assistant encoder outputs in the main generate body);
1108
1109
        # - assisted_decoding does not support `use_cache = False`
        # - assisted_decoding does not support `batch_size > 1`
1110
1111
1112

        for model_class in self.all_generative_model_classes:
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
1113
                self.skipTest("Won't fix: old model with different cache format")
1114
1115
            if any(
                model_name in model_class.__name__.lower()
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
1126
            ):
1127
                self.skipTest("May fix in the future: need model-specific fixes")
1128

1129
            # enable cache
1130
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1131

1132
1133
1134
            # NOTE: assisted generation only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
                self.skipTest("This model doesn't support caching")
1135

1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
            # Sets assisted generation arguments such that:
            # a) no EOS is generated, to ensure generation doesn't break early
            # b) the assistant model always generates two tokens when it is called, to ensure the input preparation of
            #    the assistant model is correct
            # c) there are at least two forward passes in the main model, to ensure the input preparation of
            #    the main model is correct
            generation_kwargs = {
                "eos_token_id": -1,  # see a)
                "max_new_tokens": 4,  # see c)
                "num_beams": 1,
                "do_sample": False,
                "output_scores": True,
1151
                "output_logits": True,
1152
1153
1154
1155
1156
1157
                "output_hidden_states": True,
                "output_attentions": True,
                "return_dict_in_generate": True,
            }
            output_greedy = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

1158
1159
1160
1161
1162
1163
1164
            # test with the same assistant model or randomly init one
            # in the first case all candidate tokens are accepted, in the second none is accepted
            # case when some are accepted and some not is hard to reproduce, so let's hope this catches most errors :)
            if assistant_type == "random":
                assistant_model = model_class(config).to(torch_device).eval()
            else:
                assistant_model = model
1165
1166
1167
1168
1169
1170
1171
1172
1173
            assistant_model.generation_config.num_assistant_tokens = 2  # see b)
            assistant_model.generation_config.num_assistant_tokens_schedule = "constant"  # see b)
            generation_kwargs.update({"assistant_model": assistant_model})
            output_assisted = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

            # The two outputs must match and their shape must be as expected
            self.assertListEqual(output_greedy.sequences.tolist(), output_assisted.sequences.tolist())
            for output in (output_greedy, output_assisted):
                self._check_outputs(output, input_ids, model.config, use_cache=True)
1174

1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
    @is_flaky()
    def test_prompt_lookup_decoding_matches_greedy_search(self):
        # This test ensures that the prompt lookup generation does not introduce output changes over greedy search.
        # This test is mostly a copy of test_assisted_decoding_matches_greedy_search

        for model_class in self.all_generative_model_classes:
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
                self.skipTest("Won't fix: old model with different cache format")
            if any(
                model_name in model_class.__name__.lower()
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
            ):
                self.skipTest("May fix in the future: need model-specific fixes")

            # enable cache
1199
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219

            # NOTE: assisted generation only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
                self.skipTest("This model doesn't support caching")

            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
            # Sets assisted generation arguments such that:
            # a) no EOS is generated, to ensure generation doesn't break early
            # b) the prompt lookup tries to give the model 2 tokens, to ensure the input preparation of
            #    prompt lookup is correct
            # c) there are at least two forward passes in the main model, to ensure the input preparation of
            #    the main model is correct
            generation_kwargs = {
                "eos_token_id": -1,  # see a)
                "max_new_tokens": 4,  # see c)
                "num_beams": 1,
                "do_sample": False,
                "output_scores": True,
1220
                "output_logits": True,
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
                "output_hidden_states": True,
                "output_attentions": True,
                "return_dict_in_generate": True,
            }

            output_greedy = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

            generation_kwargs.update({"prompt_lookup_num_tokens": 2})  # see b)
            output_prompt_lookup = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

            # The two outputs must match and their shape must be as expected
            self.assertListEqual(output_greedy.sequences.tolist(), output_prompt_lookup.sequences.tolist())
            for output in (output_greedy, output_prompt_lookup):
                self._check_outputs(output, input_ids, model.config, use_cache=True)

1236
    def test_assisted_decoding_sample(self):
1237
1238
1239
        # In this test we don't check assisted vs non-assisted output -- seeded assisted decoding with sample will not
        # match sample for the same seed, as the forward pass does not return the exact same logits (due to matmul with
        # different shapes, see https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535).
1240
1241
        for model_class in self.all_generative_model_classes:
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
1242
                self.skipTest("Won't fix: old model with different cache format")
1243
1244
            if any(
                model_name in model_class.__name__.lower()
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
1255
            ):
1256
                self.skipTest("May fix in the future: need model-specific fixes")
1257
1258

            # enable cache
1259
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1260
1261
1262

            # NOTE: assisted generation only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
1263
                self.skipTest("This model doesn't support caching")
1264
1265
1266
1267

            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
            # Sets assisted generation arguments such that:
            # a) no EOS is generated, to ensure generation doesn't break early
            # b) the assistant model always generates two tokens when it is called, to ensure the input preparation of
            #    the assistant model is correct
            # c) there are at least two forward passes in the main model, to ensure the input preparation of
            #    the main model is correct
            assistant_model = model
            assistant_model.generation_config.num_assistant_tokens = 2  # see b)
            assistant_model.generation_config.num_assistant_tokens_schedule = "constant"  # see b)
            generation_kwargs = {
                "eos_token_id": -1,  # see a)
                "max_new_tokens": 4,  # see c)
                "num_beams": 1,
                "do_sample": True,
                "assistant_model": assistant_model,
                "output_scores": True,
1284
                "output_logits": True,
1285
1286
1287
1288
1289
                "output_hidden_states": True,
                "output_attentions": True,
                "return_dict_in_generate": True,
            }
            output_assisted = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)
1290
1291
1292

            self._check_outputs(output_assisted, input_ids, model.config, use_cache=True)

1293
1294
1295
1296
    def test_generate_with_head_masking(self):
        """Test designed for encoder-decoder models to ensure the attention head masking is used."""
        attention_names = ["encoder_attentions", "decoder_attentions", "cross_attentions"]
        for model_class in self.all_generative_model_classes:
1297
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1298
1299
1300
            # We want to test only encoder-decoder models
            if not config.is_encoder_decoder:
                continue
Joao Gante's avatar
Joao Gante committed
1301
            model = model_class(config).to(torch_device)
1302
1303

            head_masking = {
1304
1305
1306
1307
1308
1309
1310
                "head_mask": torch.zeros(config.encoder_layers, config.encoder_attention_heads, device=torch_device),
                "decoder_head_mask": torch.zeros(
                    config.decoder_layers, config.decoder_attention_heads, device=torch_device
                ),
                "cross_attn_head_mask": torch.zeros(
                    config.decoder_layers, config.decoder_attention_heads, device=torch_device
                ),
1311
1312
1313
1314
            }

            signature = inspect.signature(model.forward)
            # We want to test only models where encoder/decoder head masking is implemented
1315
            if not set(head_masking.keys()) < {*signature.parameters.keys()}:
1316
1317
1318
1319
1320
                continue

            for attn_name, (name, mask) in zip(attention_names, head_masking.items()):
                out = model.generate(
                    input_ids,
1321
                    attention_mask=attention_mask,
1322
1323
1324
                    num_beams=1,
                    output_attentions=True,
                    return_dict_in_generate=True,
1325
                    remove_invalid_values=True,
1326
1327
1328
1329
1330
1331
                    **{name: mask},
                )
                # We check the state of decoder_attentions and cross_attentions just from the last step
                attn_weights = out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
                self.assertEqual(sum([w.sum().item() for w in attn_weights]), 0.0)

1332
    def test_left_padding_compatibility(self):
1333
1334
        # NOTE: left-padding results in small numerical differences. This is expected.
        # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
1335

1336
1337
1338
1339
1340
1341
1342
        # First, filter out models that don't support left padding
        # - The model must have generative capabilities
        if len(self.all_generative_model_classes) == 0:
            self.skipTest(reason="No generative architecture available for this model.")

        # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
        decoder_only_classes = []
1343
        for model_class in self.all_generative_model_classes:
1344
            config, _, _ = self._get_input_ids_and_config()
1345
            if config.is_encoder_decoder:
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
                continue
            else:
                decoder_only_classes.append(model_class)
        if len(decoder_only_classes) == 0:
            self.skipTest(reason="No decoder-only architecture available for this model.")

        # - Decoder-only architectures derived from encoder-decoder models could support it in theory, but we haven't
        #   added support for it yet. We skip these models for now.
        has_encoder_attributes = any(
            attr_name
            for attr_name in config.to_dict().keys()
            if attr_name.startswith("encoder") and attr_name != "encoder_no_repeat_ngram_size"
        )
        if has_encoder_attributes:
            self.skipTest(
                reason="The decoder-only derived from encoder-decoder models are not expected to support left-padding."
            )

        # Then, test left-padding
        def _prepare_model_kwargs(input_ids, attention_mask, signature):
            model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
            if "position_ids" in signature:
                position_ids = torch.cumsum(attention_mask, dim=-1) - 1
                position_ids.masked_fill_(attention_mask == 0, 1)
                model_kwargs["position_ids"] = position_ids
            if "cache_position" in signature:
                cache_position = torch.arange(input_ids.shape[-1], device=torch_device)
                model_kwargs["cache_position"] = cache_position
            return model_kwargs

        for model_class in decoder_only_classes:
1377
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1378
1379
1380
            model = model_class(config).to(torch_device).eval()
            signature = inspect.signature(model.forward).parameters.keys()

1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
            # Without padding
            model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)
            next_logits_wo_padding = model(**model_kwargs).logits[:, -1, :]

            # With left-padding (length 32)
            pad_size = (input_ids.shape[0], 32)
            padding = torch.ones(pad_size, dtype=input_ids.dtype, device=torch_device) * config.pad_token_id
            padded_input_ids = torch.cat((padding, input_ids), dim=1)
            padded_attention_mask = torch.cat((torch.zeros_like(padding), attention_mask), dim=1)
            model_kwargs = _prepare_model_kwargs(padded_input_ids, padded_attention_mask, signature)
            next_logits_with_padding = model(**model_kwargs).logits[:, -1, :]

            # They should result in very similar logits
            self.assertTrue(torch.allclose(next_logits_wo_padding, next_logits_with_padding, atol=1e-5))
1395

1396
1397
1398
1399
1400
1401
1402
1403
    def test_past_key_values_format(self):
        # Test that the KV cache is formatted correctly. Exceptions need to explicitly overwrite this test. Having a
        # standard KV cache format is important for a consistent API (and for advanced generation methods).
        for model_class in self.all_generative_model_classes:
            config, inputs = self.model_tester.prepare_config_and_inputs_for_common()

            # If it doesn't support cache, pass the test
            if not hasattr(config, "use_cache"):
1404
                self.skipTest("This model doesn't support caching")
1405
1406
1407
1408
1409
1410
1411
1412

            model = model_class(config).to(torch_device)
            if "use_cache" not in inputs:
                inputs["use_cache"] = True
            outputs = model(**inputs)

            # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
            if "past_key_values" not in outputs:
1413
                self.skipTest("This model doesn't return `past_key_values`")
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466

            num_hidden_layers = (
                getattr(config, "decoder_layers", None)
                or getattr(config, "num_decoder_layers", None)
                or config.num_hidden_layers
            )
            num_attention_heads = getattr(config, "decoder_attention_heads", config.num_attention_heads)
            embed_dim = getattr(config, "d_model", config.hidden_size)
            per_head_embed_dim = embed_dim // num_attention_heads

            past_kv = outputs["past_key_values"]
            self.assertEqual(len(past_kv), num_hidden_layers)

            # Encoder-Decoder checks
            if config.is_encoder_decoder:
                encoder_num_attention_heads = config.encoder_attention_heads
                encoder_per_head_embed_dim = embed_dim // encoder_num_attention_heads
                batch_size, seq_length = inputs["decoder_input_ids"].shape
                for i in range(num_hidden_layers):
                    self.assertEqual(len(past_kv[i]), 4)  # K V for the decoder + K V for the encoder = 4
                    self.assertEqual(
                        past_kv[i][0].shape, (batch_size, num_attention_heads, seq_length, per_head_embed_dim)
                    )
                    self.assertEqual(
                        past_kv[i][1].shape, (batch_size, num_attention_heads, seq_length, per_head_embed_dim)
                    )
                    # The sequence length for the encoder K V depends on the model. Since it is not manipulated in
                    # autoregressive generation, I'm keeping the test general and not checking the 3rd dim
                    self.assertEqual(
                        (past_kv[i][2].shape[0], past_kv[i][2].shape[1], past_kv[i][2].shape[3]),
                        (batch_size, encoder_num_attention_heads, encoder_per_head_embed_dim),
                    )
                    self.assertEqual(
                        (past_kv[i][3].shape[0], past_kv[i][3].shape[1], past_kv[i][3].shape[3]),
                        (batch_size, encoder_num_attention_heads, encoder_per_head_embed_dim),
                    )

            # Decoder-only checks
            else:
                # TODO: this line is only needed because of imagegpt, where "pixel_values" = "input_ids". Fix the
                # tests in imagegpt such that `prepare_config_and_inputs_for_common` returns the later (and the other
                # tests use it)
                key = "input_ids" if "input_ids" in inputs else "pixel_values"
                batch_size, seq_length = inputs[key].shape
                for i in range(num_hidden_layers):
                    self.assertEqual(len(past_kv[0]), 2)  # K V for the decoder = 2
                    self.assertEqual(
                        past_kv[i][0].shape, (batch_size, num_attention_heads, seq_length, per_head_embed_dim)
                    )
                    self.assertEqual(
                        past_kv[i][1].shape, (batch_size, num_attention_heads, seq_length, per_head_embed_dim)
                    )

1467
1468
1469
1470
    def test_generate_from_inputs_embeds_decoder_only(self):
        # When supported, tests that the decoder model can generate from `inputs_embeds` instead of `input_ids`
        # if fails, you should probably update the `prepare_inputs_for_generation` function
        for model_class in self.all_generative_model_classes:
1471
            config, input_ids, _ = self._get_input_ids_and_config()
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513

            # Ignore:
            # a) eos (to always output 20 tokens) and pad (so we don't try to infer the attn mask from the input_ids,
            #   which would cause a mismatch),
            config.pad_token_id = config.eos_token_id = -1
            # b) embedding scaling, the scaling factor applied after embeding from input_ids (requires knowledge of the
            #   variable that holds the scaling factor, which is model-dependent)
            if hasattr(config, "scale_embedding"):
                config.scale_embedding = False

            # This test is for decoder-only models (encoder-decoder models have native input embeddings support in the
            # decoder)
            if config.is_encoder_decoder:
                continue

            # Skip models without explicit support
            model = model_class(config).to(torch_device).eval()
            if "inputs_embeds" not in inspect.signature(model.prepare_inputs_for_generation).parameters.keys():
                continue

            # Traditional way of generating text
            outputs_from_ids = model.generate(input_ids)
            self.assertEqual(outputs_from_ids.shape, (2, 20))

            # Same thing, but from input embeddings (`input_ids` is passed so the prompt is present in the output)
            inputs_embeds = model.get_input_embeddings()(input_ids)
            outputs_from_embeds = model.generate(input_ids, inputs_embeds=inputs_embeds)
            self.assertListEqual(outputs_from_ids.tolist(), outputs_from_embeds.tolist())

            # But if we pass different inputs_embeds, we should get different outputs
            torch.manual_seed(0)
            random_embeds = torch.rand_like(inputs_embeds)
            outputs_from_rand_embeds = model.generate(input_ids, inputs_embeds=random_embeds)
            with self.assertRaises(AssertionError):
                self.assertListEqual(outputs_from_rand_embeds.tolist(), outputs_from_embeds.tolist())

            # input_ids is not a required input -- if we don't pass it, the newly generated tokens will be the same
            outputs_from_embeds_wo_ids = model.generate(
                inputs_embeds=inputs_embeds, max_new_tokens=20 - inputs_embeds.shape[1]
            )
            self.assertListEqual(
                outputs_from_embeds[:, inputs_embeds.shape[1] :].tolist(),
1514
                outputs_from_embeds_wo_ids.tolist(),
1515
1516
            )

1517
1518
1519
1520
    def test_generate_continue_from_past_key_values(self):
        # Tests that we can continue generating from past key values, returned from a previous `generate` call
        for model_class in self.all_generative_model_classes:
            if any(model_name in model_class.__name__.lower() for model_name in ["imagegpt"]):
1521
                self.skipTest("Won't fix: old model with unique inputs/caches/other")
1522
            if any(model_name in model_class.__name__.lower() for model_name in ["umt5"]):
1523
                self.skipTest("TODO: needs modeling or test input preparation fixes for compatibility")
1524
1525
1526
1527

            config, inputs = self.model_tester.prepare_config_and_inputs_for_common()

            if not hasattr(config, "use_cache"):
1528
                self.skipTest("This model doesn't support caching")
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546

            # Let's make it always:
            # 1. use cache (for obvious reasons)
            # 2. generate to max length (which can be achieved by setting the eos token to an invalid value), which
            #    would make the test flaky (e.g. EOS is generated on iteration 1 on both generations, but the
            #    continuation would force it to generate beyond an EOS token)
            # 3. ignore `token_type_ids` for simplicity
            # 4. ignore `forced_eos_token_id`, which requires further manipulation of the continuation inputs and is
            #    active by default on some models
            config.use_cache = True
            if "token_type_ids" in inputs:
                del inputs["token_type_ids"]

            model = model_class(config).to(torch_device)
            model.eval()
            model.generation_config.pad_token_id = model.generation_config.eos_token_id = -1
            model.generation_config.forced_eos_token_id = None

1547
            # If "past_key_values" is not returned, skip the test (e.g. RWKV uses a different cache name and format)
1548
1549
            outputs = model(**inputs)
            if "past_key_values" not in outputs:
1550
                self.skipTest("This model doesn't return `past_key_values`")
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592

            # Traditional way of generating text, with `return_dict_in_generate` to return the past key values
            outputs = model.generate(**inputs, do_sample=False, max_new_tokens=4, return_dict_in_generate=True)

            # Let's generate again, but passing the past key values in between (3 + 1 = 4 tokens). Note that the
            # inputs may need to be tweaked across `generate` calls (like the attention mask).
            outputs_cached = model.generate(**inputs, do_sample=False, max_new_tokens=3, return_dict_in_generate=True)

            # Continue from the tokens generated above, preparing the inputs accordingly
            inputs["past_key_values"] = outputs_cached.past_key_values
            new_attention_len = outputs_cached.sequences.shape[-1]
            if config.is_encoder_decoder:
                inputs["decoder_input_ids"] = outputs_cached.sequences
                if "decoder_attention_mask" in inputs:
                    inputs["decoder_attention_mask"] = torch.nn.functional.pad(
                        inputs["decoder_attention_mask"],
                        (0, new_attention_len - inputs["decoder_attention_mask"].shape[1]),
                        mode="constant",
                        value=1,
                    )
            else:
                inputs["input_ids"] = outputs_cached.sequences
                if "attention_mask" in inputs:
                    inputs["attention_mask"] = torch.nn.functional.pad(
                        inputs["attention_mask"],
                        (0, new_attention_len - inputs["attention_mask"].shape[1]),
                        mode="constant",
                        value=1,
                    )
            outputs_cached = model.generate(**inputs, do_sample=False, max_new_tokens=1, return_dict_in_generate=True)

            # The two sets of generated text and past kv should be equal to each other
            self.assertListEqual(outputs.sequences.tolist(), outputs_cached.sequences.tolist())
            for layer_idx in range(len(outputs_cached.past_key_values)):
                for kv_idx in range(len(outputs_cached.past_key_values[layer_idx])):
                    self.assertTrue(
                        torch.allclose(
                            outputs.past_key_values[layer_idx][kv_idx],
                            outputs_cached.past_key_values[layer_idx][kv_idx],
                        )
                    )

1593
1594
1595
1596
1597
1598
1599
1600
1601
    @parameterized.expand([(1, False), (1, True), (4, False)])
    def test_new_cache_format(self, num_beams, do_sample):
        # Tests that generating with the new format is exactly the same as the legacy one (for models that support it).
        # 馃憠 tests with and without beam search so that we can test with and without cache reordering.
        # 馃憠 tests with and without sampling so we can cover the most common use cases.
        for model_class in self.all_generative_model_classes:
            if not model_class._supports_cache_class:
                self.skipTest("This model does not support the new cache format")

1602
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
            config.use_cache = True
            config.is_decoder = True

            model = model_class(config).to(torch_device).eval()
            generation_kwargs = {
                "max_new_tokens": 5,
                "do_sample": do_sample,
                "num_beams": num_beams,
                "num_return_sequences": num_beams,
                "return_dict_in_generate": True,  # Required to return `past_key_values`
            }

            # Sets seed before calling `generate` for the case with do_sample=True
            seed = torch.randint(0, 1000000, (1,)).item()
            set_seed(seed)
            legacy_results = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)
            set_seed(seed)
            new_results = model.generate(
                input_ids, attention_mask=attention_mask, past_key_values=DynamicCache(), **generation_kwargs
            )

            # The two sets of generated sequences must match, despite the cache format between forward passes being
            # different
            self.assertListEqual(legacy_results.sequences.tolist(), new_results.sequences.tolist())
            self.assertTrue(isinstance(legacy_results.past_key_values, tuple))
            self.assertTrue(isinstance(new_results.past_key_values, DynamicCache))

            # The contents of the two caches, when converted to the same format (in both directions!), must match
            legacy_cache = legacy_results.past_key_values
            new_cache_converted = new_results.past_key_values.to_legacy_cache()
            for layer_idx in range(len(legacy_cache)):
                for kv_idx in range(len(legacy_cache[layer_idx])):
                    self.assertTrue(
                        torch.allclose(
                            legacy_cache[layer_idx][kv_idx],
                            new_cache_converted[layer_idx][kv_idx],
                        )
                    )

            new_cache = new_results.past_key_values
            legacy_cache_converted = DynamicCache.from_legacy_cache(legacy_results.past_key_values)
            for layer_idx in range(len(new_cache)):
                for kv_idx in range(len(new_cache[layer_idx])):
                    self.assertTrue(
                        torch.allclose(
                            new_cache[layer_idx][kv_idx],
                            legacy_cache_converted[layer_idx][kv_idx],
                        )
                    )

1653
1654
1655
    def _check_outputs(self, output, input_ids, config, use_cache=False, num_return_sequences=1):
        batch_size, seq_length = input_ids.shape
        num_sequences_in_output = batch_size * num_return_sequences
1656

1657
1658
1659
1660
1661
1662
1663
        gen_len = (
            output.sequences.shape[-1] - 1 if config.is_encoder_decoder else output.sequences.shape[-1] - seq_length
        )

        # scores
        self._check_scores(num_sequences_in_output, output.scores, length=gen_len, config=config)

1664
1665
1666
        # unprocessed logits
        self._check_logits(num_sequences_in_output, output.logits, config=config)

1667
1668
1669
        # Attentions
        if config.is_encoder_decoder:
            # encoder
1670
            self._check_encoder_attention_for_generate(output.encoder_attentions, batch_size, config, seq_length)
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
            # decoder
            self._check_attentions_for_generate(
                num_sequences_in_output,
                output.decoder_attentions,
                min_length=1,
                max_length=output.sequences.shape[-1],
                config=config,
                use_cache=use_cache,
            )
        else:
            # if use_cache first input is equal to no use_cache, so skip here
            attentions = output.attentions if not use_cache else output.attentions[1:]
            min_length = seq_length if not use_cache else seq_length + 1
            self._check_attentions_for_generate(
                num_sequences_in_output,
                attentions=attentions,
                min_length=min_length,
                max_length=output.sequences.shape[-1],
                config=config,
                use_cache=use_cache,
            )

        # Hidden States
        if config.is_encoder_decoder:
            # encoder
1696
1697
            self._check_encoder_hidden_states_for_generate(
                output.encoder_hidden_states, batch_size, config, seq_length
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
            )

            # decoder
            self._check_hidden_states_for_generate(
                num_sequences_in_output,
                output.decoder_hidden_states,
                min_length=1,
                max_length=output.sequences.shape[-1],
                config=config,
                use_cache=use_cache,
            )
        else:
            # if use_cache first input is equal to no use_cache, so skip here
            hidden_states = output.hidden_states if not use_cache else output.hidden_states[1:]
            min_length = seq_length if not use_cache else seq_length + 1
            self._check_hidden_states_for_generate(
                num_sequences_in_output,
                hidden_states,
                min_length=min_length,
                max_length=output.sequences.shape[-1],
                config=config,
                use_cache=use_cache,
            )

tomeras91's avatar
tomeras91 committed
1722
        # Past Key Value States -- a few notes here:
1723
1724
        # 1. Its inner sequence length is with respect to the inputs of the latest forward pass, hence the "-1"
        # 2. Some old models still return `output.past_key_values` even without `use_cache=True`
tomeras91's avatar
tomeras91 committed
1725
1726
1727
        # 3. TODO (joao): A few models have different formats/types, skipping those until the cache refactor is
        # complete
        models_without_standard_cache = ("bloom", "ctrl", "fsmt", "gptbigcode", "mega", "reformer", "jamba")
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
        has_standard_cache = not any(
            model_name in config.__class__.__name__.lower() for model_name in models_without_standard_cache
        )
        if use_cache and has_standard_cache:
            past_key_values = output.past_key_values
            past_sequence_length = output.sequences.shape[-1] - 1
            self._check_past_key_values_for_generate(
                num_sequences_in_output,
                past_key_values,
                seq_length=past_sequence_length,
                config=config,
            )

1741
1742
1743
1744
1745
1746
    def _check_scores(self, batch_size, scores, length, config):
        expected_shape = (batch_size, config.vocab_size)
        self.assertIsInstance(scores, tuple)
        self.assertEqual(len(scores), length)
        self.assertListEqual([iter_scores.shape for iter_scores in scores], [expected_shape] * len(scores))

1747
1748
1749
1750
1751
1752
1753
1754
    def _check_logits(self, batch_size, scores, config):
        self.assertIsInstance(scores, tuple)
        self.assertListEqual([iter_scores.shape[0] for iter_scores in scores], [batch_size] * len(scores))
        # vocabulary difference equal to one (imagegptmodel?) or zero (all other models)
        vocab_diff = config.vocab_size - scores[0].shape[-1]
        self.assertTrue(vocab_diff in [0, 1])
        self.assertListEqual([config.vocab_size - score.shape[-1] for score in scores], [vocab_diff] * len(scores))

1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
    def _check_attentions_for_generate(
        self, batch_size, attentions, min_length, max_length, config, use_cache=False, num_beam_groups=1
    ):
        self.assertIsInstance(attentions, tuple)
        self.assertListEqual(
            [isinstance(iter_attentions, tuple) for iter_attentions in attentions], [True] * len(attentions)
        )
        self.assertEqual(len(attentions), (max_length - min_length) * num_beam_groups)

        for idx, iter_attentions in enumerate(attentions):
            tgt_len = min_length + idx if not use_cache else 1
            src_len = min_length + idx

            expected_shape = (
                batch_size * num_beam_groups,
                config.num_attention_heads,
                tgt_len,
                src_len,
            )
            # check attn size
            self.assertListEqual(
                [layer_attention.shape for layer_attention in iter_attentions], [expected_shape] * len(iter_attentions)
            )

1779
1780
1781
1782
1783
1784
1785
1786
    def _check_encoder_attention_for_generate(self, attentions, batch_size, config, seq_length):
        encoder_expected_shape = (batch_size, config.num_attention_heads, seq_length, seq_length)
        self.assertIsInstance(attentions, tuple)
        self.assertListEqual(
            [layer_attentions.shape for layer_attentions in attentions],
            [encoder_expected_shape] * len(attentions),
        )

1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
    def _check_hidden_states_for_generate(
        self, batch_size, hidden_states, min_length, max_length, config, use_cache=False, num_beam_groups=1
    ):
        self.assertIsInstance(hidden_states, tuple)
        self.assertListEqual(
            [isinstance(iter_hidden_states, tuple) for iter_hidden_states in hidden_states],
            [True] * len(hidden_states),
        )
        self.assertEqual(len(hidden_states), (max_length - min_length) * num_beam_groups)

        for idx, iter_hidden_states in enumerate(hidden_states):
            seq_len = min_length + idx if not use_cache else 1
            expected_shape = (batch_size * num_beam_groups, seq_len, config.hidden_size)
            # check hidden size
            self.assertListEqual(
                [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states],
                [expected_shape] * len(iter_hidden_states),
            )
1805

1806
1807
1808
1809
1810
1811
1812
1813
    def _check_encoder_hidden_states_for_generate(self, hidden_states, batch_size, config, seq_length):
        encoder_expected_shape = (batch_size, seq_length, config.hidden_size)
        self.assertIsInstance(hidden_states, tuple)
        self.assertListEqual(
            [layer_hidden_states.shape for layer_hidden_states in hidden_states],
            [encoder_expected_shape] * len(hidden_states),
        )

1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
    def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config, num_beam_groups=1):
        self.assertIsInstance(past_key_values, tuple)
        self.assertListEqual(
            [isinstance(iter_past_key_values, tuple) for iter_past_key_values in past_key_values],
            [True] * len(past_key_values),
        )

        # (batch, head, seq_length, head_features)
        expected_shape = (
            batch_size * num_beam_groups,
            config.num_key_value_heads if hasattr(config, "num_key_value_heads") else config.num_attention_heads,
            seq_length,
            config.hidden_size // config.num_attention_heads,
        )
        # check shape key, value
        self.assertListEqual(
            [layer_past_key_values[0].shape for layer_past_key_values in past_key_values],
            [expected_shape] * len(past_key_values),
        )
        self.assertListEqual(
            [layer_past_key_values[1].shape for layer_past_key_values in past_key_values],
            [expected_shape] * len(past_key_values),
        )

1838
    def _check_sequence_inside_sequence(self, tensor_1, tensor_2):
1839
        # check if tensor_1 inside tensor_2 or tensor_2 inside tensor_1.
1840
1841
        # set to same device. we don't care what device.

1842
1843
1844
1845
1846
1847
        if not isinstance(tensor_1, list):
            tensor_1 = tensor_1.cpu().tolist()
        if not isinstance(tensor_2, list):
            tensor_2 = tensor_2.cpu().tolist()

        in_order = len(tensor_1) <= len(tensor_2)
1848
1849
1850
1851
        longer = tensor_2 if in_order else tensor_1
        shorter = tensor_1 if in_order else tensor_2

        flag = False
1852
1853
        chunk_size = len(shorter)
        for chunk_idx in range(len(longer) - chunk_size + 1):
1854
            subseq = longer[chunk_idx : chunk_idx + chunk_size]
1855
            if subseq == shorter:
1856
1857
1858
1859
1860
                flag = True
                break

        self.assertTrue(flag)

1861
1862
1863

@require_torch
class UtilsFunctionsTest(unittest.TestCase):
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
    def test_speculative_sampling(self):
        # assume vocab size 10, input length 5 + 3 generated candidates
        candidate_input_ids = torch.tensor([[8, 0, 3, 9, 8, 1, 4, 5]])  # input tokens
        candidate_logits = torch.tensor(
            [
                [
                    [-10.0, 10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0],  # generated 1
                    [-10.0, -10.0, -10.0, -10.0, 10.0, -10.0, -10.0, -10.0, -10.0, -10.0],  # generated 4
                    [-10.0, -10.0, -10.0, -10.0, -10.0, 10.0, -10.0, -10.0, -10.0, -10.0],  # generated 5
                ]
            ]
        )
        candidate_length = 3
        inf = float("inf")
        new_logits = torch.tensor(
            [
                [
                    [-10.0, 10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0],  # accepts 1
                    [-10.0, -10.0, -10.0, -10.0, 10.0, -10.0, -10.0, -10.0, -10.0, -10.0],  # accepts 4
                    [-inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, 10.0, -inf],  # rejects 5, accepts 8
                    [-10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0, -10.0],  # N/A
                ]
            ]
        )
        last_assistant_token_is_eos = False
        validated_tokens, n_matches = _speculative_sampling(
            candidate_input_ids,
            candidate_logits,
            candidate_length,
            new_logits,
            last_assistant_token_is_eos,
        )
        self.assertTrue(n_matches.item() == 2)
        self.assertTrue(validated_tokens.tolist()[0] == [1, 4, 8])

1899
1900

@require_torch
1901
1902
1903
1904
class GenerationIntegrationTests(unittest.TestCase, GenerationIntegrationTestsMixin):
    # setting framework_dependent_parameters needs to be gated, just like its contents' imports
    if is_torch_available():
        framework_dependent_parameters = {
1905
            "AutoModelForCausalLM": AutoModelForCausalLM,
1906
            "AutoModelForSpeechSeq2Seq": AutoModelForSpeechSeq2Seq,
1907
            "AutoModelForSeq2SeqLM": AutoModelForSeq2SeqLM,
1908
            "AutoModelForVision2Seq": AutoModelForVision2Seq,
1909
1910
            "LogitsProcessorList": LogitsProcessorList,
            "MinLengthLogitsProcessor": MinLengthLogitsProcessor,
1911
            "create_tensor_fn": torch.tensor,
1912
            "floats_tensor": floats_tensor,
1913
1914
1915
            "return_tensors": "pt",
        }

1916
1917
    @slow
    def test_diverse_beam_search(self):
1918
        # PT-only test: TF doesn't have a diverse beam search implementation
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
        article = """Justin Timberlake and Jessica Biel, welcome to parenthood.
        The celebrity couple announced the arrival of their son, Silas Randall Timberlake, in statements to People.
        "Silas was the middle name of Timberlake's maternal grandfather Bill Bomar, who died in 2012, while Randall is the musician's own middle name, as well as his father's first," People reports.
        The couple announced the pregnancy in January, with an Instagram post. It is the first baby for both."""

        bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
        bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn").to(torch_device)
        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids.to(torch_device)

        outputs = bart_model.generate(
1929
1930
1931
1932
1933
1934
            input_ids,
            num_beams=4,
            num_return_sequences=2,
            num_beam_groups=4,
            diversity_penalty=2.0,
            remove_invalid_values=True,
1935
1936
1937
1938
1939
1940
1941
        )

        generated_text = bart_tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
Sylvain Gugger's avatar
Sylvain Gugger committed
1942
1943
1944
1945
1946
1947
                "The couple announced the birth of their son, Silas Randall Timberlake, in a statement. Silas was the"
                " middle name of Timberlake's maternal grandfather Bill Bomar. Randall is the musician's own middle"
                " name, as well as his father's first. It is the first baby for both of them.",
                "Justin Timberlake and Jessica Biel have a son. The baby is named Silas Randall Timberlake. It is the"
                " first child for both. The couple announced the pregnancy in January. The name Silas is the middle"
                " name of Timberlake's maternal grandfather. It's also his own middle name.",
1948
1949
            ],
        )
1950

1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
    def test_max_length_if_input_embeds(self):
        # PT-only test: TF doesn't have StoppingCriteria
        article = "Today a dragon flew over Paris."
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        input_ids = tokenizer(article, return_tensors="pt").input_ids.to(torch_device)
        inputs_embeds = model.get_input_embeddings()(input_ids)

        max_length = 20
        input_len = input_ids.shape[-1]
        out_gen = model.generate(input_ids=input_ids, max_length=max_length)
        out_gen_embeds = model.generate(inputs_embeds=inputs_embeds, max_length=max_length)
        self.assertEqual(out_gen.shape[-1], input_len + out_gen_embeds.shape[-1])

1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
    def test_min_length_if_input_embeds(self):
        # PT-only test: TF doesn't have StoppingCriteria
        article = "Today a dragon flew over Paris."
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        input_ids = tokenizer(article, return_tensors="pt").input_ids.to(torch_device)
        inputs_embeds = model.get_input_embeddings()(input_ids)

        min_length = 10
        input_len = input_ids.shape[-1]
        out_gen = model.generate(input_ids=input_ids, min_length=min_length)
        out_gen_embeds = model.generate(inputs_embeds=inputs_embeds, min_length=min_length)
        self.assertEqual(out_gen.shape[-1], input_len + out_gen_embeds.shape[-1])

1979
    def test_custom_stopping_criteria_overload_error(self):
1980
        # PT-only test: TF doesn't have StoppingCriteria
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
        article = """Justin Timberlake and Jessica Biel, welcome to parenthood."""
        bart_tokenizer = BartTokenizer.from_pretrained("sshleifer/bart-tiny-random")
        bart_model = BartForConditionalGeneration.from_pretrained("sshleifer/bart-tiny-random").to(torch_device)

        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids.to(torch_device)
        stopping_criteria = StoppingCriteriaList()
        stopping_criteria.append(MaxLengthCriteria(max_length=42))
        with self.assertRaises(ValueError):
            bart_model.generate(input_ids, stopping_criteria=stopping_criteria)
        with self.assertRaises(ValueError):
            bart_model.generate(input_ids, stopping_criteria=stopping_criteria, max_length=32)

    def test_custom_stopping_criteria(self):
1994
        # PT-only test: TF doesn't have StoppingCriteria
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
        article = """Justin Timberlake and Jessica Biel, welcome to parenthood."""
        bart_tokenizer = BartTokenizer.from_pretrained("sshleifer/bart-tiny-random")
        bart_model = BartForConditionalGeneration.from_pretrained("sshleifer/bart-tiny-random").to(torch_device)
        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids.to(torch_device)

        class DummyCriteria(StoppingCriteria):
            def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
                return input_ids.shape[-1] >= 20

        stopping_criteria = StoppingCriteriaList()
        stopping_criteria.append(DummyCriteria())

        self.assertEqual(
            list(bart_model.generate(input_ids, stopping_criteria=stopping_criteria, max_length=22).shape),
            [1, 20],
        )
        self.assertEqual(
            list(bart_model.generate(input_ids, stopping_criteria=stopping_criteria, max_length=18).shape),
            [1, 18],
        )

2016
    def test_stop_sequence_stopping_criteria(self):
2017
        # PT-only test: TF doesn't have StoppingCriteria
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
        prompt = """Hello I believe in"""
        generator = pipeline("text-generation", model="hf-internal-testing/tiny-random-bart")
        output = generator(prompt)
        self.assertEqual(
            output,
            [
                {
                    "generated_text": (
                        "Hello I believe in in in number number number number number number number number number"
                    )
                }
            ],
        )

        output = generator(prompt, stop_sequence=" number")
        self.assertEqual(output, [{"generated_text": "Hello I believe in in in number"}])

2035
    def test_generate_non_nlp_input_ids_as_kwarg(self):
2036
        # PT-only test: AFAIK there's no non-NLP model architecture in TF that supports `input_ids` as its only input
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
        model = ImageGPTForCausalImageModeling.from_pretrained(
            "hf-internal-testing/tiny-random-imagegpt", max_length=10
        ).to(torch_device)
        input_ids = ids_tensor((3, 5), vocab_size=10)

        output_sequences_kwargs = model.generate(input_ids=input_ids).cpu()
        output_sequences = model.generate(input_ids).cpu()

        self.assertListEqual(output_sequences.tolist(), output_sequences_kwargs.tolist())
        self.assertEqual(output_sequences.shape, (3, 10))

2048
    def test_generate_input_values_as_encoder_kwarg(self):
2049
        # PT-only test: AFAIK there's no generate-capable architecture in TF that supports `input_values` as its input
2050
2051
2052
2053
2054
2055
2056
2057
2058
        input_values = floats_tensor((2, 250))
        model = SpeechEncoderDecoderModel.from_pretrained("hf-internal-testing/tiny-random-speech-encoder-decoder")
        model = model.to(torch_device)
        output_sequences_kwargs = model.generate(input_values=input_values, max_length=5).cpu()
        output_sequences = model.generate(input_values, max_length=5).cpu()

        self.assertListEqual(output_sequences.tolist(), output_sequences_kwargs.tolist())
        self.assertEqual(output_sequences.shape, (2, 5))

2059
    def test_transition_scores_group_beam_search_encoder_decoder(self):
2060
        # PT-only test: TF doesn't have group beam search
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
        articles = [
            "Justin Timberlake and Jessica Biel, welcome to parenthood.",
            "Michael Phelps is arguably the most decorated Olympian of all time.",
        ]
        tokenizer = BartTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
        model = BartForConditionalGeneration.from_pretrained(
            "hf-internal-testing/tiny-random-bart",
            max_length=10,
            num_beams=2,
            num_beam_groups=2,
            num_return_sequences=2,
2072
            diversity_penalty=1.0,
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
            eos_token_id=None,
            return_dict_in_generate=True,
            output_scores=True,
            length_penalty=0.0,
        )
        model = model.to(torch_device)

        input_ids = tokenizer(articles, return_tensors="pt", padding=True).input_ids.to(torch_device)
        outputs = model.generate(input_ids=input_ids)

2083
        transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, outputs.beam_indices)
2084
2085
2086
        transition_scores_sum = transition_scores.sum(-1)

        self.assertTrue(torch.allclose(transition_scores_sum, outputs.sequences_scores, atol=1e-3))
2087

2088
    def test_beam_search_low_memory(self):
2089
2090
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
        model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
        tokenizer.pad_token_id = tokenizer.eos_token_id
        model_inputs = tokenizer("I", return_tensors="pt")["input_ids"]

        low_output = model.generate(model_inputs, max_new_tokens=40, num_beams=5, early_stopping=True, low_memory=True)

        high_output = model.generate(
            model_inputs, max_new_tokens=40, num_beams=5, early_stopping=True, low_memory=False
        )
        self.assertListEqual(low_output.tolist(), high_output.tolist())

2101
2102
    @slow
    def test_beam_search_example_integration(self):
2103
        # PT-only test: TF doesn't have a BeamSearchScorer
2104
2105
        # exactly the example provided in the docstrings of beam search, which previously
        # failed after directly copying from it. Refer to PR #15555
2106
2107
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2108
2109
2110
2111
2112
2113
2114

        encoder_input_str = "translate English to German: How old are you?"
        encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

        # lets run beam search using 3 beams
        num_beams = 3
        # define decoder start token ids
2115
        input_ids = torch.ones((1, 1), device=model.device, dtype=torch.long)
2116
2117
2118
        input_ids = input_ids * model.config.decoder_start_token_id

        # add encoder_outputs to model keyword arguments
2119
        model_kwargs = {"encoder_outputs": model.get_encoder()(encoder_input_ids, return_dict=True)}
2120

2121
2122
        outputs = model.generate(
            input_ids, num_beams=num_beams, min_length=5, eos_token_id=model.config.eos_token_id, **model_kwargs
2123
2124
2125
2126
2127
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(outputs, ["Wie alt bist du?"])

2128
2129
    @slow
    def test_constrained_beam_search(self):
2130
        # PT-only test: TF doesn't have constrained beam search
2131
2132
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2133

2134
2135
        force_tokens = tokenizer("scared", add_prefix_space=True, add_special_tokens=False).input_ids
        force_tokens_2 = tokenizer("big weapons", add_prefix_space=True, add_special_tokens=False).input_ids
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160

        constraints = [
            PhrasalConstraint(force_tokens),
            PhrasalConstraint(force_tokens_2),
        ]

        starting_text = ["The soldiers were not prepared and"]

        input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)

        outputs = model.generate(
            input_ids,
            constraints=constraints,
            num_beams=10,
            num_return_sequences=1,
            no_repeat_ngram_size=1,
            max_length=30,
            remove_invalid_values=True,
        )

        generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
2161
2162
                "The soldiers were not prepared and didn't know what to do. They had no idea how they would react if"
                " the enemy attacked them, big weapons scared"
2163
2164
2165
            ],
        )

2166
2167
    @slow
    def test_constrained_beam_search_mixed(self):
2168
        # PT-only test: TF doesn't have constrained beam search
2169
2170
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200

        force_phrase = tokenizer("scared", add_prefix_space=True, add_special_tokens=False).input_ids
        flexible_phrases = tokenizer(
            ["scream", "screams", "screaming", "screamed"], add_prefix_space=True, add_special_tokens=False
        ).input_ids

        constraints = [
            PhrasalConstraint(force_phrase),
            DisjunctiveConstraint(flexible_phrases),
        ]

        starting_text = ["The soldiers", "The child"]

        input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)

        outputs = model.generate(
            input_ids,
            constraints=constraints,
            num_beams=10,
            num_return_sequences=1,
            no_repeat_ngram_size=1,
            # max_length=20,
            remove_invalid_values=True,
        )

        generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
2201
2202
2203
                "The soldiers, who had been stationed at the base for more than a year before being evacuated"
                " screaming scared",
                "The child was taken to a local hospital where he died.\n 'I don't think screaming scared",
2204
2205
2206
2207
2208
            ],
        )

    @slow
    def test_constrained_beam_search_mixed_mixin(self):
2209
        # PT-only test: TF doesn't have constrained beam search
2210
2211
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238

        force_word = "scared"
        force_flexible = ["scream", "screams", "screaming", "screamed"]

        force_words_ids = [
            tokenizer([force_word], add_prefix_space=True, add_special_tokens=False).input_ids,
            tokenizer(force_flexible, add_prefix_space=True, add_special_tokens=False).input_ids,
        ]

        starting_text = ["The soldiers", "The child"]

        input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)

        outputs = model.generate(
            input_ids,
            force_words_ids=force_words_ids,
            num_beams=10,
            num_return_sequences=1,
            no_repeat_ngram_size=1,
            remove_invalid_values=True,
        )

        generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
2239
2240
2241
                "The soldiers, who had been stationed at the base for more than a year before being evacuated"
                " screaming scared",
                "The child was taken to a local hospital where he died.\n 'I don't think screaming scared",
2242
2243
2244
            ],
        )

2245
2246
    @slow
    def test_cfg_mixin(self):
2247
2248
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284

        input = tokenizer(["The dragon flew over Paris,"], return_tensors="pt", return_attention_mask=True)
        input["input_ids"] = input["input_ids"].to(torch_device)
        input["attention_mask"] = input["attention_mask"].to(torch_device)

        outputs = model.generate(**input, max_new_tokens=32, guidance_scale=1.5)
        generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
                "The dragon flew over Paris, landing in the Rue de la Bastille. The crowd was so excited "
                'that they had to leave the city.\n\n"We\'re going to Paris!"\n'
            ],
        )

        neg = tokenizer(["France,"], return_tensors="pt", return_attention_mask=True)
        neg["input_ids"] = neg["input_ids"].to(torch_device)
        neg["attention_mask"] = neg["attention_mask"].to(torch_device)
        outputs = model.generate(
            **input,
            max_new_tokens=32,
            guidance_scale=1.5,
            negative_prompt_ids=neg["input_ids"],
            negative_prompt_attention_mask=neg["attention_mask"],
        )
        generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

        self.assertListEqual(
            generated_text,
            [
                'The dragon flew over Paris, landing on the pavement.\n\n"Paris!"\n\n"Paris!"\n\n"'
                'Paris!"\n\n"Paris!"\n\n"Paris!"\n\n'
            ],
        )

2285
2286
    @slow
    def test_constrained_beam_search_example_translation_mixin(self):
2287
        # PT-only test: TF doesn't have constrained beam search
2288
2289
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307

        encoder_input_str = "translate English to German: How old are you?"
        force_words = ["sind"]

        input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids
        force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids

        outputs = model.generate(
            input_ids,
            force_words_ids=force_words_ids,
            num_beams=10,
            num_return_sequences=1,
            no_repeat_ngram_size=1,
            remove_invalid_values=True,
        )

        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

2308
        self.assertListEqual(outputs, ["Wie alt sind Sie?"])
2309

2310
2311
    @slow
    def test_constrained_beam_search_example_integration(self):
2312
        # PT-only test: TF doesn't have constrained beam search
2313
2314
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2315
2316
2317
2318
2319
2320
2321

        encoder_input_str = "translate English to German: How old are you?"
        encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

        # lets run beam search using 5 beams
        num_beams = 5
        # define decoder start token ids
2322
        input_ids = torch.ones((1, 1), device=model.device, dtype=torch.long)
2323
2324
2325
        input_ids = input_ids * model.config.decoder_start_token_id

        # add encoder_outputs to model keyword arguments
2326
        model_kwargs = {"encoder_outputs": model.get_encoder()(encoder_input_ids, return_dict=True)}
2327
2328
2329
2330

        constraint_str = "sind"
        constraint_token_ids = tokenizer.encode(constraint_str)[:-1]  # remove eos token

2331
2332
2333
2334
2335
2336
2337
        outputs = model.generate(
            input_ids,
            num_beams=num_beams,
            force_words_ids=[constraint_token_ids],
            min_length=5,
            eos_token_id=model.config.eos_token_id,
            **model_kwargs,
2338
2339
2340
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

2341
        self.assertListEqual(outputs, ["Wie alt sind Sie?"])
2342

2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
    @slow
    def test_per_row_stopping_criteria(self):
        text = [
            "They completed the challenging puzzle, revealing the hidden",
            "Today a dragon flew over France",
            "The aroma of freshly baked pizza filled the kitchen",
        ]
        stop_strings = ["secrets"]

        model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
        tokenizer.padding_side = "left"
        tokenizer.pad_token_id = tokenizer.eos_token_id
        input_ids = tokenizer(text, return_tensors="pt", padding="longest", add_special_tokens=False).input_ids.to(
            torch_device
        )

        # normal generation with one stopping criteria
        out = model.generate(input_ids, max_length=15)
        out_text = tokenizer.batch_decode(out)
        expected_out = [
            "They completed the challenging puzzle, revealing the hidden secrets of the world.\n",
            "<|endoftext|><|endoftext|><|endoftext|>Today a dragon flew over France and the French government was forced",
            "The aroma of freshly baked pizza filled the kitchen with a sense of freshness",
        ]
        self.assertListEqual(out_text, expected_out)

        # generation should stop at "secrets" for first batch only, filling the rest with eos tokens
        out = model.generate(input_ids, max_length=15, stop_strings=stop_strings, tokenizer=tokenizer)
        out_text = tokenizer.batch_decode(out)
        expected_out = [
            "They completed the challenging puzzle, revealing the hidden secrets<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>",
            "<|endoftext|><|endoftext|><|endoftext|>Today a dragon flew over France and the French government was forced",
            "The aroma of freshly baked pizza filled the kitchen with a sense of freshness",
        ]
        self.assertListEqual(out_text, expected_out)

2380
    def test_constrained_beam_search_mixin_type_checks(self):
2381
        # PT-only test: TF doesn't have constrained beam search
2382
2383
        tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/t5-tiny-random")
        model = AutoModelForSeq2SeqLM.from_pretrained("patrickvonplaten/t5-tiny-random")
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419

        encoder_input_str = "translate English to German: How old are you?"
        input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

        with self.assertRaises(ValueError):
            force_words = ["sind"]
            force_words_ids = tokenizer(force_words, return_tensors="pt").input_ids
            model.generate(
                input_ids,
                force_words_ids=force_words_ids,
                num_beams=10,
                num_return_sequences=1,
                no_repeat_ngram_size=1,
                remove_invalid_values=True,
            )

        with self.assertRaises(ValueError):
            force_words = ["sind"]
            force_words_ids = [tokenizer(force_words, return_tensors="pt").input_ids]
            model.generate(
                input_ids,
                force_words_ids=force_words_ids,
                num_beams=10,
                num_return_sequences=1,
                no_repeat_ngram_size=1,
                remove_invalid_values=True,
            )

        with self.assertRaises(ValueError):
            model.generate(input_ids, force_words_ids=[])

        with self.assertRaises(ValueError):
            model.generate(input_ids, force_words_ids=[[-1]])

        with self.assertRaises(ValueError):
            model.generate(input_ids, force_words_ids=[[[-1]]])
2420

2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
    def test_batched_decoder_start_id(self):
        # PT-only test: TF doesn't support batched_decoder_start_id
        articles = [
            "Justin Timberlake and Jessica Biel, welcome to parenthood.",
            "Michael Phelps is arguably the most decorated Olympian of all time.",
        ]
        bart_tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
        bart_model = BartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart").to(
            torch_device
        )
        input_ids = bart_tokenizer(articles, return_tensors="pt", padding=True).input_ids.to(torch_device)
        decoder_start_token_id = bart_model.generation_config.decoder_start_token_id
        decoder_start_token_id_batch = [decoder_start_token_id] * input_ids.shape[0]

        outputs = bart_model.generate(input_ids, decoder_start_token_id=decoder_start_token_id)

        outputs_batched_ids = bart_model.generate(input_ids, decoder_start_token_id=decoder_start_token_id_batch)

        self.assertListEqual(outputs.tolist(), outputs_batched_ids.tolist())

2441
    def test_contrastive_search_batched(self):
2442
        # PT-only test: TF doesn't have constrained beam search
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
        # Tests that contrastive search works with batched inputs (i.e. has the same output as for non-batched inputs)
        articles = ["Foo", "Bar Baz"]
        tokenizer = BartTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
        model = BartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart").to(torch_device)

        model.config.eos_token_id = None
        input_ids_batched = tokenizer(articles, padding=True, return_tensors="pt").input_ids.to(torch_device)
        input_ids = tokenizer(articles[1], return_tensors="pt").input_ids.to(torch_device)

        output_sequences_batched = model.generate(
            input_ids=input_ids_batched, penalty_alpha=0.6, top_k=4, return_dict_in_generate=True, output_scores=True
        )
        output_sequences = model.generate(
            input_ids=input_ids, penalty_alpha=0.6, top_k=4, return_dict_in_generate=True, output_scores=True
        )

        batched_out = tokenizer.decode(output_sequences_batched.sequences[1], skip_special_tokens=True)
        out = tokenizer.decode(output_sequences.sequences[0], skip_special_tokens=True)
        self.assertEqual(batched_out, out)

        # output_sequences_batched.scores[0][1] -> 1st set of logits, 2nd sequence
        max_score_diff = (output_sequences_batched.scores[0][1] - output_sequences.scores[0][0]).abs().max()
        self.assertTrue(max_score_diff < 1e-5)

2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
    def test_logits_processor_not_inplace(self):
        # PT-only test: TF fixes were not made
        article = "Today a dragon flew over Paris."
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        input_ids = tokenizer(article, return_tensors="pt").input_ids.to(torch_device)

        out = model.generate(input_ids, output_logits=True, output_scores=True, return_dict_in_generate=True)
        out_with_temp = model.generate(
            input_ids,
            temperature=0.5,
            do_sample=True,
            output_logits=True,
            output_scores=True,
            return_dict_in_generate=True,
        )

        # if no logits processor is used, scores == logits. Otherwise, the processor has to modify the scores
        self.assertListEqual(out.logits[-1].tolist(), out.scores[-1].tolist())
        self.assertNotEqual(out_with_temp.logits[-1].tolist(), out_with_temp.scores[-1].tolist())

2488
    def test_eos_token_id_int_and_list_top_k_top_sampling(self):
2489
        # Has TF equivalent: this test relies on random sampling
2490
2491
2492
2493
2494
2495
2496
        generation_kwargs = {
            "do_sample": True,
            "num_beams": 1,
            "top_p": 0.7,
            "top_k": 10,
            "temperature": 0.7,
        }
2497
        expectation = 20
2498

2499
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
2500
        text = """Hello, my dog is cute and"""
2501
        tokens = tokenizer(text, return_tensors="pt").to(torch_device)
2502
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
2503

2504
2505
2506
        # Only some seeds will work both on CPU/GPU for a fixed `expectation` value.
        # The selected seed is not guaranteed to work on all torch versions.
        torch.manual_seed(1)
2507
2508
2509
2510
        eos_token_id = 846
        generated_tokens = model.generate(**tokens, eos_token_id=eos_token_id, **generation_kwargs)
        self.assertTrue(expectation == len(generated_tokens[0]))

2511
        torch.manual_seed(1)
2512
        eos_token_id = [846, 198]
2513
2514
        generated_tokens = model.generate(**tokens, eos_token_id=eos_token_id, **generation_kwargs)
        self.assertTrue(expectation == len(generated_tokens[0]))
2515

2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
    def test_model_kwarg_encoder_signature_filtering(self):
        # Has TF equivalent: ample use of framework-specific code
        bart_tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
        article = """Hugging Face is a technology company based in New York and Paris."""
        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids.to(torch_device)
        bart_model = BartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart").to(
            torch_device
        )
        output = bart_model.generate(input_ids).cpu().numpy()

        # Let's create a fake model that has a different signature. In particular, this fake model accepts "foo" as an
        # argument. Because "foo" is not in the encoder signature and doesn't start with "decoder_", it will be part of
        # the encoder kwargs prior to signature filtering, which would lead to an exception. But filtering kicks in and
        # saves the day.
        class FakeBart(BartForConditionalGeneration):
            def forward(self, input_ids, foo=None, **kwargs):
                return super().forward(input_ids, **kwargs)

        bart_model = FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart").to(torch_device)
        fake_output = bart_model.generate(input_ids, foo="bar").cpu().numpy()
        self.assertTrue(np.array_equal(output, fake_output))

        # Encoder signature filtering only kicks in if it doesn't accept wildcard kwargs. The following test will fail
        # because it doesn't do signature filtering.
        class FakeEncoder(bart_model.model.encoder.__class__):
            def forward(self, input_ids, **kwargs):
                return super().forward(input_ids, **kwargs)

        fake_encoder = FakeEncoder(bart_model.config, bart_model.model.shared).to(torch_device)
        bart_model.model.encoder = fake_encoder

        # Normal generation still works (the output will be different because the encoder weights are different)
        fake_output = bart_model.generate(input_ids).cpu().numpy()
        with self.assertRaises(TypeError):
            # FakeEncoder.forward() accepts **kwargs -> no filtering -> type error due to unexpected input "foo"
            bart_model.generate(input_ids, foo="bar")
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572

    def test_default_max_length_warning(self):
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        model.config.pad_token_id = tokenizer.eos_token_id

        text = "Hello world"
        tokenized_inputs = tokenizer([text], return_tensors="pt")
        input_ids = tokenized_inputs.input_ids.to(torch_device)

        # Default generation config value of 20 -> emits warning
        with self.assertWarns(UserWarning):
            model.generate(input_ids)

        # Explicitly setting max_length to 20 -> no warning
        with warnings.catch_warnings(record=True) as warning_list:
            model.generate(input_ids, max_length=20)
            self.assertEqual(len(warning_list), 0)

        # Generation config max_length != 20 -> no warning
        with warnings.catch_warnings(record=True) as warning_list:
2573
            # generation_config is modified -> legacy mode is disabled = generation_config takes precedence
2574
2575
2576
            model.generation_config.max_length = 10
            model.generate(input_ids)
            self.assertEqual(len(warning_list), 0)
2577

2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
    def test_length_warning_assisted_generation(self):
        # PT-only test: TF doesn't support assisted decoding yet.
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        model.config.pad_token_id = tokenizer.eos_token_id
        assistant.config.pad_token_id = tokenizer.eos_token_id

        text = "Hello world"
        tokenized_inputs = tokenizer([text], return_tensors="pt")
        input_ids = tokenized_inputs.input_ids.to(torch_device)

        # This should not raise any warning that min length is not feasible in candidate generation
        with warnings.catch_warnings(record=True) as warning_list:
            model.generate(
                input_ids,
                assistant_model=assistant,
                min_new_tokens=10,
                max_length=20,
            )
            self.assertEqual(len(warning_list), 0)

    def test_generated_length_assisted_generation(self):
        # PT-only test: TF doesn't support assisted decoding yet.
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        model.config.pad_token_id = tokenizer.eos_token_id
        assistant.config.pad_token_id = tokenizer.eos_token_id

        text = "Hello world"
        tokenized_inputs = tokenizer([text], return_tensors="pt")
        input_ids = tokenized_inputs.input_ids.to(torch_device)
        input_length = input_ids.shape[-1]

        out = model.generate(
            input_ids,
            assistant_model=assistant,
            min_new_tokens=10,
            max_new_tokens=20,
        )
        self.assertTrue((10 + input_length) <= out.shape[-1] <= (20 + input_length))

        out = model.generate(
            input_ids,
            assistant_model=assistant,
            min_new_tokens=10,
        )
        self.assertTrue((input_length + 10) <= out.shape[-1] <= 20)

2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
    def test_model_kwarg_assisted_decoding_decoder_only(self):
        # PT-only test: TF doesn't support assisted decoding yet.
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        model.config.pad_token_id = tokenizer.eos_token_id

        text = "Hello world"
        tokenized_inputs = tokenizer([text], return_tensors="pt")
        input_ids = tokenized_inputs.input_ids.to(torch_device)

        # Traditional way of generating text
        outputs_normal = model.generate(input_ids)
        self.assertEqual(outputs_normal.shape, (1, 20))

        # Should be different with token_type_ids
        outputs_tti = model.generate(
            input_ids,
            token_type_ids=torch.zeros(input_ids.shape, dtype=torch.long).to(torch_device),
        )
        with self.assertRaises(AssertionError):
            self.assertListEqual(outputs_tti.tolist(), outputs_normal.tolist())

        # Assistant model
        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        assistant.config.pad_token_id = tokenizer.eos_token_id

        # If assisted generation passes model_kwargs correctly, should be same as previous
        outputs_assisted = model.generate(
            input_ids,
            token_type_ids=torch.zeros(input_ids.shape, dtype=torch.long).to(torch_device),
            assistant_model=assistant,
        )
        self.assertListEqual(outputs_assisted.tolist(), outputs_tti.tolist())

    def test_model_kwarg_assisted_decoding_encoder_decoder(self):
2663
2664
2665
2666
2667
2668
2669
2670
        """
        Tests that the following scenario is compatible with assisted generation:
        1. encoder-decoder main model
        2. encoder-decoder assistant model
        3. both have a custom input
        (e.g. Whisper)
        """

2671
2672
2673
        # PT-only test: TF doesn't support assisted decoding yet.
        # Bart subclass with a kwarg that distorts the output
        class FakeBart(BartForConditionalGeneration):
2674
2675
            def forward(self, input_ids, past_key_values, foo=False, **kwargs):
                outs = super().forward(input_ids, past_key_values=past_key_values, **kwargs)
2676
2677
2678
2679
                if foo:
                    outs["logits"][:, :, :] = 0.0
                return outs

2680
2681
            def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None, **kwargs):
                kwargs["encoder_outputs"] = encoder_outputs
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
                inputs = super().prepare_inputs_for_generation(*args, **kwargs)
                inputs["foo"] = foo
                return inputs

        model = FakeBart.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
            torch_device
        )
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration")

        text = "Hello world"
        tokenized_inputs = tokenizer([text], return_tensors="pt")
        input_ids = tokenized_inputs.input_ids.to(torch_device)

        # Traditional way of generating text
        outputs_normal = model.generate(input_ids)
        self.assertEqual(outputs_normal.shape, (1, 20))

        # Should be different with foo
2700
        outputs_foo = model.generate(input_ids, foo=True)
2701
2702
2703
2704
        with self.assertRaises(AssertionError):
            self.assertListEqual(outputs_foo.tolist(), outputs_normal.tolist())

        # Assistant model
2705
2706
2707
        assistant = FakeBart.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
            torch_device
        )
2708
2709
2710
2711
2712
2713
2714
2715

        # If assisted generation passes model_kwargs correctly, should be same as previous
        outputs_assisted = model.generate(
            input_ids,
            foo=True,
            assistant_model=assistant,
        )
        self.assertListEqual(outputs_assisted.tolist(), outputs_foo.tolist())
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726

        # Check that passing encoder_outputs directly also works as expected
        encoder_outputs = assistant.get_encoder()(input_ids)

        outputs_assisted = model.generate(
            foo=True,
            assistant_model=assistant,
            encoder_outputs=encoder_outputs,
            assistant_encoder_outputs=encoder_outputs,
        )
        self.assertListEqual(outputs_assisted.tolist(), outputs_foo.tolist())
2727
2728

    def test_assisted_decoding_encoder_decoder_shared_encoder(self):
2729
2730
2731
2732
2733
2734
2735
2736
        """
        Tests that the following scenario is compatible with assisted generation:
        1. encoder-decoder main model
        2. decoder-only assistant model
        3. both have a custom input
        (e.g. DistilWhisper)
        """

2737
2738
        # PT-only test: TF doesn't support assisted decoding yet.
        # Bart subclass with a kwarg called foo that distorts the output
2739
        class FakeBartSeq2Seq(BartForConditionalGeneration):
2740
2741
2742
2743
2744
2745
2746
2747
2748
            def forward(self, input_ids, foo=False, **kwargs):
                outs = super().forward(input_ids, **kwargs)
                if foo:
                    outs["logits"][:, :, :] = 0.0
                return outs

            def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None, **kwargs):
                kwargs["encoder_outputs"] = encoder_outputs
                inputs = super().prepare_inputs_for_generation(*args, **kwargs)
2749
2750
2751
2752
2753
2754
2755
2756
2757
                inputs["foo"] = foo
                return inputs

        class FakeBartCausalLM(BartForCausalLM):
            def forward(self, input_ids, attention_mask, past_key_values, foo=False, **kwargs):
                outs = super().forward(input_ids, attention_mask, past_key_values=past_key_values, **kwargs)
                if foo:
                    outs["logits"][:, :, :] = 0.0
                return outs
2758

2759
2760
2761
            def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None, **kwargs):
                kwargs["encoder_outputs"] = encoder_outputs
                inputs = super().prepare_inputs_for_generation(*args, **kwargs)
2762
2763
2764
                inputs["foo"] = foo
                return inputs

2765
        model = FakeBartSeq2Seq.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
            torch_device
        )
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration")

        text = "Hello world"
        tokenized_inputs = tokenizer([text], return_tensors="pt")
        input_ids = tokenized_inputs.input_ids.to(torch_device)

        # Traditional way of generating text
        outputs_normal = model.generate(input_ids)
        self.assertEqual(outputs_normal.shape, (1, 20))

        # Should be different with foo
        outputs_foo = model.generate(input_ids, foo=True)
        with self.assertRaises(AssertionError):
            self.assertListEqual(outputs_foo.tolist(), outputs_normal.tolist())

        # Assistant model
2784
2785
2786
        assistant = FakeBartCausalLM.from_pretrained(
            "hf-internal-testing/tiny-random-BartForConditionalGeneration"
        ).to(torch_device)
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804

        # If assisted generation passes model_kwargs correctly, should be same as previous
        outputs_assisted = model.generate(
            input_ids,
            foo=True,
            assistant_model=assistant,
        )
        self.assertListEqual(outputs_assisted.tolist(), outputs_foo.tolist())

        # Check that passing encoder_outputs directly also works as expected
        encoder_outputs = model.get_encoder()(input_ids)

        outputs_assisted = model.generate(
            foo=True,
            assistant_model=assistant,
            encoder_outputs=encoder_outputs,
        )
        self.assertListEqual(outputs_assisted.tolist(), outputs_foo.tolist())
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850

    def test_assisted_decoding_num_assistant_tokens_heuristic_schedule(self):
        # This test ensures that the assisted generation num_assistant_tokens 'heuristic' schedule works properly.

        prompt = "Alice and Bob"
        checkpoint = "EleutherAI/pythia-160m-deduped"
        tokenizer = AutoTokenizer.from_pretrained(checkpoint)
        inputs = tokenizer(prompt, return_tensors="pt")

        model = AutoModelForCausalLM.from_pretrained(checkpoint)

        assistant_model = model
        assistant_model.generation_config.num_assistant_tokens = 5
        assistant_model.generation_config.num_assistant_tokens_schedule = "heuristic"
        generation_kwargs = {
            "eos_token_id": -1,
            "max_new_tokens": 5,
            "do_sample": False,
            "assistant_model": assistant_model,
        }
        model.generate(**inputs, **generation_kwargs)
        # update_candidate_strategy is called only once and therefore, assistant_model.generation_config.num_assistant_tokens should be either 4 or 7
        self.assertTrue(assistant_model.generation_config.num_assistant_tokens in (4, 7))

    def test_assisted_decoding_num_assistant_tokens_heuristic_transient_schedule(self):
        # This test ensures that the assisted generation num_assistant_tokens 'heuristic' schedule works properly.

        prompt = "Alice and Bob"
        checkpoint = "EleutherAI/pythia-160m-deduped"
        tokenizer = AutoTokenizer.from_pretrained(checkpoint)
        inputs = tokenizer(prompt, return_tensors="pt")

        model = AutoModelForCausalLM.from_pretrained(checkpoint)

        assistant_model = model
        assistant_model.generation_config.num_assistant_tokens = 5
        assistant_model.generation_config.num_assistant_tokens_schedule = "heuristic_transient"
        generation_kwargs = {
            "eos_token_id": -1,
            "max_new_tokens": 5,
            "do_sample": False,
            "assistant_model": assistant_model,
        }
        model.generate(**inputs, **generation_kwargs)
        # update_candidate_strategy is called once but assistant_model.generation_config.num_assistant_tokens should stay 5
        self.assertEqual(assistant_model.generation_config.num_assistant_tokens, 5)
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907

    def test_compare_unprocessed_logit_scores(self):
        # Get unprocessed logit scores back from model generate function.
        # Assert that unprocessed logits from generate() are same as those from modal eval()

        # tell model to generate text and return unprocessed/unwarped logit scores
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        text = "generate yes or no: "
        input_ids = tokenizer([text], return_tensors="pt").input_ids.to(torch_device)

        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)

        with torch.no_grad():
            # Get logits for the next token from fwd pass
            logits_fwd = model(input_ids).logits[:, -1, :][0]

        # Get logits for the next token from generate function
        outputs = model.generate(
            input_ids=input_ids,
            return_dict_in_generate=True,
            output_logits=True,
            max_new_tokens=1,
            do_sample=True,
        )
        logits_gen = outputs.logits[0][0]

        # assert that unprocessed logits from generate() are same as those from modal eval()
        self.assertListEqual(logits_fwd.tolist(), logits_gen.tolist())

    def test_return_unprocessed_logit_scores(self):
        # tell model to generate text and return unprocessed/unwarped logit scores
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        text = "generate yes or no: "
        input_ids = tokenizer([text], return_tensors="pt").input_ids.to(torch_device)
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)

        outputs = model.generate(
            input_ids=input_ids, return_dict_in_generate=True, output_logits=True, max_new_tokens=3
        )

        # perform dummy check if unpreprocessed logits make sense.
        # do preselection on high probabilities; find scores of y and n tokens
        probs_all = torch.nn.functional.softmax(outputs.logits[2][0], dim=-1)
        indices = torch.argwhere(probs_all > 0.001)
        indices = indices[:, -1]
        tokens_max = tokenizer.batch_decode(indices, skip_special_tokens=True)
        probs_max = probs_all[probs_all > 0.001]

        self.assertTrue(len(indices) >= 2)
        next_token_dict = {str(t): p for t, p in zip(tokens_max, probs_max)}
        self.assertTrue("n" in next_token_dict)
        self.assertTrue("y" in next_token_dict)
        y_prob = next_token_dict["y"]
        n_prob = next_token_dict["n"]

        self.assertTrue(y_prob > 0.001 and n_prob > 0.001)
        self.assertTrue(y_prob <= 1.0 and n_prob <= 1.0)
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920

    def test_generate_from_inputs_embeds_with_bos_token_id_is_none(self):
        article = "Today a dragon flew over Paris."
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        input_ids = tokenizer(article, return_tensors="pt").input_ids.to(torch_device)
        inputs_embeds = model.get_input_embeddings()(input_ids)

        model.generate(inputs_embeds=inputs_embeds, max_length=20, bos_token_id=None)

        # bos_token_id is required when no input ids nor inputs_embeds is passed
        with self.assertRaises(ValueError):
            model.generate(max_length=20, bos_token_id=None)