Repository: huggingface/transformers
File: tests/generation/test_utils.py (Python)

Code before:
        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])

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

    def test_model_kwarg_assisted_decoding_decoder_only(self):
        model = AutoModelForCausalLM.from_pretrained(
            "hf-internal-testing/tiny-random-gpt2", attn_implementation="eager"
        ).to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        model.generation_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 + input_ids.shape[1]))

        # 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", attn_implementation="eager"
        ).to(torch_device)
        assistant.config.pad_token_id = tokenizer.eos_token_id


Review comment (performance):
Can we check if we could use something along https://github.com/huggingface/transformers/blob/f2f0d17d160bc239de46771d6ad84a4e5c2bc2e2/tests/test_modeling_common.py#L3644-L3645 instead?

We would have `with sdpa_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):` instead

Code after:

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

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

    def test_model_kwarg_assisted_decoding_decoder_only(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.generation_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)

        with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):
            # Traditional way of generating text
            outputs_normal = model.generate(input_ids)
            self.assertEqual(outputs_normal.shape, (1, 20 + input_ids.shape[1]))

            # 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(