Skip to content

[generate] fix breaking change for patch #29976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ArthurZucker merged 12 commits into from
Apr 2, 2024
Merged

[generate] fix breaking change for patch #29976

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
5a38dcf
fix bug and add tests
ArthurZucker Apr 1, 2024
2eab875
nit
ArthurZucker Apr 1, 2024
c6248dd
otherway to get the cur len instead of attention mask
ArthurZucker Apr 1, 2024
dc4e768
more places where this might have been broken
ArthurZucker Apr 1, 2024
bb34689
nit
ArthurZucker Apr 1, 2024
aafd697
oups
ArthurZucker Apr 1, 2024
11f12fc
inputs_embeds vs input_embeds
ArthurZucker Apr 1, 2024
62b8dfb
test generated outptus
ArthurZucker Apr 1, 2024
c1bd53e
style
ArthurZucker Apr 1, 2024
bb72f89
nit
ArthurZucker Apr 1, 2024
bc36242
fix
ArthurZucker Apr 1, 2024
4567319
skip failing biogpt
ArthurZucker Apr 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/transformers/generation/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3034,6 +3034,8 @@ def _beam_search(
num_beams = beam_scorer.num_beams

batch_beam_size, cur_len = input_ids.shape
if "inputs_embeds" in model_kwargs:
cur_len = model_kwargs["inputs_embeds"].shape[1]
model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)

if num_beams * batch_size != batch_beam_size:
Expand Down Expand Up @@ -3437,6 +3439,8 @@ def _beam_sample(
num_beams = beam_scorer.num_beams

batch_beam_size, cur_len = input_ids.shape
if "inputs_embeds" in model_kwargs:
cur_len = model_kwargs["inputs_embeds"].shape[1]
model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)

# init attention / hidden states / scores tuples
Expand Down Expand Up @@ -3795,6 +3799,8 @@ def _group_beam_search(
device = input_ids.device

batch_beam_size, cur_len = input_ids.shape
if "inputs_embeds" in model_kwargs:
cur_len = model_kwargs["inputs_embeds"].shape[1]
model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)

if return_dict_in_generate and output_scores:
Expand Down Expand Up @@ -4211,6 +4217,8 @@ def _constrained_beam_search(
num_beams = constrained_beam_scorer.num_beams

batch_beam_size, cur_len = input_ids.shape
if "inputs_embeds" in model_kwargs:
cur_len = model_kwargs["inputs_embeds"].shape[1]
model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)

if num_beams * batch_size != batch_beam_size:
Expand Down
12 changes: 12 additions & 0 deletions tests/generation/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -718,6 +718,18 @@ def test_beam_sample_generate(self):

self.assertTrue(output_generate.shape[-1] == max_length)

input_embeds = model.get_input_embeddings()(input_ids)
beam_kwargs.update({"input_embeds": input_embeds})
output_generate = self._beam_search_generate(
model=model,
input_ids=None,
attention_mask=attention_mask,
max_length=max_length,
beam_kwargs=beam_kwargs,
)

self.assertTrue(output_generate.shape[-1] == max_length)

def test_beam_sample_generate_dict_output(self):
for model_class in self.all_generative_model_classes:
config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
Expand Down