Fix beam search with batch processing in Whisper decoding

[zuazo]

When using the Whisper model's decode() method with both beam_size and batch_size greater than 1, a dimension mismatch error occurs. This issue was apparently introduced accidentally in PR #1483.

Example to Reproduce the Error

import torch
import whisper

device = "cuda" if torch.cuda.is_available() else "cpu"
model = whisper.load_model("tiny").to(device)

# load audio and pad/trim it to fit 30 seconds
audio = whisper.load_audio("tests/jfk.flac")
audio = whisper.pad_or_trim(audio)
mel = whisper.log_mel_spectrogram(audio).to(model.device)

# Create a small batch with 2 examples
batch_mel = mel.unsqueeze(0).repeat(2, 1, 1)

# decode the audio with beam size > 1
options = whisper.DecodingOptions(beam_size=5)
results = model.decode(batch_mel, options)

# print the recognized text
for result in results:
    print(result.text)

This outputs the following error:

Traceback (most recent call last):
  File "decode_example.py", line 17, in <module>
    results = model.decode(batch_mel, options)
  File "~/.anaconda3/envs/whisper-dev/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "~/src/whisper/whisper/decoding.py", line 824, in decode
    result = DecodingTask(model, options).run(mel)
  File "~/.anaconda3/envs/whisper-dev/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "~/src/whisper/whisper/decoding.py", line 737, in run
    tokens, sum_logprobs, no_speech_probs = self._main_loop(audio_features, tokens)
  File "~/src/whisper/whisper/decoding.py", line 687, in _main_loop
    logits = self.inference.logits(tokens, audio_features)
  File "~/src/whisper/whisper/decoding.py", line 163, in logits
    return self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache)
  File "~/.anaconda3/envs/whisper-dev/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/src/whisper/whisper/model.py", line 211, in forward
    x = block(x, xa, mask=self.mask, kv_cache=kv_cache)
  File "~/.anaconda3/envs/whisper-dev/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/src/whisper/whisper/model.py", line 138, in forward
    x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
  File "~/.anaconda3/envs/whisper-dev/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/src/whisper/whisper/model.py", line 90, in forward
    wv, qk = self.qkv_attention(q, k, v, mask)
  File "~/src/whisper/whisper/model.py", line 102, in qkv_attention
    qk = q @ k
RuntimeError: The size of tensor a (10) must match the size of tensor b (2) at non-singleton dimension 0

Solution

• Ensure that audio features are correctly duplicated across beams for each batch item.
• Add a test for decode() that includes a regression test for this.
• Update .github/workflows/test.yml to run the new test for decode().

[nevisende]

The problem is solved.