test_utils.py 145 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
    require_accelerate,
Ahmed Moubtahij's avatar
Ahmed Moubtahij committed
30
    require_auto_gptq,
31
    require_quanto,
32
    require_torch,
jiqing-feng's avatar
jiqing-feng committed
33
    require_torch_gpu,
34
    require_torch_multi_accelerator,
jiqing-feng's avatar
jiqing-feng committed
35
    require_torch_multi_gpu,
36
37
38
    slow,
    torch_device,
)
39

40
from ..test_modeling_common import floats_tensor, ids_tensor
41
from .test_framework_agnostic import GenerationIntegrationTestsMixin
42

43
44
45
46

if is_torch_available():
    import torch

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


class GenerationTesterMixin:
    model_tester = None
    all_generative_model_classes = ()
Suraj Patil's avatar
Suraj Patil committed
93
    input_name = "input_ids"
94
    max_new_tokens = 3
95

96
    def _get_input_ids_and_config(self, batch_size=2):
97
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
Suraj Patil's avatar
Suraj Patil committed
98
        input_ids = inputs_dict[self.input_name]
99

100
        input_ids = input_ids[:batch_size]
101
102
103

        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()`
104
105
106
            if isinstance(config.eos_token_id, int):
                config.eos_token_id = [config.eos_token_id]
            config.pad_token_id = config.eos_token_id[0]
107
108
109
110
111

        if self.has_attentions:
            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        else:
            attention_mask = None
112

113
114
115
116
117
        # 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

118
        return config, input_ids, attention_mask
119
120

    @staticmethod
121
    def _get_logits_processor_and_warper_kwargs(
122
123
124
125
        input_length,
        forced_bos_token_id=None,
        forced_eos_token_id=None,
    ):
126
127
128
        process_kwargs = {
            "bad_words_ids": [[1, 0]],
            "repetition_penalty": 1.2,
129
            "remove_invalid_values": True,
130
        }
131
132
133
134
        # 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

135
        warp_kwargs = {"top_k": 10, "top_p": 0.7, "temperature": 0.7}
136
        return process_kwargs, warp_kwargs
137
138

    @staticmethod
139
    def _get_beam_kwargs(num_return_sequences=1):
140
141
142
143
144
145
        beam_kwargs = {
            "early_stopping": False,
            "length_penalty": 2.0,
            "num_beams": 2,
            "num_return_sequences": num_return_sequences,
        }
146
        return beam_kwargs
147

148
    @staticmethod
149
    def _get_diverse_beam_kwargs(num_return_sequences=1):
150
151
152
153
154
155
156
157
        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,
        }
158
        return beam_kwargs
159

160
    @staticmethod
161
    def _get_constrained_beam_kwargs(num_return_sequences=1):
162
163
164
165
166
167
        beam_kwargs = {
            "early_stopping": False,
            "length_penalty": 2.0,
            "num_beams": num_return_sequences * 4,
            "num_return_sequences": num_return_sequences,
        }
168
        return beam_kwargs
169

170
    @staticmethod
171
172
173
    def _get_encoder_outputs(
        model, input_ids, attention_mask, output_attentions=None, output_hidden_states=None, num_interleave=1
    ):
174
        encoder = model.get_encoder()
175
176
177
178
179
180
        encoder_outputs = encoder(
            input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
181
182
183
        encoder_outputs["last_hidden_state"] = encoder_outputs.last_hidden_state.repeat_interleave(
            num_interleave, dim=0
        )
184
185
186
        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
187
188
189
        attention_mask = None
        return encoder_outputs, input_ids, attention_mask

190
191
192
193
194
195
    def _greedy_generate(
        self,
        model,
        input_ids,
        attention_mask,
        output_scores=False,
196
        output_logits=False,
197
198
199
200
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
201
        logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
202
203
204
            input_ids.shape[-1],
            forced_bos_token_id=model.config.forced_bos_token_id,
            forced_eos_token_id=model.config.forced_eos_token_id,
205
206
        )

207
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
208
209
210
211
        output_generate = model.generate(
            input_ids,
            do_sample=False,
            num_beams=1,
212
            max_new_tokens=self.max_new_tokens,
213
214
215
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_scores=output_scores,
216
            output_logits=output_logits,
217
218
            return_dict_in_generate=return_dict_in_generate,
            **logits_process_kwargs,
219
            **model_kwargs,
220
221
        )

222
        return output_generate
223
224
225
226
227
228
229
230
231
232

    def _sample_generate(
        self,
        model,
        input_ids,
        attention_mask,
        num_return_sequences,
        logits_warper_kwargs,
        process_kwargs,
        output_scores=False,
233
        output_logits=False,
234
235
236
237
238
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        torch.manual_seed(0)
239
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
240
241
242
243
        output_generate = model.generate(
            input_ids,
            do_sample=True,
            num_beams=1,
244
            max_new_tokens=self.max_new_tokens,
245
246
            num_return_sequences=num_return_sequences,
            output_scores=output_scores,
247
            output_logits=output_logits,
248
249
250
251
252
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **logits_warper_kwargs,
            **process_kwargs,
253
            **model_kwargs,
254
255
        )

256
        return output_generate
257
258
259
260
261
262
263
264
265

    def _beam_search_generate(
        self,
        model,
        input_ids,
        attention_mask,
        beam_kwargs,
        logits_process_kwargs,
        output_scores=False,
266
        output_logits=False,
267
268
269
270
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
271
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
272
273
274
        output_generate = model.generate(
            input_ids,
            do_sample=False,
275
            max_new_tokens=self.max_new_tokens,
276
            output_scores=output_scores,
277
            output_logits=output_logits,
278
279
280
281
282
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **beam_kwargs,
            **logits_process_kwargs,
283
            **model_kwargs,
284
285
        )

286
        return output_generate
287
288
289
290
291
292
293
294
295

    def _beam_sample_generate(
        self,
        model,
        input_ids,
        attention_mask,
        beam_kwargs,
        logits_warper_kwargs,
        output_scores=False,
296
        output_logits=False,
297
298
299
300
301
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        torch.manual_seed(0)
302
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
303
304
305
        output_generate = model.generate(
            input_ids,
            do_sample=True,
306
            max_new_tokens=self.max_new_tokens,
307
            output_scores=output_scores,
308
            output_logits=output_logits,
309
310
311
312
313
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **beam_kwargs,
            **logits_warper_kwargs,
314
            **model_kwargs,
315
316
        )

317
        return output_generate
318
319
320
321
322
323
324
325
326

    def _group_beam_search_generate(
        self,
        model,
        input_ids,
        attention_mask,
        beam_kwargs,
        logits_process_kwargs,
        output_scores=False,
327
        output_logits=False,
328
329
330
331
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
332
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
333
334
335
        output_generate = model.generate(
            input_ids,
            do_sample=False,
336
            max_new_tokens=self.max_new_tokens,
337
            output_scores=output_scores,
338
            output_logits=output_logits,
339
340
341
342
343
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=return_dict_in_generate,
            **beam_kwargs,
            **logits_process_kwargs,
344
            **model_kwargs,
345
346
        )

347
        return output_generate
348

349
350
351
352
353
354
355
356
357
    def _constrained_beam_search_generate(
        self,
        model,
        input_ids,
        attention_mask,
        constraints,
        beam_kwargs,
        logits_process_kwargs,
        output_scores=False,
358
        output_logits=False,
359
360
361
362
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
363
        model_kwargs = {"attention_mask": attention_mask} if attention_mask is not None else {}
364
365
366
        output_generate = model.generate(
            input_ids,
            do_sample=False,
367
            max_new_tokens=self.max_new_tokens,
368
            output_scores=output_scores,
369
            output_logits=output_logits,
370
371
372
373
374
375
            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,
376
            **model_kwargs,
377
378
        )

379
        return output_generate
380

381
382
383
384
385
386
    def _contrastive_generate(
        self,
        model,
        input_ids,
        attention_mask,
        output_scores=False,
387
        output_logits=False,
388
389
390
391
392
393
394
395
396
        output_attentions=False,
        output_hidden_states=False,
        return_dict_in_generate=False,
    ):
        contrastive_search_kwargs = {
            "penalty_alpha": 0.6,
            "top_k": 5,
        }

397
        logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
398
399
400
401
402
403
404
405
406
407
            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,
408
            max_new_tokens=self.max_new_tokens,
409
410
411
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_scores=output_scores,
412
            output_logits=output_logits,
413
414
415
416
417
418
            return_dict_in_generate=return_dict_in_generate,
            **logits_process_kwargs,
            **model_kwargs,
            **contrastive_search_kwargs,
        )

419
        return output_generate
420

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

425
            model = model_class(config).to(torch_device).eval()
426
            output_generate = self._greedy_generate(model=model, input_ids=input_ids, attention_mask=attention_mask)
427

428
429
430
431
            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])
432

433
434
    def test_greedy_generate_dict_outputs(self):
        for model_class in self.all_generative_model_classes:
435
            config, input_ids, attention_mask = self._get_input_ids_and_config()
436

437
438
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
439
            output_generate = self._greedy_generate(
440
441
442
443
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_scores=True,
444
                output_logits=True,
445
                output_hidden_states=True,
446
                output_attentions=self.has_attentions,
447
448
                return_dict_in_generate=True,
            )
449
450

            if model.config.is_encoder_decoder:
451
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
452
453
                self.assertIsInstance(output_generate, GenerateEncoderDecoderOutput)
                # Retrocompatibility check
454
455
                self.assertIsInstance(output_generate, GreedySearchEncoderDecoderOutput)
            else:
456
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
457
458
                self.assertIsInstance(output_generate, GenerateDecoderOnlyOutput)
                # Retrocompatibility check
459
                self.assertIsInstance(output_generate, GreedySearchDecoderOnlyOutput)
460

461
            self._check_outputs(output_generate, input_ids, model.config)
462
463
464

    def test_greedy_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
465
            config, input_ids, attention_mask = self._get_input_ids_and_config()
466
467

            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
468
                self.skipTest(reason="This model doesn't support caching")
469
            if any(model_name in model_class.__name__.lower() for model_name in ["rwkv"]):
amyeroberts's avatar
amyeroberts committed
470
                self.skipTest(reason="Won't fix: model with non-standard dictionary output shapes")
471
472

            config.use_cache = True
473
            config.is_decoder = True
474
            model = model_class(config).to(torch_device).eval()
475
            output_generate = self._greedy_generate(
476
477
                model=model,
                input_ids=input_ids,
478
                attention_mask=attention_mask,
479
                output_scores=True,
480
                output_logits=True,
481
                output_hidden_states=True,
482
                output_attentions=self.has_attentions,
483
                return_dict_in_generate=True,
484
            )
485

486
487
488
489
            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])
490
            self._check_outputs(output_generate, input_ids, model.config, use_cache=True)
491
492
493

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

496
497
            model = model_class(config).to(torch_device).eval()
            process_kwargs, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(
498
499
500
501
502
                input_ids.shape[-1],
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
            )

503
            output_generate = self._sample_generate(
504
505
506
507
508
509
510
511
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                num_return_sequences=1,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
            )

512
513
514
515
            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])
516

517
518
    def test_sample_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
519
            config, input_ids, attention_mask = self._get_input_ids_and_config()
520

521
522
            config.use_cache = False
            model = model_class(config).to(torch_device).eval()
523

524
            process_kwargs, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(
525
526
527
                input_ids.shape[-1],
                forced_bos_token_id=model.config.forced_bos_token_id,
                forced_eos_token_id=model.config.forced_eos_token_id,
528
            )
529

530
            output_generate = self._sample_generate(
531
532
                model=model,
                input_ids=input_ids,
533
                attention_mask=attention_mask,
534
535
536
537
                num_return_sequences=2,
                logits_warper_kwargs=logits_warper_kwargs,
                process_kwargs=process_kwargs,
                output_scores=True,
538
                output_logits=True,
539
                output_hidden_states=True,
540
                output_attentions=self.has_attentions,
541
                return_dict_in_generate=True,
542
543
544
            )

            if model.config.is_encoder_decoder:
545
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
546
547
                self.assertIsInstance(output_generate, GenerateEncoderDecoderOutput)
                # Retrocompatibility check
548
                self.assertIsInstance(output_generate, SampleEncoderDecoderOutput)
549
            else:
550
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
551
552
                self.assertIsInstance(output_generate, GenerateDecoderOnlyOutput)
                # Retrocompatibility check
553
554
                self.assertIsInstance(output_generate, SampleDecoderOnlyOutput)

555
            self._check_outputs(output_generate, input_ids, model.config, num_return_sequences=2)
556
557
558

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

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

563
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
564
565
566
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
567
            )
568
            beam_kwargs = self._get_beam_kwargs()
569

570
            output_generate = self._beam_search_generate(
571
572
                model=model,
                input_ids=input_ids,
573
                attention_mask=attention_mask,
574
575
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
576
            )
577

578
579
580
581
            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])
582
583
584

    def test_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
585
            config, input_ids, attention_mask = self._get_input_ids_and_config()
586
587

            # disable cache
588
            config.use_cache = False
589

590
            model = model_class(config).to(torch_device).eval()
591
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
592
593
594
595
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )
596
597
            beam_kwargs = self._get_beam_kwargs()
            output_generate = self._beam_search_generate(
598
599
                model=model,
                input_ids=input_ids,
600
                attention_mask=attention_mask,
601
602
603
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
604
                output_logits=True,
605
                output_hidden_states=True,
606
                output_attentions=self.has_attentions,
607
                return_dict_in_generate=True,
608
609
            )
            if model.config.is_encoder_decoder:
610
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
611
612
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
613
                self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput)
614
            else:
615
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
616
617
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
618
619
                self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput)

620
621
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
622
623
624
625
626
            )

    def test_beam_search_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
            # enable cache
627
            config, input_ids, attention_mask = self._get_input_ids_and_config()
628
629

            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
630
                self.skipTest(reason="This model doesn't support caching")
631
            if any(model_name in model_class.__name__.lower() for model_name in ["rwkv"]):
amyeroberts's avatar
amyeroberts committed
632
                self.skipTest(reason="Won't fix: model with non-standard dictionary output shapes")
633
634

            model = model_class(config).to(torch_device).eval()
635
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
636
637
638
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
639
640
            )

641
            beam_kwargs = self._get_beam_kwargs()
642
643

            config.use_cache = True
644
            config.is_decoder = True
645
            model = model_class(config).to(torch_device).eval()
646
            output_generate = self._beam_search_generate(
647
648
649
650
651
652
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
653
                output_logits=True,
654
                output_hidden_states=True,
655
                output_attentions=self.has_attentions,
656
657
658
                return_dict_in_generate=True,
            )

659
660
661
662
            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])
663
664
665
            self._check_outputs(
                output_generate, input_ids, model.config, use_cache=True, num_return_sequences=beam_kwargs["num_beams"]
            )
666

667
    @require_accelerate
668
    @require_torch_multi_accelerator
669
670
    def test_model_parallel_beam_search(self):
        for model_class in self.all_generative_model_classes:
671
            if "xpu" in torch_device:
amyeroberts's avatar
amyeroberts committed
672
                return unittest.skip(reason="device_map='auto' does not work with XPU devices")
673

674
675
676
            if model_class._no_split_modules is None:
                continue

677
            config, input_ids, attention_mask = self._get_input_ids_and_config()
678
679
680
681
682
683
684
685
686

            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,
687
                    max_new_tokens=self.max_new_tokens,
688
689
690
                    num_beams=2,
                )

691
692
    def test_beam_sample_generate(self):
        for model_class in self.all_generative_model_classes:
693
            config, input_ids, attention_mask = self._get_input_ids_and_config()
694

695
            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])
696

697
            model = model_class(config).to(torch_device).eval()
698
            beam_kwargs = self._get_beam_kwargs()
699

700
            output_generate = self._beam_sample_generate(
701
702
                model=model,
                input_ids=input_ids,
703
                attention_mask=attention_mask,
704
705
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
706
            )
707

708
709
710
711
712
            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])

713
714
715
716
717
718
719
720
721
722
723
724
            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)
725
726
727

    def test_beam_sample_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
728
            config, input_ids, attention_mask = self._get_input_ids_and_config()
729
730

            # disable cache
731
            config.use_cache = False
732

733
            model = model_class(config).to(torch_device).eval()
734
735
            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])
            beam_kwargs = self._get_beam_kwargs()
736

737
            output_generate = self._beam_sample_generate(
738
739
740
741
742
743
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
                output_scores=True,
744
                output_logits=True,
745
                output_hidden_states=True,
746
                output_attentions=self.has_attentions,
747
748
749
750
                return_dict_in_generate=True,
            )

            if model.config.is_encoder_decoder:
751
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
752
753
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
754
                self.assertIsInstance(output_generate, BeamSampleEncoderDecoderOutput)
755
            else:
756
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
757
758
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
759
                self.assertIsInstance(output_generate, BeamSampleDecoderOnlyOutput)
760

761
762
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
763
            )
764

765
    def test_generate_without_input_ids(self):
766
        config, _, _ = self._get_input_ids_and_config()
767

768
769
        # if no bos token id => cannot generate from None
        if config.bos_token_id is None:
amyeroberts's avatar
amyeroberts committed
770
            self.skipTest(reason="bos_token_id is None")
771

772
773
774
775
        # 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

776
777
778
        for model_class in self.all_generative_model_classes:
            model = model_class(config).to(torch_device)
            model.eval()
779

780
781
782
            output_ids_generate = model.generate(
                do_sample=False, max_new_tokens=self.max_new_tokens, remove_invalid_values=True
            )
783
            self.assertIsNotNone(output_ids_generate)
784

785
786
    def test_group_beam_search_generate(self):
        for model_class in self.all_generative_model_classes:
787
            config, input_ids, attention_mask = self._get_input_ids_and_config()
788

789
            model = model_class(config).to(torch_device).eval()
790
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
791
792
793
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
794
795
796
            )

            # check `generate()` and `group_beam_search()` are equal
797
798
            beam_kwargs = self._get_diverse_beam_kwargs()
            output_generate = self._group_beam_search_generate(
799
800
                model=model,
                input_ids=input_ids,
801
                attention_mask=attention_mask,
802
803
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
804
            )
805
806
807
808
            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])
809

810
            # check `group_beam_search` for higher than 1 `num_return_sequences`
811
            num_return_sequences = 2
812
813
            beam_kwargs = self._get_diverse_beam_kwargs(num_return_sequences=num_return_sequences)
            output_generate = self._group_beam_search_generate(
814
815
816
817
818
819
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
820
821
822
823
            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])
824

825
826
    def test_group_beam_search_generate_dict_output(self):
        for model_class in self.all_generative_model_classes:
827
            config, input_ids, attention_mask = self._get_input_ids_and_config()
828
            config.use_cache = False
829

830
            model = model_class(config).to(torch_device).eval()
831
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
832
833
834
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
835
836
            )

837
838
            beam_kwargs = self._get_diverse_beam_kwargs()
            output_generate = self._group_beam_search_generate(
839
840
                model=model,
                input_ids=input_ids,
841
                attention_mask=attention_mask,
842
843
844
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
                output_scores=True,
845
                output_logits=True,
846
                output_hidden_states=True,
847
                output_attentions=self.has_attentions,
848
                return_dict_in_generate=True,
849
850
            )
            if model.config.is_encoder_decoder:
851
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
852
853
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
854
                self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput)
855
            else:
856
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
857
858
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
859
860
                self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput)

861
862
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
863
864
            )

865
866
    # TODO: @gante
    @is_flaky()
867
868
    def test_constrained_beam_search_generate(self):
        for model_class in self.all_generative_model_classes:
869
            config, input_ids, attention_mask = self._get_input_ids_and_config()
870
871
872

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

873
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
874
875
876
877
878
879
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )

            # Sample constraints
880
881
            min_id = 3
            max_id = config.vocab_size
882

883
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
884
885
886
887
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

888
889
            beam_kwargs = self._get_constrained_beam_kwargs()
            output_generate = self._constrained_beam_search_generate(
890
891
892
893
894
895
896
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                constraints=constraints,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
897
898
899
900
901
902

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

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

906
            # check`constrained_beam_search` for higher than 1 `num_return_sequences`
907
            # Sample constraints
908
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
909
910
911
912
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

913
            beam_kwargs = self._get_constrained_beam_kwargs(num_return_sequences=2)
914

915
            output_generate = self._constrained_beam_search_generate(
916
917
918
919
920
921
922
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                constraints=constraints,
                beam_kwargs=beam_kwargs,
                logits_process_kwargs=logits_process_kwargs,
            )
923
924
925
926
927

            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])
928
929
930
931
932
933

            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:
934
            config, input_ids, attention_mask = self._get_input_ids_and_config()
935
936
937
938
939

            # disable cache
            config.use_cache = False

            model = model_class(config).to(torch_device).eval()
940
            logits_process_kwargs, _ = self._get_logits_processor_and_warper_kwargs(
941
942
943
944
945
946
                input_ids.shape[-1],
                config.forced_bos_token_id,
                config.forced_eos_token_id,
            )

            # Sample constraints
947
948
            min_id = 3
            max_id = model.config.vocab_size
949
            force_tokens = torch.randint(min_id, max_id, (1, 2)).tolist()[0]
950
951
952
953
            constraints = [
                PhrasalConstraint(force_tokens),
            ]

954
955
            beam_kwargs = self._get_constrained_beam_kwargs()
            output_generate = self._constrained_beam_search_generate(
956
957
958
959
960
961
962
                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,
963
                output_logits=True,
964
                output_hidden_states=True,
965
                output_attentions=self.has_attentions,
966
967
968
969
                return_dict_in_generate=True,
            )

            if model.config.is_encoder_decoder:
970
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + 1)
971
972
                self.assertIsInstance(output_generate, GenerateBeamEncoderDecoderOutput)
                # Retrocompatibility check
973
974
                self.assertIsInstance(output_generate, BeamSearchEncoderDecoderOutput)
            else:
975
                self.assertTrue(output_generate.sequences.shape[-1] == self.max_new_tokens + input_ids.shape[-1])
976
977
                self.assertIsInstance(output_generate, GenerateBeamDecoderOnlyOutput)
                # Retrocompatibility check
978
979
                self.assertIsInstance(output_generate, BeamSearchDecoderOnlyOutput)

980
981
            self._check_outputs(
                output_generate, input_ids, model.config, num_return_sequences=beam_kwargs["num_beams"]
982
983
            )

984
985
    def test_contrastive_generate(self):
        for model_class in self.all_generative_model_classes:
986
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
987
                self.skipTest(reason="Stateful models don't support contrastive search generation")
988

989
            # won't fix: FSMT and Reformer have a different cache variable type (and format).
990
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
amyeroberts's avatar
amyeroberts committed
991
                self.skipTest(reason="Won't fix: old model with different cache format")
992

993
            config, input_ids, attention_mask = self._get_input_ids_and_config()
994
995
996

            # NOTE: contrastive search only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
997
                self.skipTest(reason="This model doesn't support caching")
998
999
1000
1001
1002
            config.use_cache = True
            config.is_decoder = True

            # test old generation output for backwards compatibility
            model = model_class(config).to(torch_device).eval()
1003
            output_generate = self._contrastive_generate(
1004
                model=model, input_ids=input_ids, attention_mask=attention_mask
1005
            )
1006
1007
1008
1009
            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])
1010
1011
1012

    def test_contrastive_generate_dict_outputs_use_cache(self):
        for model_class in self.all_generative_model_classes:
1013
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
1014
                self.skipTest(reason="Stateful models don't support contrastive search generation")
1015

1016
            # won't fix: FSMT and Reformer have a different cache variable type (and format).
1017
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
amyeroberts's avatar
amyeroberts committed
1018
                self.skipTest(reason="Won't fix: old model with different cache format")
1019

1020
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1021
1022
1023

            # NOTE: contrastive search only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
1024
                self.skipTest(reason="This model doesn't support caching")
1025
1026
1027
1028
            config.use_cache = True
            config.is_decoder = True

            model = model_class(config).to(torch_device).eval()
1029
            output_generate = self._contrastive_generate(
1030
1031
1032
1033
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_scores=True,
1034
                output_logits=True,
1035
                output_hidden_states=True,
1036
                output_attentions=self.has_attentions,
1037
1038
1039
                return_dict_in_generate=True,
            )

1040
1041
1042
1043
            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])
1044
            self._check_outputs(output_generate, input_ids, model.config, use_cache=True)
1045

1046
1047
1048
    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:
1049
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
1050
                self.skipTest(reason="Stateful models don't support contrastive search generation")
1051

1052
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer", "speech2text"]):
amyeroberts's avatar
amyeroberts committed
1053
                self.skipTest(reason="Won't fix: old model with different cache format")
1054
            if any(model_name in model_class.__name__.lower() for model_name in ["gptbigcode"]):
amyeroberts's avatar
amyeroberts committed
1055
                self.skipTest(reason="TODO: fix me")
1056

1057
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1058
1059
1060

            # NOTE: contrastive search only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
1061
                self.skipTest(reason="This model doesn't support caching")
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073

            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,
1074
                max_new_tokens=self.max_new_tokens,
1075
1076
1077
1078
1079
1080
1081
1082
                attention_mask=attention_mask,
            )

            high_output = model.generate(
                input_ids,
                top_k=4,
                penalty_alpha=0.6,
                low_memory=False,
1083
                max_new_tokens=self.max_new_tokens,
1084
1085
1086
1087
                attention_mask=attention_mask,
            )
            self.assertListEqual(low_output.tolist(), high_output.tolist())

1088
1089
1090
    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:
1091
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
1092
                self.skipTest(reason="May fix in the future: need custom cache handling")
1093
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
amyeroberts's avatar
amyeroberts committed
1094
                self.skipTest(reason="Won't fix: old model with different cache format")
1095
1096
1097
1098
1099
1100
1101
1102
1103
            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
1104
                    "jamba",
1105
1106
                ]
            ):
amyeroberts's avatar
amyeroberts committed
1107
                self.skipTest(reason="May fix in the future: need model-specific fixes")
1108
            config, input_ids, _ = self._get_input_ids_and_config(batch_size=2)
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
            # 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())

1124
    @parameterized.expand([("random",), ("same",)])
1125
    @is_flaky()  # Read NOTE (1) below. If there are API issues, all attempts will fail.
1126
    def test_assisted_decoding_matches_greedy_search(self, assistant_type):
1127
        # This test ensures that the assisted generation does not introduce output changes over greedy search.
1128
1129
1130
1131
1132
        # 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:
1133
        # - assisted_decoding, contrarily to the other methods, can't be called on its own (e.g. needs to
1134
        # prepare the assistant encoder outputs in the main generate body);
1135
1136
        # - assisted_decoding does not support `use_cache = False`
        # - assisted_decoding does not support `batch_size > 1`
1137
1138

        for model_class in self.all_generative_model_classes:
1139
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
1140
                self.skipTest(reason="Stateful models don't support assisted generation")
1141
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
amyeroberts's avatar
amyeroberts committed
1142
                self.skipTest(reason="Won't fix: old model with different cache format")
1143
1144
            if any(
                model_name in model_class.__name__.lower()
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
1155
            ):
amyeroberts's avatar
amyeroberts committed
1156
                self.skipTest(reason="May fix in the future: need model-specific fixes")
1157

1158
            # enable cache
1159
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1160

1161
1162
            # NOTE: assisted generation only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
1163
                self.skipTest(reason="This model doesn't support caching")
1164

1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
            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,
1180
                "output_logits": True,
1181
                "output_hidden_states": True,
1182
                "output_attentions": self.has_attentions,
1183
1184
1185
1186
                "return_dict_in_generate": True,
            }
            output_greedy = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

1187
1188
1189
1190
1191
1192
1193
            # 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
1194
1195
1196
1197
1198
1199
1200
1201
1202
            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)
1203

1204
1205
1206
1207
1208
1209
    @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:
1210
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
1211
                self.skipTest(reason="Stateful models don't support assisted generation")
1212
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
amyeroberts's avatar
amyeroberts committed
1213
                self.skipTest(reason="Won't fix: old model with different cache format")
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
            if any(
                model_name in model_class.__name__.lower()
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
            ):
amyeroberts's avatar
amyeroberts committed
1227
                self.skipTest(reason="May fix in the future: need model-specific fixes")
1228
1229

            # enable cache
1230
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1231
1232
1233

            # NOTE: assisted generation only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
1234
                self.skipTest(reason="This model doesn't support caching")
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250

            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,
1251
                "output_logits": True,
1252
                "output_hidden_states": True,
1253
                "output_attentions": self.has_attentions,
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
                "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)

1267
    def test_assisted_decoding_sample(self):
1268
1269
1270
        # 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).
1271
        for model_class in self.all_generative_model_classes:
1272
            if model_class._is_stateful:
amyeroberts's avatar
amyeroberts committed
1273
                self.skipTest(reason="Stateful models don't support assisted generation")
1274
            if any(model_name in model_class.__name__.lower() for model_name in ["fsmt", "reformer"]):
amyeroberts's avatar
amyeroberts committed
1275
                self.skipTest(reason="Won't fix: old model with different cache format")
1276
1277
            if any(
                model_name in model_class.__name__.lower()
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
                for model_name in [
                    "bigbirdpegasus",
                    "led",
                    "mega",
                    "speech2text",
                    "git",
                    "prophetnet",
                    "seamlessm4t",
                    "clvp",
                ]
1288
            ):
amyeroberts's avatar
amyeroberts committed
1289
                self.skipTest(reason="May fix in the future: need model-specific fixes")
1290
1291

            # enable cache
1292
            config, input_ids, attention_mask = self._get_input_ids_and_config(batch_size=1)
1293
1294
1295

            # NOTE: assisted generation only works with cache on at the moment.
            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
1296
                self.skipTest(reason="This model doesn't support caching")
1297
1298
1299
1300

            config.use_cache = True
            config.is_decoder = True
            model = model_class(config).to(torch_device).eval()
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
            # 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,
1317
                "output_logits": True,
1318
                "output_hidden_states": True,
1319
                "output_attentions": self.has_attentions,
1320
1321
1322
                "return_dict_in_generate": True,
            }
            output_assisted = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)
1323
1324
1325

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

1326
1327
1328
1329
    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:
1330
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1331
1332
1333
            # We want to test only encoder-decoder models
            if not config.is_encoder_decoder:
                continue
Joao Gante's avatar
Joao Gante committed
1334
            model = model_class(config).to(torch_device)
1335
1336

            head_masking = {
1337
1338
1339
1340
1341
1342
1343
                "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
                ),
1344
1345
1346
1347
            }

            signature = inspect.signature(model.forward)
            # We want to test only models where encoder/decoder head masking is implemented
1348
            if not set(head_masking.keys()) < {*signature.parameters.keys()}:
1349
1350
1351
1352
1353
                continue

            for attn_name, (name, mask) in zip(attention_names, head_masking.items()):
                out = model.generate(
                    input_ids,
1354
                    attention_mask=attention_mask,
1355
                    num_beams=1,
1356
                    output_attentions=self.has_attentions,
1357
                    return_dict_in_generate=True,
1358
                    remove_invalid_values=True,
1359
1360
1361
1362
1363
1364
                    **{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)

1365
    def test_left_padding_compatibility(self):
1366
1367
        # NOTE: left-padding results in small numerical differences. This is expected.
        # See https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535
1368

1369
1370
1371
1372
1373
        # 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.")

1374
1375
1376
1377
        # - The model must support padding
        if not self.has_attentions:
            self.skipTest(reason="This model doesn't support padding.")

1378
1379
        # - The model must be a decoder-only architecture (encoder-based architectures use right-padding)
        decoder_only_classes = []
1380
        for model_class in self.all_generative_model_classes:
1381
            config, _, _ = self._get_input_ids_and_config()
1382
            if config.is_encoder_decoder:
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
                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:
1414
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1415
1416
1417
            model = model_class(config).to(torch_device).eval()
            signature = inspect.signature(model.forward).parameters.keys()

1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
            # 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))
1432

1433
1434
1435
1436
1437
1438
1439
1440
    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"):
amyeroberts's avatar
amyeroberts committed
1441
                self.skipTest(reason="This model doesn't support caching")
1442
1443
1444
1445
1446
1447
1448
1449

            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:
amyeroberts's avatar
amyeroberts committed
1450
                self.skipTest(reason="This model doesn't return `past_key_values`")
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503

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

1504
1505
1506
1507
    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:
1508
            config, input_ids, _ = self._get_input_ids_and_config()
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550

            # 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(),
1551
                outputs_from_embeds_wo_ids.tolist(),
1552
1553
            )

1554
1555
1556
1557
    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"]):
amyeroberts's avatar
amyeroberts committed
1558
                self.skipTest(reason="Won't fix: old model with unique inputs/caches/other")
1559
            if any(model_name in model_class.__name__.lower() for model_name in ["umt5"]):
amyeroberts's avatar
amyeroberts committed
1560
                self.skipTest(reason="TODO: needs modeling or test input preparation fixes for compatibility")
1561
1562
1563
1564

            config, inputs = self.model_tester.prepare_config_and_inputs_for_common()

            if not hasattr(config, "use_cache"):
amyeroberts's avatar
amyeroberts committed
1565
                self.skipTest(reason="This model doesn't support caching")
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583

            # 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

1584
            # If "past_key_values" is not returned, skip the test (e.g. RWKV uses a different cache name and format)
1585
1586
            outputs = model(**inputs)
            if "past_key_values" not in outputs:
amyeroberts's avatar
amyeroberts committed
1587
                self.skipTest(reason="This model doesn't return `past_key_values`")
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629

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

1630
1631
1632
1633
1634
1635
1636
    @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:
amyeroberts's avatar
amyeroberts committed
1637
                self.skipTest(reason="This model does not support the new cache format")
1638

1639
            config, input_ids, attention_mask = self._get_input_ids_and_config()
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
            config.use_cache = 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)
1656
1657
1658
1659
1660
1661
            if config.is_encoder_decoder:
                cache_cls = EncoderDecoderCache
                past_key_values = cache_cls(DynamicCache(), DynamicCache())
            else:
                cache_cls = DynamicCache
                past_key_values = cache_cls()
1662
            new_results = model.generate(
1663
                input_ids, attention_mask=attention_mask, past_key_values=past_key_values, **generation_kwargs
1664
1665
1666
1667
1668
1669
            )

            # 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))
1670
            self.assertTrue(isinstance(new_results.past_key_values, cache_cls))
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684

            # 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
1685
            legacy_cache_converted = cache_cls.from_legacy_cache(legacy_results.past_key_values)
1686
1687
1688
1689
1690
1691
1692
1693
1694
            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],
                        )
                    )

1695
1696
1697
1698
    @require_quanto
    def test_generate_with_quant_cache(self):
        for model_class in self.all_generative_model_classes:
            if not model_class._supports_quantized_cache:
amyeroberts's avatar
amyeroberts committed
1699
                self.skipTest(reason="This model does not support the quantized cache format")
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727

            config, input_ids, attention_mask = self._get_input_ids_and_config()
            config.use_cache = True
            config.is_decoder = True

            model = model_class(config).to(torch_device).eval()
            generation_kwargs = {
                "max_new_tokens": 5,
                "cache_implementation": "quantized",
                # careful with group size, should be divisor of model's hidden size
                "cache_config": {"backend": "quanto", "nbits": 2, "q_group_size": 8, "residual_length": 128},
                "return_dict_in_generate": True,  # Required to return `past_key_values`
            }

            results = model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)
            self.assertTrue(isinstance(results.past_key_values, QuantoQuantizedCache))

            # passing past key values of different type should raise Error
            with self.assertRaises(ValueError):
                model.generate(
                    input_ids, attention_mask=attention_mask, past_key_valyes=DynamicCache(), **generation_kwargs
                )

            # setting incorrect cache_config args should raise an Error, i.e. nbits=60 does not make sense
            generation_kwargs["cache_config"] = {"nbits": 60, "q_group_size": 8, "residual_length": 128}
            with self.assertRaises(ValueError):
                model.generate(input_ids, attention_mask=attention_mask, **generation_kwargs)

1728
1729
1730
    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
1731

1732
1733
1734
1735
1736
1737
1738
        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)

1739
1740
1741
        # unprocessed logits
        self._check_logits(num_sequences_in_output, output.logits, config=config)

1742
        # Attentions
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
        if self.has_attentions:
            if config.is_encoder_decoder:
                # encoder
                self._check_encoder_attention_for_generate(output.encoder_attentions, batch_size, config, seq_length)
                # 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,
                )
1768
1769
1770
1771

        # Hidden States
        if config.is_encoder_decoder:
            # encoder
1772
1773
            self._check_encoder_hidden_states_for_generate(
                output.encoder_hidden_states, batch_size, config, seq_length
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
            )

            # 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
1798
        # Past Key Value States -- a few notes here:
1799
1800
        # 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
1801
1802
        # 3. TODO (joao): A few models have different formats/types, skipping those until the cache refactor is
        # complete
1803
        models_without_standard_cache = ("bloom", "ctrl", "fsmt", "gptbigcode", "mega", "reformer", "jamba", "mamba")
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
        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,
            )

1817
1818
1819
1820
1821
1822
    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))

1823
1824
1825
1826
1827
1828
1829
1830
    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))

1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
    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)
            )

1855
1856
1857
1858
1859
1860
1861
1862
    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),
        )

1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
    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),
            )
1881

1882
1883
1884
1885
1886
1887
1888
1889
    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),
        )

1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
    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),
        )

1914
    def _check_sequence_inside_sequence(self, tensor_1, tensor_2):
1915
        # check if tensor_1 inside tensor_2 or tensor_2 inside tensor_1.
1916
1917
        # set to same device. we don't care what device.

1918
1919
1920
1921
1922
1923
        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)
1924
1925
1926
1927
        longer = tensor_2 if in_order else tensor_1
        shorter = tensor_1 if in_order else tensor_2

        flag = False
1928
1929
        chunk_size = len(shorter)
        for chunk_idx in range(len(longer) - chunk_size + 1):
1930
            subseq = longer[chunk_idx : chunk_idx + chunk_size]
1931
            if subseq == shorter:
1932
1933
1934
1935
1936
                flag = True
                break

        self.assertTrue(flag)

1937
1938
1939

@require_torch
class UtilsFunctionsTest(unittest.TestCase):
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
    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])

1975
1976

@require_torch
1977
1978
1979
1980
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 = {
1981
            "AutoModelForCausalLM": AutoModelForCausalLM,
1982
            "AutoModelForSpeechSeq2Seq": AutoModelForSpeechSeq2Seq,
1983
            "AutoModelForSeq2SeqLM": AutoModelForSeq2SeqLM,
1984
            "AutoModelForVision2Seq": AutoModelForVision2Seq,
1985
1986
            "LogitsProcessorList": LogitsProcessorList,
            "MinLengthLogitsProcessor": MinLengthLogitsProcessor,
1987
            "create_tensor_fn": torch.tensor,
1988
            "floats_tensor": floats_tensor,
1989
1990
1991
            "return_tensors": "pt",
        }

1992
1993
    @slow
    def test_diverse_beam_search(self):
1994
        # PT-only test: TF doesn't have a diverse beam search implementation
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
        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(
2005
2006
2007
2008
2009
2010
            input_ids,
            num_beams=4,
            num_return_sequences=2,
            num_beam_groups=4,
            diversity_penalty=2.0,
            remove_invalid_values=True,
2011
2012
2013
2014
2015
2016
2017
        )

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

        self.assertListEqual(
            generated_text,
            [
Sylvain Gugger's avatar
Sylvain Gugger committed
2018
2019
2020
2021
2022
2023
                "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.",
2024
2025
            ],
        )
2026

2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
    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])

2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
    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])

2055
    def test_custom_stopping_criteria_overload_error(self):
2056
        # PT-only test: TF doesn't have StoppingCriteria
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
        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):
2070
        # PT-only test: TF doesn't have StoppingCriteria
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
        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],
        )

2092
    # TODO (joao): replace `stop_sequence` in the pipeline by the more recent `generate` functionality
2093
    def test_stop_sequence_stopping_criteria(self):
2094
        # PT-only test: TF doesn't have StoppingCriteria
2095
2096
2097
2098
2099
        prompt = """Hello I believe in"""
        generator = pipeline("text-generation", model="hf-internal-testing/tiny-random-bart")
        output = generator(prompt)
        self.assertEqual(
            output,
2100
            [{"generated_text": ("Hello I believe in we we we we we we we we we")}],
2101
2102
        )

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

2106
    def test_generate_non_nlp_input_ids_as_kwarg(self):
2107
        # PT-only test: AFAIK there's no non-NLP model architecture in TF that supports `input_ids` as its only input
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
        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))

2119
    def test_generate_input_values_as_encoder_kwarg(self):
2120
        # PT-only test: AFAIK there's no generate-capable architecture in TF that supports `input_values` as its input
2121
2122
2123
2124
2125
2126
2127
2128
2129
        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))

2130
    def test_transition_scores_group_beam_search_encoder_decoder(self):
2131
        # PT-only test: TF doesn't have group beam search
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
        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,
2143
            diversity_penalty=1.0,
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
            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)

2154
        transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, outputs.beam_indices)
2155
2156
2157
        transition_scores_sum = transition_scores.sum(-1)

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

2159
    def test_beam_search_low_memory(self):
2160
2161
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
        model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
        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())

2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
    @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)

2184
2185
        # We will not check watermarked text, since we check it in `logits_processors` tests
        # Checking if generated ids are as expected fails on different hardware
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
        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)

2196
        # Check that the detector is detecting watermarked text
2197
2198
2199
2200
2201
2202
2203
        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)

        self.assertListEqual(detection_out_watermarked.prediction.tolist(), [True])
        self.assertListEqual(detection_out.prediction.tolist(), [False])

2204
2205
    @slow
    def test_beam_search_example_integration(self):
2206
        # PT-only test: TF doesn't have a BeamSearchScorer
2207
2208
        # exactly the example provided in the docstrings of beam search, which previously
        # failed after directly copying from it. Refer to PR #15555
2209
2210
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2211
2212
2213
2214
2215
2216
2217

        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
2218
        input_ids = torch.ones((1, 1), device=model.device, dtype=torch.long)
2219
2220
2221
        input_ids = input_ids * model.config.decoder_start_token_id

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

2224
2225
        outputs = model.generate(
            input_ids, num_beams=num_beams, min_length=5, eos_token_id=model.config.eos_token_id, **model_kwargs
2226
2227
2228
2229
2230
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

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

2231
2232
    @slow
    def test_constrained_beam_search(self):
2233
        # PT-only test: TF doesn't have constrained beam search
2234
2235
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2236

2237
2238
        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
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263

        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,
            [
2264
2265
                "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"
2266
2267
2268
            ],
        )

2269
2270
    @slow
    def test_constrained_beam_search_mixed(self):
2271
        # PT-only test: TF doesn't have constrained beam search
2272
2273
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303

        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,
            [
2304
2305
2306
                "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",
2307
2308
2309
2310
2311
            ],
        )

    @slow
    def test_constrained_beam_search_mixed_mixin(self):
2312
        # PT-only test: TF doesn't have constrained beam search
2313
2314
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341

        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,
            [
2342
2343
2344
                "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",
2345
2346
2347
            ],
        )

2348
2349
    @slow
    def test_cfg_mixin(self):
2350
2351
        model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(torch_device)
        tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387

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

2388
2389
    @slow
    def test_constrained_beam_search_example_translation_mixin(self):
2390
        # PT-only test: TF doesn't have constrained beam search
2391
2392
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410

        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)

2411
        self.assertListEqual(outputs, ["Wie alt sind Sie?"])
2412

2413
2414
    @slow
    def test_constrained_beam_search_example_integration(self):
2415
        # PT-only test: TF doesn't have constrained beam search
2416
2417
        tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
        model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
2418
2419
2420
2421
2422
2423
2424

        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
2425
        input_ids = torch.ones((1, 1), device=model.device, dtype=torch.long)
2426
2427
2428
        input_ids = input_ids * model.config.decoder_start_token_id

        # add encoder_outputs to model keyword arguments
2429
        model_kwargs = {"encoder_outputs": model.get_encoder()(encoder_input_ids, return_dict=True)}
2430
2431
2432
2433

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

2434
2435
2436
2437
2438
2439
2440
        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,
2441
2442
2443
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

2444
        self.assertListEqual(outputs, ["Wie alt sind Sie?"])
2445

2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
    @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)

2483
    def test_constrained_beam_search_mixin_type_checks(self):
2484
        # PT-only test: TF doesn't have constrained beam search
2485
2486
        tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/t5-tiny-random")
        model = AutoModelForSeq2SeqLM.from_pretrained("patrickvonplaten/t5-tiny-random")
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522

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

2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
    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())

2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
    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))

2573
    def test_contrastive_search_batched(self):
2574
        # PT-only test: TF doesn't have constrained beam search
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
        # 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)

2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
    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())

2620
    def test_eos_token_id_int_and_list_top_k_top_sampling(self):
2621
        # Has TF equivalent: this test relies on random sampling
2622
2623
2624
2625
2626
2627
2628
        generation_kwargs = {
            "do_sample": True,
            "num_beams": 1,
            "top_p": 0.7,
            "top_k": 10,
            "temperature": 0.7,
        }
2629
        expectation = 20
2630

2631
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
2632
        text = """Hello, my dog is cute and"""
2633
        tokens = tokenizer(text, return_tensors="pt").to(torch_device)
2634
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
2635

2636
2637
2638
        # 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)
2639
2640
2641
2642
        eos_token_id = 846
        generated_tokens = model.generate(**tokens, eos_token_id=eos_token_id, **generation_kwargs)
        self.assertTrue(expectation == len(generated_tokens[0]))

2643
        torch.manual_seed(1)
2644
        eos_token_id = [846, 198]
2645
2646
        generated_tokens = model.generate(**tokens, eos_token_id=eos_token_id, **generation_kwargs)
        self.assertTrue(expectation == len(generated_tokens[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
    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")
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704

    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:
2705
            # generation_config is modified -> legacy mode is disabled = generation_config takes precedence
2706
2707
2708
            model.generation_config.max_length = 10
            model.generate(input_ids)
            self.assertEqual(len(warning_list), 0)
2709

2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
    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)

2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
    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):
2795
2796
2797
2798
2799
2800
2801
2802
        """
        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)
        """

2803
2804
2805
        # PT-only test: TF doesn't support assisted decoding yet.
        # Bart subclass with a kwarg that distorts the output
        class FakeBart(BartForConditionalGeneration):
2806
2807
            def forward(self, input_ids, past_key_values, foo=False, **kwargs):
                outs = super().forward(input_ids, past_key_values=past_key_values, **kwargs)
2808
2809
2810
2811
                if foo:
                    outs["logits"][:, :, :] = 0.0
                return outs

2812
2813
            def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None, **kwargs):
                kwargs["encoder_outputs"] = encoder_outputs
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
                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
2832
        outputs_foo = model.generate(input_ids, foo=True)
2833
2834
2835
2836
        with self.assertRaises(AssertionError):
            self.assertListEqual(outputs_foo.tolist(), outputs_normal.tolist())

        # Assistant model
2837
2838
2839
        assistant = FakeBart.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
            torch_device
        )
2840
2841
2842
2843
2844
2845
2846
2847

        # 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())
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858

        # 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())
2859
2860

    def test_assisted_decoding_encoder_decoder_shared_encoder(self):
2861
2862
2863
2864
2865
2866
2867
2868
        """
        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)
        """

2869
2870
        # PT-only test: TF doesn't support assisted decoding yet.
        # Bart subclass with a kwarg called foo that distorts the output
2871
        class FakeBartSeq2Seq(BartForConditionalGeneration):
2872
2873
2874
2875
2876
2877
2878
2879
2880
            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)
2881
2882
2883
2884
2885
2886
2887
2888
2889
                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
2890

2891
2892
2893
            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)
2894
2895
2896
                inputs["foo"] = foo
                return inputs

2897
        model = FakeBartSeq2Seq.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
            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
2916
2917
2918
        assistant = FakeBartCausalLM.from_pretrained(
            "hf-internal-testing/tiny-random-BartForConditionalGeneration"
        ).to(torch_device)
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936

        # 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())
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
2978
2979
2980
2981
2982

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

2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
3001
3002
3003
3004
3005
3006
3007
3008
3009
3010
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
3021
3022
3023
3024
3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
3041
3042
3043
3044
    @slow
    def test_validate_assistant(self):
        # Generate a random sample:
        inputs = np.random.rand(160000)

        # Load a main encoder-decoder model:
        model_id = "openai/whisper-large-v2"
        processor = AutoProcessor.from_pretrained(model_id)
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        )
        model.to(torch_device)

        # process the input:
        features = processor(inputs, return_tensors="pt").to(torch_device)

        # Load an encoder-decoder assistant with same encoder as the main model:
        assistant_distil_model_id = "distil-whisper/distil-large-v2"
        assistant_seq_to_seq = AutoModelForSpeechSeq2Seq.from_pretrained(
            assistant_distil_model_id,
            use_safetensors=True,
        ).to(torch_device)
        self.assertTrue(model.generate(**features, assistant_model=assistant_seq_to_seq).sum())

        # Load its decoder only version:
        assistant_causal_lm = AutoModelForCausalLM.from_pretrained(
            assistant_distil_model_id,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        ).to(torch_device)
        self.assertTrue(model.generate(**features, assistant_model=assistant_causal_lm).sum())

        # Load an encoder-decoder assistant with a different encoder than the main model:
        assistant_distil_model_id = "openai/whisper-tiny"
        assistant_seq_to_seq = AutoModelForSpeechSeq2Seq.from_pretrained(
            assistant_distil_model_id,
            use_safetensors=True,
        ).to(torch_device)
        self.assertTrue(model.generate(**features, assistant_model=assistant_seq_to_seq).sum())

        # Load its decoder only version:
        assistant_causal_lm = AutoModelForCausalLM.from_pretrained(
            assistant_distil_model_id,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        ).to(torch_device)
        # It will raise an error as the encoder of the main and assistant model are not compatible:
        with self.assertRaises(ValueError):
            model.generate(**features, assistant_model=assistant_causal_lm)

        # Load an encoder-decoder model with a different tokenizer than the main model:
        assistant_distil_model_id = "hf-internal-testing/tiny-random-SeamlessM4Tv2ForSpeechToText"
        assistant_seq_to_seq = AutoModelForSpeechSeq2Seq.from_pretrained(
            assistant_distil_model_id,
        ).to(torch_device)
        # This should raise an error as the main and assistant model don't use the same tokenizer:
        with self.assertRaises(ValueError):
            model.generate(**features, assistant_model=assistant_seq_to_seq)

3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
3099
3100
    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)
3101

jiqing-feng's avatar
jiqing-feng committed
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
    @slow
    @require_torch_multi_gpu
    def test_assisted_decoding_in_different_gpu(self):
        # PT-only test: TF doesn't support assisted decoding yet.
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to("cuda:0")
        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to(
            "cuda:1"
        )
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
        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,
            max_new_tokens=20,
        )
        self.assertTrue(input_length <= out.shape[-1] <= input_length + 20)

    @slow
    @require_torch_gpu
    def test_assisted_decoding_in_gpu_cpu(self):
        # PT-only test: TF doesn't support assisted decoding yet.
        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to("cuda")
        assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to(
            "cpu"
        )
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
        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,
            max_new_tokens=20,
        )
        self.assertTrue(input_length <= out.shape[-1] <= input_length + 20)

Ahmed Moubtahij's avatar
Ahmed Moubtahij committed
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186

@require_torch
class TokenHealingTestCase(unittest.TestCase):
    @parameterized.expand(
        [
            (
                "square_bracket",
                'An example ["like this"] and another example [',
                'An example ["like this"] and another example ["',
            ),
            ("url", 'The link is <a href="http:', 'The link is <a href="http://'),
            # aggressive_healing: "http" shouldn't be replaced with "https"
            ("aggressive_healing", 'The link is <a href="http', 'The link is <a href="http'),
            ("trailing_whitespace", "I read a book about ", "I read a book about"),
            ("nothing_to_heal", "I read a book about", "I read a book about"),
            ("single_token", "I", "I"),
            ("empty_prompt", "", ""),
        ]
    )
    @require_auto_gptq
    def test_prompts(self, name, input, expected):
        model_name_or_path = "TheBloke/deepseek-llm-7B-base-GPTQ"
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
        completion_model = AutoModelForCausalLM.from_pretrained(
            model_name_or_path,
            device_map="auto",
            trust_remote_code=False,
            revision="main",
            use_cache=True,
        )
        input_ids = tokenizer(input, return_tensors="pt").input_ids.to(completion_model.device)

        healed_ids = completion_model.heal_tokens(input_ids)
        predicted = tokenizer.decode(healed_ids[0], skip_special_tokens=True)

        self.assertEqual(predicted, expected)

3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
    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)