diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index d131b2f8d5..f8fb086cba 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -441,6 +441,9 @@ class GenerationMixin: if isinstance(value, torch.Tensor): batch_size = value.shape[0] break + + if "inputs_embeds" in model_kwargs: + return torch.ones((batch_size, 0), dtype=torch.long, device=self.device) return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id def _prepare_attention_mask_for_generation( @@ -1421,6 +1424,14 @@ class GenerationMixin: ) generation_config.max_length = generation_config.max_new_tokens + input_ids_length + # otherwise the total length [inputs-embeds-len + new-tokens-len] will go beyond indicated `max_length`` + elif ( + model_input_name == "inputs_embeds" + and inputs_tensor.shape[:-1] != input_ids.shape + and not self.config.is_encoder_decoder + ): + generation_config.max_length -= inputs_tensor.shape[1] + # if we don't pass `past_key_values` and a cache_implementation is specified if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING and not model_kwargs.get( "past_key_values", False diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index c91ff7993a..18e7eb481f 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -1963,7 +1963,7 @@ class GenerationTesterMixin: ) self.assertListEqual( outputs_from_embeds[:, inputs_embeds.shape[1] :].tolist(), - outputs_from_embeds_wo_ids[:, 1:].tolist(), + outputs_from_embeds_wo_ids.tolist(), ) def test_generate_continue_from_past_key_values(self): @@ -2730,6 +2730,20 @@ class GenerationIntegrationTests(unittest.TestCase, GenerationIntegrationTestsMi **model_kwargs, ) + 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]) + def test_custom_stopping_criteria_overload_error(self): # PT-only test: TF doesn't have StoppingCriteria article = """Justin Timberlake and Jessica Biel, welcome to parenthood."""