Skip to content
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

Fix missing sequences_scores in the Whisper beam search output #32970

Merged
merged 5 commits into from
Sep 17, 2024
Merged

Fix missing sequences_scores in the Whisper beam search output #32970

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fb9cf6a
added sequences_scores to the output
Nik-Kras Aug 22, 2024
c1b7b9f
added beam_indices to output
Nik-Kras Sep 10, 2024
ad3e341
added test to check for beam_indices, sequences_scores and their shape
Nik-Kras Sep 10, 2024
f02ba02
removed redundant whitespaces
Nik-Kras Sep 10, 2024
1d00f7e
make fixup
Nik-Kras Sep 12, 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: 5 additions & 3 deletions src/transformers/models/whisper/generation_whisper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -994,19 +994,21 @@ def _stack_split_outputs(self, seek_outputs, model_output_type, device, kwargs):
for key in seek_outputs[0].keys():
if key == "sequences":
outputs[key] = torch.stack([v[key] for v in seek_outputs], dim=0).to(device)
if key in ["scores", "encoder_attentions", "encoder_hidden_states", "logits"]:
elif key in ["scores", "encoder_attentions", "encoder_hidden_states", "logits"]:
outputs[key] = tuple(
torch.stack([v[key][i] for v in seek_outputs]).to(device) for i in range(len(seek_outputs[0][key]))
)
if key in ["decoder_attentions", "decoder_hidden_states", "cross_attentions"]:
elif key == "sequences_scores":
outputs[key] = torch.stack([v[key] for v in seek_outputs], dim=0).to(device)
elif key in ["decoder_attentions", "decoder_hidden_states", "cross_attentions"]:
outputs[key] = tuple(
tuple(
torch.stack([v[key][i][j] for v in seek_outputs]).squeeze(1).to(device)
for j in range(len(seek_outputs[0][key][0]))
)
for i in range(len(seek_outputs[0][key]))
)
if key == "past_key_values":
elif key == "past_key_values":
past_key_value_type = kwargs.get("past_key_values")
if seek_outputs[0][key] is not None:
outputs[key] = tuple(
Expand Down