[Tests] Diverse Whisper fixes (#33665)

* fix beam indices in token_timestamps

* fix attention_mask in FA2

* correct translation example with the right example

* correct how somes tests are using outputs + correct num_frames

* fix shortform batch prev cond tests

* make fix-copies

* make fix-copies

* take care of shifting beam indices

* [run-slow] whisper

* [run-slow] whisper
This commit is contained in:
Yoach Lacombe
2024-10-03 15:59:01 +02:00
committed by GitHub
parent ab97a78130
commit bf0ffe3d29
4 changed files with 33 additions and 18 deletions

View File

@@ -291,7 +291,7 @@ class Qwen2AudioFlashAttention2(Qwen2AudioAttention):
causal_mask = attention_mask
if attention_mask is not None: # no matter the length, we just slice it
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
causal_mask = attention_mask[:, : key_states.shape[-2]]
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
# therefore the input hidden states gets silently casted in float32. Hence, we need

View File

@@ -173,7 +173,9 @@ def _pad_to_max_length(
class WhisperGenerationMixin(GenerationMixin):
def _extract_token_timestamps(self, generate_outputs, alignment_heads, time_precision=0.02, num_frames=None):
def _extract_token_timestamps(
self, generate_outputs, alignment_heads, time_precision=0.02, num_frames=None, num_input_ids=None
):
"""
Calculates token-level timestamps using the encoder-decoder cross-attentions and dynamic time-warping (DTW) to
map each output token to a position in the input audio. If `num_frames` is specified, the encoder-decoder
@@ -200,11 +202,18 @@ class WhisperGenerationMixin(GenerationMixin):
# since the beam search strategy chooses the most probable sequences at the end of the search.
# In that case, the cross_attentions weights are too long and we have to make sure that they have the right output_length
weight_length = (generate_outputs.beam_indices != -1).sum(-1).max()
weight_length = weight_length if num_input_ids is None else weight_length + num_input_ids
# beam search takes `decoder_input_ids` into account in the `beam_indices` length
# but forgot to shift the beam_indices by the number of `decoder_input_ids`
beam_indices = torch.zeros_like(generate_outputs.beam_indices[:, :weight_length])
# we actually shif the beam indices here
beam_indices[:, num_input_ids:] = generate_outputs.beam_indices[:, : weight_length - num_input_ids]
weights = weights[:, :, :weight_length]
# If beam index is still -1, it means that the associated token id is EOS
# We need to replace the index with 0 since index_select gives an error if any of the indexes is -1.
beam_indices = generate_outputs.beam_indices[:, :weight_length]
beam_indices = beam_indices.masked_fill(beam_indices == -1, 0)
# Select the cross attention from the right beam for each output sequences
@@ -218,8 +227,10 @@ class WhisperGenerationMixin(GenerationMixin):
# make sure timestamps are as long as weights
input_length = weight_length or cross_attentions[0].shape[2]
timestamps = torch.zeros_like(generate_outputs.sequences, dtype=torch.float32)[:, : input_length + 1]
batch_size = timestamps.shape[0]
batch_size = generate_outputs.sequences.shape[0]
timestamps = torch.zeros(
(batch_size, input_length + 1), dtype=torch.float32, device=generate_outputs.sequences.device
)
if num_frames is not None:
# two cases:
@@ -239,6 +250,7 @@ class WhisperGenerationMixin(GenerationMixin):
else:
# num_frames is of shape (batch_size,) whereas batch_size is truely batch_size*num_return_sequences
repeat_time = batch_size if isinstance(num_frames, int) else batch_size // len(num_frames)
num_frames = num_frames.cpu() if isinstance(num_frames, (torch.Tensor)) else num_frames
num_frames = np.repeat(num_frames, repeat_time)
if num_frames is None or isinstance(num_frames, int):
@@ -948,7 +960,10 @@ class WhisperGenerationMixin(GenerationMixin):
if return_token_timestamps and hasattr(generation_config, "alignment_heads"):
num_frames = getattr(generation_config, "num_frames", None)
seek_outputs["token_timestamps"] = self._extract_token_timestamps(
seek_outputs, generation_config.alignment_heads, num_frames=num_frames
seek_outputs,
generation_config.alignment_heads,
num_frames=num_frames,
num_input_ids=decoder_input_ids.shape[-1],
)
seek_outputs["token_timestamps"] = seek_outputs["token_timestamps"][:, start_idx:]

View File

@@ -422,7 +422,7 @@ class WhisperFlashAttention2(WhisperAttention):
causal_mask = attention_mask
if attention_mask is not None: # no matter the length, we just slice it
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
causal_mask = attention_mask[:, : key_states.shape[-2]]
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
# therefore the input hidden states gets silently casted in float32. Hence, we need

View File

@@ -1916,14 +1916,14 @@ class WhisperModelIntegrationTests(unittest.TestCase):
input_features, do_sample=False, max_length=20, language="<|de|>", task="transcribe"
)
transcript = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
EXPECTED_TRANSCRIPT = " Mein sechster Sohn scheint, wenigstens auf den ersten Blick,"
EXPECTED_TRANSCRIPT = " Denken Sie, soeben walten meine Gedanken bei Ihnen in Adela"
self.assertEqual(transcript, EXPECTED_TRANSCRIPT)
generated_ids = model.generate(
input_features, do_sample=False, max_length=20, language="<|de|>", task="translate"
)
transcript = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
EXPECTED_TRANSCRIPT = " My sixth son seems, at least at first glance, the most deeply-minded"
EXPECTED_TRANSCRIPT = " Think, my thoughts were just rolling with you in Adelaide, and I"
self.assertEqual(transcript, EXPECTED_TRANSCRIPT)
@slow
@@ -2238,7 +2238,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
input_features, max_length=448, return_timestamps=True, return_token_timestamps=True
)
self.assertEqual(generate_outputs.sequences.shape, generate_outputs.token_timestamps.shape)
self.assertEqual(generate_outputs["sequences"].shape, generate_outputs["token_timestamps"].shape)
# fmt: off
EXPECTED_OUTPUT = torch.tensor([
@@ -2249,7 +2249,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
])
# fmt: on
self.assertTrue(torch.allclose(generate_outputs.token_timestamps.to("cpu"), EXPECTED_OUTPUT))
self.assertTrue(torch.allclose(generate_outputs["token_timestamps"].to("cpu"), EXPECTED_OUTPUT))
@slow
def test_large_token_timestamp_generation(self):
@@ -2268,7 +2268,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
**input_features, max_length=448, return_timestamps=True, return_token_timestamps=True
)
self.assertEqual(generate_outputs.sequences.shape, generate_outputs.token_timestamps.shape)
self.assertEqual(generate_outputs["sequences"].shape, generate_outputs["token_timestamps"].shape)
# fmt: off
EXPECTED_OUTPUT = torch.tensor([
@@ -2279,7 +2279,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
])
# fmt: on
self.assertTrue(torch.allclose(generate_outputs.token_timestamps.to("cpu"), EXPECTED_OUTPUT))
self.assertTrue(torch.allclose(generate_outputs["token_timestamps"].to("cpu"), EXPECTED_OUTPUT))
@slow
def test_tiny_token_timestamp_batch_generation(self):
@@ -2306,9 +2306,9 @@ class WhisperModelIntegrationTests(unittest.TestCase):
)
# task id and lang id prompts should not have timestamp tokens
self.assertEqual(generate_outputs.sequences.shape[-1] - 2, generate_outputs.token_timestamps.shape[-1])
self.assertEqual(generate_outputs["sequences"].shape[-1] - 2, generate_outputs["token_timestamps"].shape[-1])
self.assertEqual(len(generate_outputs.sequences), num_return_sequences * num_samples)
self.assertEqual(len(generate_outputs["sequences"]), num_return_sequences * num_samples)
@slow
def test_tiny_token_timestamp_generation_longform(self):
@@ -2799,7 +2799,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
torch.manual_seed(0)
result = model.generate(input_features, **gen_kwargs)
decoded = processor.batch_decode(result.sequences, skip_special_tokens=True)
decoded = processor.batch_decode(result, skip_special_tokens=True)
assert decoded == EXPECTED_TEXT
@@ -2814,7 +2814,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
torch.manual_seed(0)
result = model.generate(input_features, **gen_kwargs)
decoded = processor.batch_decode(result.sequences, skip_special_tokens=True)
decoded = processor.batch_decode(result, skip_special_tokens=True)
assert decoded == EXPECTED_TEXT1
@@ -3114,7 +3114,7 @@ class WhisperModelIntegrationTests(unittest.TestCase):
}
result = model.generate(**inputs, **gen_kwargs)
decoded_all = processor.batch_decode(result.sequences, skip_special_tokens=True)
decoded_all = processor.batch_decode(result, skip_special_tokens=True)
for i in range(num_samples):
if isinstance(EXPECTED_TEXT[i], str):