"vscode:/vscode.git/clone" did not exist on "9c9fe89f84f7aa8ec29f19c39a1bf7f1bca82fc3"
test_utils.py 135 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
        GenerationConfig,
69
70
        GreedySearchDecoderOnlyOutput,
        GreedySearchEncoderDecoderOutput,
71
        LogitsProcessorList,
72
        MaxLengthCriteria,
73
        MinLengthLogitsProcessor,
74
75
76
77
78
        PhrasalConstraint,
        SampleDecoderOnlyOutput,
        SampleEncoderDecoderOutput,
        StoppingCriteria,
        StoppingCriteriaList,
79
80
        WatermarkDetector,
        WatermarkingConfig,
81
    )
82
    from transformers.generation.utils import _speculative_sampling
83
84
85
86
87


class GenerationTesterMixin:
    model_tester = None
    all_generative_model_classes = ()
Suraj Patil's avatar
Suraj Patil committed
88
    input_name = "input_ids"
89
    max_new_tokens = 3
90

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

95
        input_ids = input_ids[:batch_size]
96
97
98

        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()`
99
100
101
            if isinstance(config.eos_token_id, int):
                config.eos_token_id = [config.eos_token_id]
            config.pad_token_id = config.eos_token_id[0]
102
        attention_mask = torch.ones_like(input_ids, dtype=torch.long)
103

104
105
106
107
108
        # 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

109
        return config, input_ids, attention_mask
110
111

    @staticmethod
112
    def _get_logits_processor_and_warper_kwargs(
113
114
115
116
        input_length,
        forced_bos_token_id=None,
        forced_eos_token_id=None,
    ):
117
118
119
        process_kwargs = {
            "bad_words_ids": [[1, 0]],
            "repetition_penalty": 1.2,
120
            "remove_invalid_values": True,
121
        }
122
123
124
125
        # 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

126
        warp_kwargs = {"top_k": 10, "top_p": 0.7, "temperature": 0.7}
127
        return process_kwargs, warp_kwargs
128
129

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

139
    @staticmethod
140
    def _get_diverse_beam_kwargs(num_return_sequences=1):
141
142
143
144
145
146
147
148
        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,
        }
149
        return beam_kwargs
150

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

161
    @staticmethod
162
163
164
    def _get_encoder_outputs(
        model, input_ids, attention_mask, output_attentions=None, output_hidden_states=None, num_interleave=1
    ):
165
        encoder = model.get_encoder()
166
167
168
169
170
171
        encoder_outputs = encoder(
            input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
172
173
174
        encoder_outputs["last_hidden_state"] = encoder_outputs.last_hidden_state.repeat_interleave(
            num_interleave, dim=0
        )
175
176
177
        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
178
179
180
        attention_mask = None
        return encoder_outputs, input_ids, attention_mask

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

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

213
        return output_generate
214
215
216
217
218
219
220
221
222
223

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

247
        return output_generate
248
249
250
251
252
253
254
255
256

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

277
        return output_generate
278
279
280
281
282
283
284
285
286

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

308
        return output_generate
309
310
311
312
313
314
315
316
317

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

338
        return output_generate
339

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

370
        return output_generate
371

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

388
        logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
389
390
391
392
393
394
395
396
397
398
            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,
399
            max_new_tokens=self.max_new_tokens,
400
401
402
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_scores=output_scores,
403
            output_logits=output_logits,
404
405
406
407
408
409
            return_dict_in_generate=return_dict_in_generate,
            **logits_process_kwargs,
            **model_kwargs,
            **contrastive_search_kwargs,
        )

410
        return output_generate
411

412
413
    def test_greedy_generate(self):
        for model_class in self.all_generative_model_classes:
414
            config, input_ids, attention_mask = self._get_input_ids_and_config()
415

416
            model = model_class(config).to(torch_device).eval()
417
            output_generate = self._greedy_generate(model=model, input_ids=input_ids, attention_mask=attention_mask)
418

419
420
421
422
            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])
423

424
425
    def test_greedy_generate_dict_outputs(self):
        for model_class in self.all_generative_model_classes:
426
            config, input_ids, attention_mask = self._get_input_ids_and_config()
427

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

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

452
            self._check_outputs(output_generate, input_ids, model.config)
453
454
455

    def test_greedy_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
456
            config, input_ids, attention_mask = self._get_input_ids_and_config()
457
458

            if not hasattr(config, "use_cache"):
459
                self.skipTest("This model doesn't support caching")
460
461

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

475
476
477
478
            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])
479
            self._check_outputs(output_generate, input_ids, model.config, use_cache=True)
480
481
482

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

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

492
            output_generate = self._sample_generate(
493
494
495
496
497
498
499
500
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                num_return_sequences=1,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
            )

501
502
503
504
            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])
505

506
507
    def test_sample_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
508
            config, input_ids, attention_mask = self._get_input_ids_and_config()
509

510
511
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
512

513
            process_kwargs, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(
514
515
516
                input_ids.shape[-1],
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
517
            )
518

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

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

544
            self._check_outputs(output_generate, input_ids, model.config, num_return_sequences=2)
545
546
547

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

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

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

559
            output_generate = self._beam_search_generate(
560
561
                model=model,
                input_ids=input_ids,
562
                attention_mask=attention_mask,
563
564
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
565
            )
566

567
568
569
570
            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])
571
572
573

    def test_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
574
            config, input_ids, attention_mask = self._get_input_ids_and_config()
575
576

            # disable cache
577
            config.use_cache = False
578

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

609
610
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
611
612
613
614
615
            )

    def test_beam_search_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
            # enable cache
616
            config, input_ids, attention_mask = self._get_input_ids_and_config()
617
618

            if not hasattr(config, "use_cache"):
619
                self.skipTest("This model doesn't support caching")
620
621

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

628
            beam_kwargs = self._get_beam_kwargs()
629
630

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

646
647
648
649
            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])
650
651
652
            self._check_outputs(
                output_generate, input_ids, model.config, use_cache=True, num_return_sequences=beam_kwargs["num_beams"]
            )
653

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

661
662
663
            if model_class._no_split_modules is None:
                continue

664
            config, input_ids, attention_mask = self._get_input_ids_and_config()
665
666
667
668
669
670
671
672
673

            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,
674
                    max_new_tokens=self.max_new_tokens,
675
676
677
                    num_beams=2,
                )

678
679
    def test_beam_sample_generate(self):
        for model_class in self.all_generative_model_classes:
680
            config, input_ids, attention_mask = self._get_input_ids_and_config()
681

682
            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])
683

684
            model = model_class(config).to(torch_device).eval()
685
            beam_kwargs = self._get_beam_kwargs()
686

687
            output_generate = self._beam_sample_generate(
688
689
                model=model,
                input_ids=input_ids,
690
                attention_mask=attention_mask,
691
692
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
693
            )
694

695
696
697
698
699
            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])

700
701
702
703
704
705
706
707
708
709
710
711
            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)
712
713
714

    def test_beam_sample_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
715
            config, input_ids, attention_mask = self._get_input_ids_and_config()
716
717

            # disable cache
718
            config.use_cache = False
719

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

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

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

748
749
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
750
            )
751

752
    def test_generate_without_input_ids(self):
753
        config, _, _ = self._get_input_ids_and_config()
754

755
756
757
        # if no bos token id => cannot generate from None
        if config.bos_token_id is None:
            return
758

759
760
761
762
        # 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

763
764
765
        for model_class in self.all_generative_model_classes:
            model = model_class(config).to(torch_device)
            model.eval()
766

767
768
769
            output_ids_generate = model.generate(
                do_sample=False, max_new_tokens=self.max_new_tokens, remove_invalid_values=True
            )
770
            self.assertIsNotNone(output_ids_generate)
771

772
773
    def test_group_beam_search_generate(self):
        for model_class in self.all_generative_model_classes:
774
            config, input_ids, attention_mask = self._get_input_ids_and_config()
775

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

            # check `generate()` and `group_beam_search()` are equal
784
785
            beam_kwargs = self._get_diverse_beam_kwargs()
            output_generate = self._group_beam_search_generate(
786
787
                model=model,
                input_ids=input_ids,
788
                attention_mask=attention_mask,
789
790
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
791
            )
792
793
794
795
            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])
796

797
            # check `group_beam_search` for higher than 1 `num_return_sequences`
798
            num_return_sequences = 2
799
800
            beam_kwargs = self._get_diverse_beam_kwargs(num_return_sequences=num_return_sequences)
            output_generate = self._group_beam_search_generate(
801
802
803
804
805
806
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
807
808
809
810
            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])
811

812
813
    def test_group_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
814
            config, input_ids, attention_mask = self._get_input_ids_and_config()
815
            config.use_cache = False
816

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

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

848
849
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
850
851
            )

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

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

860
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
861
862
863
864
865
866
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )

            # Sample constraints
867
868
            min_id = 3
            max_id = config.vocab_size
869

870
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
871
872
873
874
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

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

            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])

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

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

900
            beam_kwargs = self._get_constrained_beam_kwargs(num_return_sequences=2)
901

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

            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])
915
916
917
918
919
920

            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:
921
            config, input_ids, attention_mask = self._get_input_ids_and_config()
922
923
924
925
926

            # disable cache
            config.use_cache = False

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

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

941
942
            beam_kwargs = self._get_constrained_beam_kwargs()
            output_generate = self._constrained_beam_search_generate(
943
944
945
946
947
948
949
                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,
950
                output_logits=True,
951
952
953
954
955
956
                output_hidden_states=True,
                output_attentions=True,
                return_dict_in_generate=True,
            )

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

967
968
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
969
970
            )

971
972
973
    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).
974
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
975
                self.skipTest("Won't fix: old model with different cache format")
976

977
            config, input_ids, attention_mask = self._get_input_ids_and_config()
978
979
980

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

            # test old generation output for backwards compatibility
            model = model_class(config).to(torch_device).eval()
987
            output_generate = self._contrastive_generate(
988
                model=model, input_ids=input_ids, attention_mask=attention_mask
989
            )
990
991
992
993
            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])
994
995
996
997

    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).
998
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
999
                self.skipTest("Won't fix: old model with different cache format")
1000

1001
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1002
1003
1004

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

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

1021
1022
1023
1024
            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])
1025
            self._check_outputs(output_generate, input_ids, model.config, use_cache=True)
1026

1027
1028
1029
    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:
1030
1031
            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
1032
            if any(model_name in model_class.__name__.lower() for model_name in ["gptbigcode", "jamba"]):
1033
                self.skipTest("TODO: fix me")
1034

1035
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1036
1037
1038

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

            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,
1052
                max_new_tokens=self.max_new_tokens,
1053
1054
1055
1056
1057
1058
1059
1060
                attention_mask=attention_mask,
            )

            high_output = model.generate(
                input_ids,
                top_k=4,
                penalty_alpha=0.6,
                low_memory=False,
1061
                max_new_tokens=self.max_new_tokens,
1062
1063
1064
1065
                attention_mask=attention_mask,
            )
            self.assertListEqual(low_output.tolist(), high_output.tolist())

1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
    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
1080
                    "jamba",
1081
1082
1083
                ]
            ):
                self.skipTest("May fix in the future: need model-specific fixes")
1084
            config, input_ids, _ = self._get_input_ids_and_config(batch_size=2)
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
            # 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())

1100
    @parameterized.expand([("random",), ("same",)])
1101
    @is_flaky()  # Read NOTE (1) below. If there are API issues, all attempts will fail.
1102
    def test_assisted_decoding_matches_greedy_search(self, assistant_type):
1103
        # This test ensures that the assisted generation does not introduce output changes over greedy search.
1104
1105
1106
1107
1108
        # 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:
1109
        # - assisted_decoding, contrarily to the other methods, can't be called on its own (e.g. needs to
1110
        # prepare the assistant encoder outputs in the main generate body);
1111
1112
        # - assisted_decoding does not support `use_cache = False`
        # - assisted_decoding does not support `batch_size > 1`
1113
1114
1115

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

1132
            # enable cache
1133
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1134

1135
1136
1137
            # 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")
1138

1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
            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,
1154
                "output_logits": True,
1155
1156
1157
1158
1159
1160
                "output_hidden_states": True,
                "output_attentions": True,
                "return_dict_in_generate": True,
            }
            output_greedy = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

1161
1162
1163
1164
1165
1166
1167
            # 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
1168
1169
1170
1171
1172
1173
1174
1175
1176
            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)
1177

1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
    @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
1202
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222

            # 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,
1223
                "output_logits": True,
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
                "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)

1239
    def test_assisted_decoding_sample(self):
1240
1241
1242
        # 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).
1243
1244
        for model_class in self.all_generative_model_classes:
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
1245
                self.skipTest("Won't fix: old model with different cache format")
1246
1247
            if any(
                model_name in model_class.__name__.lower()
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
1258
            ):
1259
                self.skipTest("May fix in the future: need model-specific fixes")
1260
1261

            # enable cache
1262
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1263
1264
1265

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

            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
            # 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,
1287
                "output_logits": True,
1288
1289
1290
1291
1292
                "output_hidden_states": True,
                "output_attentions": True,
                "return_dict_in_generate": True,
            }
            output_assisted = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)
1293
1294
1295

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

1296
1297
1298
1299
    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:
1300
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1301
1302
1303
            # We want to test only encoder-decoder models
            if not config.is_encoder_decoder:
                continue
Joao Gante's avatar
Joao Gante committed
1304
            model = model_class(config).to(torch_device)
1305
1306

            head_masking = {
1307
1308
1309
1310
1311
1312
1313
                "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
                ),
1314
1315
1316
1317
            }

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

            for attn_name, (name, mask) in zip(attention_names, head_masking.items()):
                out = model.generate(
                    input_ids,
1324
                    attention_mask=attention_mask,
1325
1326
1327
                    num_beams=1,
                    output_attentions=True,
                    return_dict_in_generate=True,
1328
                    remove_invalid_values=True,
1329
1330
1331
1332
1333
1334
                    **{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)

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

1339
1340
1341
1342
1343
1344
1345
        # 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 = []
1346
        for model_class in self.all_generative_model_classes:
1347
            config, _, _ = self._get_input_ids_and_config()
1348
            if config.is_encoder_decoder:
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
1377
1378
1379
                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:
1380
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1381
1382
1383
            model = model_class(config).to(torch_device).eval()
            signature = inspect.signature(model.forward).parameters.keys()

1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
            # 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))
1398

1399
1400
1401
1402
1403
1404
1405
1406
    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"):
1407
                self.skipTest("This model doesn't support caching")
1408
1409
1410
1411
1412
1413
1414
1415

            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:
1416
                self.skipTest("This model doesn't return `past_key_values`")
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
1467
1468
1469

            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)
                    )

1470
1471
1472
1473
    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:
1474
            config, input_ids, _ = self._get_input_ids_and_config()
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
1514
1515
1516

            # 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(),
1517
                outputs_from_embeds_wo_ids.tolist(),
1518
1519
            )

1520
1521
1522
1523
    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"]):
1524
                self.skipTest("Won't fix: old model with unique inputs/caches/other")
1525
            if any(model_name in model_class.__name__.lower() for model_name in ["umt5"]):
1526
                self.skipTest("TODO: needs modeling or test input preparation fixes for compatibility")
1527
1528
1529
1530

            config, inputs = self.model_tester.prepare_config_and_inputs_for_common()

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

            # 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

1550
            # If "past_key_values" is not returned, skip the test (e.g. RWKV uses a different cache name and format)
1551
1552
            outputs = model(**inputs)
            if "past_key_values" not in outputs:
1553
                self.skipTest("This model doesn't return `past_key_values`")
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
1593
1594
1595

            # 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],
                        )
                    )

1596
1597
1598
1599
1600
1601
1602
1603
1604
    @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")

1605
            config, input_ids, attention_mask = self._get_input_ids_and_config()
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
1653
1654
1655
            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],
                        )
                    )

1656
1657
1658
    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
1659

1660
1661
1662
1663
1664
1665
1666
        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)

1667
1668
1669
        # unprocessed logits
        self._check_logits(num_sequences_in_output, output.logits, config=config)

1670
1671
1672
        # Attentions
        if config.is_encoder_decoder:
            # encoder
1673
            self._check_encoder_attention_for_generate(output.encoder_attentions, batch_size, config, seq_length)
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
            # 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
1699
1700
            self._check_encoder_hidden_states_for_generate(
                output.encoder_hidden_states, batch_size, config, seq_length
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
            )

            # 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
1725
        # Past Key Value States -- a few notes here:
1726
1727
        # 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
1728
1729
1730
        # 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")
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
        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,
            )

1744
1745
1746
1747
1748
1749
    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))

1750
1751
1752
1753
1754
1755
1756
1757
    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))

1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
    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)
            )

1782
1783
1784
1785
1786
1787
1788
1789
    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),
        )

1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
    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),
            )
1808

1809
1810
1811
1812
1813
1814
1815
1816
    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),
        )

1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
    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),
        )

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

1845
1846
1847
1848
1849
1850
        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)
1851
1852
1853
1854
        longer = tensor_2 if in_order else tensor_1
        shorter = tensor_1 if in_order else tensor_2

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

        self.assertTrue(flag)

1864
1865
1866

@require_torch
class UtilsFunctionsTest(unittest.TestCase):
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
1899
1900
1901
    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])

1902
1903

@require_torch
1904
1905
1906
1907
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 = {
1908
            "AutoModelForCausalLM": AutoModelForCausalLM,
1909
            "AutoModelForSpeechSeq2Seq": AutoModelForSpeechSeq2Seq,
1910
            "AutoModelForSeq2SeqLM": AutoModelForSeq2SeqLM,
1911
            "AutoModelForVision2Seq": AutoModelForVision2Seq,
1912
1913
            "LogitsProcessorList": LogitsProcessorList,
            "MinLengthLogitsProcessor": MinLengthLogitsProcessor,
1914
            "create_tensor_fn": torch.tensor,
1915
            "floats_tensor": floats_tensor,
1916
1917
1918
            "return_tensors": "pt",
        }

1919
1920
    @slow
    def test_diverse_beam_search(self):
1921
        # PT-only test: TF doesn't have a diverse beam search implementation
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
        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(
1932
1933
1934
1935
1936
1937
            input_ids,
            num_beams=4,
            num_return_sequences=2,
            num_beam_groups=4,
            diversity_penalty=2.0,
            remove_invalid_values=True,
1938
1939
1940
1941
1942
1943
1944
        )

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

        self.assertListEqual(
            generated_text,
            [
Sylvain Gugger's avatar
Sylvain Gugger committed
1945
1946
1947
1948
1949
1950
                "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.",
1951
1952
            ],
        )
1953

1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
    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])

1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
    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])

1982
    def test_custom_stopping_criteria_overload_error(self):
1983
        # PT-only test: TF doesn't have StoppingCriteria
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
        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):
1997
        # PT-only test: TF doesn't have StoppingCriteria
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
        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],
        )

2019
    def test_stop_sequence_stopping_criteria(self):
2020
        # PT-only test: TF doesn't have StoppingCriteria
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
        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"}])

2038
    def test_generate_non_nlp_input_ids_as_kwarg(self):
2039
        # PT-only test: AFAIK there's no non-NLP model architecture in TF that supports `input_ids` as its only input
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
        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))

2051
    def test_generate_input_values_as_encoder_kwarg(self):
2052
        # PT-only test: AFAIK there's no generate-capable architecture in TF that supports `input_values` as its input
2053
2054
2055
2056
2057
2058
2059
2060
2061
        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))

2062
    def test_transition_scores_group_beam_search_encoder_decoder(self):
2063
        # PT-only test: TF doesn't have group beam search
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
        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,
2075
            diversity_penalty=1.0,
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
            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)

2086
        transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, outputs.beam_indices)
2087
2088
2089
        transition_scores_sum = transition_scores.sum(-1)

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

2091
    def test_beam_search_low_memory(self):
2092
2093
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
        model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
        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())

2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
    @slow
    def test_watermark_generation(self):
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
        model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer.pad_token_id = tokenizer.eos_token_id
        model_inputs = tokenizer("I will be", return_tensors="pt").to(torch_device)
        input_len = model_inputs["input_ids"].shape[-1]

        # generation should work with both input types: WatermarkingConfig or Dict, so let's check it here :)
        watermark_config = WatermarkingConfig(bias=2.5, seeding_scheme="selfhash")
        _ = model.generate(**model_inputs, watermarking_config=watermark_config, do_sample=False, max_length=15)

        args = {
            "bias": 2.0,
            "context_width": 1,
            "seeding_scheme": "selfhash",
            "greenlist_ratio": 0.25,
            "hashing_key": 15485863,
        }
        output = model.generate(**model_inputs, do_sample=False, max_length=15)
        output_selfhash = model.generate(**model_inputs, watermarking_config=args, do_sample=False, max_length=15)

        # check that the watermarked text is generating what is should
        self.assertListEqual(
            output.tolist(), [[40, 481, 307, 262, 717, 284, 9159, 326, 314, 716, 407, 257, 4336, 286, 262]]
        )
        self.assertListEqual(
            output_selfhash.tolist(), [[40, 481, 307, 2263, 616, 640, 284, 651, 616, 1621, 503, 612, 553, 531, 367]]
        )

        detector = WatermarkDetector(model_config=model.config, device=torch_device, watermarking_config=args)
        detection_out_watermarked = detector(output_selfhash[:, input_len:], return_dict=True)
        detection_out = detector(output[:, input_len:], return_dict=True)

        # check that the detector is detecting watermarked text
        self.assertListEqual(detection_out_watermarked.prediction.tolist(), [True])
        self.assertListEqual(detection_out.prediction.tolist(), [False])

2142
2143
    @slow
    def test_beam_search_example_integration(self):
2144
        # PT-only test: TF doesn't have a BeamSearchScorer
2145
2146
        # exactly the example provided in the docstrings of beam search, which previously
        # failed after directly copying from it. Refer to PR #15555
2147
2148
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2149
2150
2151
2152
2153
2154
2155

        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
2156
        input_ids = torch.ones((1, 1), device=model.device, dtype=torch.long)
2157
2158
2159
        input_ids = input_ids * model.config.decoder_start_token_id

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

2162
2163
        outputs = model.generate(
            input_ids, num_beams=num_beams, min_length=5, eos_token_id=model.config.eos_token_id, **model_kwargs
2164
2165
2166
2167
2168
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

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

2169
2170
    @slow
    def test_constrained_beam_search(self):
2171
        # PT-only test: TF doesn't have constrained beam search
2172
2173
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2174

2175
2176
        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
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201

        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,
            [
2202
2203
                "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"
2204
2205
2206
            ],
        )

2207
2208
    @slow
    def test_constrained_beam_search_mixed(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
2239
2240
2241

        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,
            [
2242
2243
2244
                "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",
2245
2246
2247
2248
2249
            ],
        )

    @slow
    def test_constrained_beam_search_mixed_mixin(self):
2250
        # PT-only test: TF doesn't have constrained beam search
2251
2252
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
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

        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,
            [
2280
2281
2282
                "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",
2283
2284
2285
            ],
        )

2286
2287
    @slow
    def test_cfg_mixin(self):
2288
2289
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325

        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'
            ],
        )

2326
2327
    @slow
    def test_constrained_beam_search_example_translation_mixin(self):
2328
        # PT-only test: TF doesn't have constrained beam search
2329
2330
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348

        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)

2349
        self.assertListEqual(outputs, ["Wie alt sind Sie?"])
2350

2351
2352
    @slow
    def test_constrained_beam_search_example_integration(self):
2353
        # PT-only test: TF doesn't have constrained beam search
2354
2355
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2356
2357
2358
2359
2360
2361
2362

        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
2363
        input_ids = torch.ones((1, 1), device=model.device, dtype=torch.long)
2364
2365
2366
        input_ids = input_ids * model.config.decoder_start_token_id

        # add encoder_outputs to model keyword arguments
2367
        model_kwargs = {"encoder_outputs": model.get_encoder()(encoder_input_ids, return_dict=True)}
2368
2369
2370
2371

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

2372
2373
2374
2375
2376
2377
2378
        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,
2379
2380
2381
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

2382
        self.assertListEqual(outputs, ["Wie alt sind Sie?"])
2383

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
2420
    @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)

2421
    def test_constrained_beam_search_mixin_type_checks(self):
2422
        # PT-only test: TF doesn't have constrained beam search
2423
2424
        tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/t5-tiny-random")
        model = AutoModelForSeq2SeqLM.from_pretrained("patrickvonplaten/t5-tiny-random")
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460

        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]]])
2461

2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
    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())

2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
    def test_decoder_start_id_from_config(self):
        # Refer to: (#30899)
        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

        # we should be able to take `decoder_start_token_id` from model's generation config if user passes a `GenerationConfig` type
        outputs = bart_model.generate(input_ids, generation_config=GenerationConfig(do_sample=False))

        # If the generatoin config has no `decoder_start_token_id` or `bos_token_id`, we will raise an error unless user passes it in config
        bart_model.generation_config.decoder_start_token_id = None
        bart_model.generation_config.bos_token_id = None
        outputs_with_user_id = bart_model.generate(
            input_ids,
            generation_config=GenerationConfig(do_sample=False, decoder_start_token_id=decoder_start_token_id),
        )

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

        with self.assertRaises(ValueError):
            outputs = bart_model.generate(input_ids, generation_config=GenerationConfig(do_sample=False))

2511
    def test_contrastive_search_batched(self):
2512
        # PT-only test: TF doesn't have constrained beam search
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
        # 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)

2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
    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())

2558
    def test_eos_token_id_int_and_list_top_k_top_sampling(self):
2559
        # Has TF equivalent: this test relies on random sampling
2560
2561
2562
2563
2564
2565
2566
        generation_kwargs = {
            "do_sample": True,
            "num_beams": 1,
            "top_p": 0.7,
            "top_k": 10,
            "temperature": 0.7,
        }
2567
        expectation = 20
2568

2569
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
2570
        text = """Hello, my dog is cute and"""
2571
        tokens = tokenizer(text, return_tensors="pt").to(torch_device)
2572
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
2573

2574
2575
2576
        # 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)
2577
2578
2579
2580
        eos_token_id = 846
        generated_tokens = model.generate(**tokens, eos_token_id=eos_token_id, **generation_kwargs)
        self.assertTrue(expectation == len(generated_tokens[0]))

2581
        torch.manual_seed(1)
2582
        eos_token_id = [846, 198]
2583
2584
        generated_tokens = model.generate(**tokens, eos_token_id=eos_token_id, **generation_kwargs)
        self.assertTrue(expectation == len(generated_tokens[0]))
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
    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")
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642

    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:
2643
            # generation_config is modified -> legacy mode is disabled = generation_config takes precedence
2644
2645
2646
            model.generation_config.max_length = 10
            model.generate(input_ids)
            self.assertEqual(len(warning_list), 0)
2647

2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
    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)

2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
    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):
2733
2734
2735
2736
2737
2738
2739
2740
        """
        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)
        """

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

2750
2751
            def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None, **kwargs):
                kwargs["encoder_outputs"] = encoder_outputs
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
                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
2770
        outputs_foo = model.generate(input_ids, foo=True)
2771
2772
2773
2774
        with self.assertRaises(AssertionError):
            self.assertListEqual(outputs_foo.tolist(), outputs_normal.tolist())

        # Assistant model
2775
2776
2777
        assistant = FakeBart.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
            torch_device
        )
2778
2779
2780
2781
2782
2783
2784
2785

        # 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())
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796

        # 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())
2797
2798

    def test_assisted_decoding_encoder_decoder_shared_encoder(self):
2799
2800
2801
2802
2803
2804
2805
2806
        """
        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)
        """

2807
2808
        # PT-only test: TF doesn't support assisted decoding yet.
        # Bart subclass with a kwarg called foo that distorts the output
2809
        class FakeBartSeq2Seq(BartForConditionalGeneration):
2810
2811
2812
2813
2814
2815
2816
2817
2818
            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)
2819
2820
2821
2822
2823
2824
2825
2826
2827
                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
2828

2829
2830
2831
            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)
2832
2833
2834
                inputs["foo"] = foo
                return inputs

2835
        model = FakeBartSeq2Seq.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
            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
2854
2855
2856
        assistant = FakeBartCausalLM.from_pretrained(
            "hf-internal-testing/tiny-random-BartForConditionalGeneration"
        ).to(torch_device)
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874

        # 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())
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
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920

    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)
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977

    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)
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990

    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)