Simple Fix to an Annoying Issue While Batch Decoding with Beam Search

[farisalasmary]

Hi All,

I was trying to do batch inference and decoding using Beam Search Decoder using the following code snippet:

from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE, get_tokenizer
from whisper.decoding import DecodingOptions, DecodingResult
from whisper.audio import (
    FRAMES_PER_SECOND,
    HOP_LENGTH,
    N_FRAMES,
    N_SAMPLES,
    SAMPLE_RATE,
    log_mel_spectrogram,
    pad_or_trim,
)

def load_audio(audio_file_path, start_time=0.0, duration=None, sample_rate=16000):
    audio_signal, _ = librosa.load(
        audio_file_path,
        sr=sample_rate,
        mono=True,
        offset=start_time,
        duration=duration
    )
    return audio_signal

def get_batch_features(audio_signals, n_mels=128):
    mel_segment_list = []
    for audio_signal in audio_signals:
        mel = log_mel_spectrogram(audio_signal, n_mels, padding=N_SAMPLES)
        mel_segment = pad_or_trim(mel, N_FRAMES)
        mel_segment_list.append(mel_segment)
    mel_segments = torch.stack(mel_segment_list)
    return mel_segments

audio_files = ['a.wav', 'b.wav', 'c.wav', 'd.wav', 'e.wav', 'f.wav']
audio_signals = [load_audio(f, sample_rate=16000) for f in audio_files]
mel_segments = get_batch_features(audio_signals, n_mels=128)
mel_segments = mel_segments.to(model.device)

beam_size = 5
options = DecodingOptions(beam_size=beam_size, language='ar', task='transcribe', temperature=0)
decode_result = model.decode(mel_segments, options)

and I faced this annoying issue:

Traceback (most recent call last):
  File "run_whisper_batch.py", line 72, in <module>
    decode_result = model.decode(mel_segments, options)
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 34, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/whisper/decoding.py", line 826, in decode
    result = DecodingTask(model, options).run(mel)
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 34, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/whisper/decoding.py", line 739, in run
    tokens, sum_logprobs, no_speech_probs = self._main_loop(audio_features, tokens)
  File "/usr/local/lib/python3.8/dist-packages/whisper/decoding.py", line 687, in _main_loop
    logits = self.inference.logits(tokens, audio_features)
  File "/usr/local/lib/python3.8/dist-packages/whisper/decoding.py", line 163, in logits
    return self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1480, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/whisper/model.py", line 242, in forward
    x = block(x, xa, mask=self.mask, kv_cache=kv_cache)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1480, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/whisper/model.py", line 169, in forward
    x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1480, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/whisper/model.py", line 111, in forward
    wv, qk = self.qkv_attention(q, k, v, mask)
  File "/usr/local/lib/python3.8/dist-packages/whisper/model.py", line 130, in qkv_attention
    qk = (q * scale) @ (k * scale).transpose(-1, -2)
RuntimeError: The size of tensor a (30) must match the size of tensor b (6) at non-singleton dimension 0

After spending some time and wasting a lot of energy debugging the issue, I found the bug.
In the following code snippet, audio_features should be repeated using repeat_interleave method similar to tokens:

whisper/whisper/decoding.py

Lines 732 to 735 in c0d2f62

	# repeat text tensors by the group size, for beam search or best-of-n sampling
	tokens = tokens.repeat_interleave(self.n_group, dim=0).to(audio_features.device)

The final code should look like this:

        # repeat text tensors by the group size, for beam search or best-of-n sampling
        tokens = tokens.repeat_interleave(self.n_group, dim=0).to(audio_features.device)
        audio_features = audio_features.repeat_interleave(self.n_group, dim=0).to(audio_features.device)

and that's it!