From 417ee2f17d6ade110ab48a9b5fea7fe6d1de0113 Mon Sep 17 00:00:00 2001 From: Sahil Faizal Date: Mon, 17 Aug 2026 19:23:53 -0700 Subject: [PATCH] Add Whisper beam-search decoding --- ACKNOWLEDGMENTS.md | 3 +- whisper/README.md | 13 ++++ whisper/mlx_whisper/cli.py | 6 ++ whisper/mlx_whisper/decoding.py | 131 +++++++++++++++++++++++++++++++- whisper/test.py | 19 +++++ 5 files changed, 170 insertions(+), 2 deletions(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index c6853710d..047a5e0d5 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -14,4 +14,5 @@ MLX Examples was developed with contributions from the following individuals: - Markus Enzweiler: Added the `cvae` examples. - Prince Canuma: Helped add support for `Starcoder2` models. - Shiyu Li: Added the `Segment Anything Model`. -- Gökdeniz Gülmez: Added support for `MiniCPM`, `Helium`, `Mamba version 1`, `OLMoE` archtectures and support for `full-fine-tuning`. \ No newline at end of file +- Gökdeniz Gülmez: Added support for `MiniCPM`, `Helium`, `Mamba version 1`, `OLMoE` archtectures and support for `full-fine-tuning`. +- Sahil Faizal: Implemented beam-search decoding support for Whisper. diff --git a/whisper/README.md b/whisper/README.md index cd3bc684a..e669e98f3 100644 --- a/whisper/README.md +++ b/whisper/README.md @@ -35,6 +35,13 @@ Use `-f` to specify the output format and `--model` to specify the model. There are many other supported command line options. To see them all, run `mlx_whisper -h`. +For deterministic beam-search decoding, set `--beam-size` (and optionally +`--patience` or `--length-penalty`): + +```sh +mlx_whisper audio_file.mp3 --beam-size 5 +``` + You can also pipe the audio content of other programs via stdin: ```sh @@ -61,6 +68,12 @@ setting `path_or_hf_repo`. For example: result = mlx_whisper.transcribe(speech_file, path_or_hf_repo="models/large") ``` +The same option is available through the Python API: + +```python +result = mlx_whisper.transcribe(speech_file, beam_size=5) +``` + This will load the model contained in `models/large`. The `path_or_hf_repo` can also point to an MLX-style Whisper model on the Hugging Face Hub. In this case, the model will be automatically downloaded. A [collection of pre-converted diff --git a/whisper/mlx_whisper/cli.py b/whisper/mlx_whisper/cli.py index ee8212648..f1437cab3 100644 --- a/whisper/mlx_whisper/cli.py +++ b/whisper/mlx_whisper/cli.py @@ -92,6 +92,12 @@ def str2bool(string): default=5, help="Number of candidates when sampling with non-zero temperature", ) + parser.add_argument( + "--beam-size", + type=optional_int, + default=None, + help="Number of beams to use for deterministic decoding", + ) parser.add_argument( "--patience", type=float, diff --git a/whisper/mlx_whisper/decoding.py b/whisper/mlx_whisper/decoding.py index 814dc95ca..fb5a64a45 100644 --- a/whisper/mlx_whisper/decoding.py +++ b/whisper/mlx_whisper/decoding.py @@ -189,8 +189,11 @@ def scores(logprobs, lengths): class TokenDecoder: + source_indices: Optional[List[int]] = None + def reset(self): """Initialize any stateful variables for decoding a new sequence""" + self.source_indices = None def update( self, tokens: mx.array, logits: mx.array, sum_logprobs: mx.array @@ -283,6 +286,123 @@ def finalize(self, tokens: mx.array, sum_logprobs: mx.array): return tokens, sum_logprobs +class BeamSearchDecoder(TokenDecoder): + """Conventional beam search with optional patience-based early stopping.""" + + def __init__(self, beam_size: int, eot: int, patience: Optional[float] = None): + self.beam_size = beam_size + self.eot = eot + self.max_candidates = round(beam_size * (patience or 1.0)) + if self.max_candidates < beam_size: + raise ValueError("patience must be greater than or equal to 1") + self.finished_sequences: Optional[List[Dict[Tuple[int, ...], float]]] = None + self.source_indices = None + + def reset(self): + self.finished_sequences = None + self.source_indices = None + + def update( + self, tokens: mx.array, logits: mx.array, sum_logprobs: mx.array + ) -> Tuple[mx.array, bool, mx.array]: + if tokens.shape[0] % self.beam_size != 0: + raise ValueError("the number of sequences must be a multiple of beam_size") + + n_audio = tokens.shape[0] // self.beam_size + if self.finished_sequences is None: + self.finished_sequences = [{} for _ in range(n_audio)] + + logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) + candidate_tokens = mx.argpartition( + -logprobs, kth=self.beam_size, axis=-1 + )[:, : self.beam_size + 1] + candidate_logprobs = logprobs[ + mx.arange(logprobs.shape[0])[:, None], candidate_tokens + ] + mx.eval(candidate_tokens, candidate_logprobs, tokens, sum_logprobs) + + candidate_tokens = np.array(candidate_tokens) + candidate_logprobs = np.array(candidate_logprobs) + current_tokens = tokens.tolist() + current_logprobs = sum_logprobs.tolist() + + next_tokens: List[List[int]] = [] + next_logprobs: List[float] = [] + source_indices: List[int] = [] + for audio_index in range(n_audio): + scores: Dict[Tuple[int, ...], float] = {} + sources: Dict[Tuple[int, ...], int] = {} + for beam_index in range(self.beam_size): + index = audio_index * self.beam_size + beam_index + for token, logprob in zip( + candidate_tokens[index], candidate_logprobs[index] + ): + if not np.isfinite(logprob): + continue + sequence = tuple(current_tokens[index] + [int(token)]) + score = current_logprobs[index] + float(logprob) + if score > scores.get(sequence, -np.inf): + scores[sequence] = score + sources[sequence] = index + + saved = 0 + for sequence in sorted(scores, key=scores.get, reverse=True): + if sequence[-1] == self.eot: + self.finished_sequences[audio_index][sequence] = scores[sequence] + else: + next_tokens.append(list(sequence)) + next_logprobs.append(scores[sequence]) + source_indices.append(sources[sequence]) + saved += 1 + if saved == self.beam_size: + break + + # Keep the batch shape stable when every top candidate reached EOT. + if saved < self.beam_size: + for sequence in sorted(scores, key=scores.get, reverse=True): + if sequence[-1] != self.eot or saved >= self.beam_size: + continue + next_tokens.append(list(sequence[:-1]) + [self.eot]) + next_logprobs.append(scores[sequence]) + source_indices.append(sources[sequence]) + saved += 1 + + self.source_indices = source_indices + completed = all( + len(sequences) >= self.max_candidates + for sequences in self.finished_sequences + ) + return mx.array(next_tokens), completed, mx.array(next_logprobs) + + def finalize(self, tokens: mx.array, sum_logprobs: mx.array): + if self.finished_sequences is None: + raise RuntimeError("beam search has not been initialized") + + mx.eval(tokens, sum_logprobs) + active_tokens = tokens.tolist() + active_logprobs = sum_logprobs.tolist() + candidates: List[List[List[int]]] = [] + candidate_logprobs: List[List[float]] = [] + max_length = 0 + + for audio_index, finished in enumerate(self.finished_sequences): + sequences = dict(finished) + for beam_index in range(self.beam_size): + sequence = tuple(active_tokens[audio_index][beam_index] + [self.eot]) + sequences.setdefault(sequence, active_logprobs[audio_index][beam_index]) + + best = sorted(sequences, key=sequences.get, reverse=True)[: self.beam_size] + candidates.append([list(sequence) for sequence in best]) + candidate_logprobs.append([sequences[sequence] for sequence in best]) + max_length = max(max_length, *(len(sequence) for sequence in best)) + + padded = [ + [sequence + [self.eot] * (max_length - len(sequence)) for sequence in group] + for group in candidates + ] + return mx.array(padded), mx.array(candidate_logprobs) + + class LogitFilter: def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: """Apply any filtering or masking to logits @@ -434,7 +554,9 @@ def __init__(self, model: "Whisper", options: DecodingOptions): # decoder: implements how to select the next tokens, given the autoregressive distribution if options.beam_size is not None: - raise NotImplementedError("Beam search decoder is not yet implemented") + self.decoder = BeamSearchDecoder( + options.beam_size, tokenizer.eot, options.patience + ) else: self.decoder = GreedyDecoder(options.temperature, tokenizer.eot) @@ -470,6 +592,10 @@ def _verify_options(self, options: DecodingOptions) -> DecodingOptions: raise ValueError("best_of with greedy sampling (T=0) is not compatible") if options.patience is not None and options.beam_size is None: raise ValueError("patience requires beam_size to be given") + if options.beam_size is not None and options.beam_size <= 0: + raise ValueError("beam_size must be a positive integer") + if options.patience is not None and options.patience < 1: + raise ValueError("patience must be greater than or equal to 1") if options.length_penalty is not None and not ( 0 <= options.length_penalty <= 1 ): @@ -587,6 +713,8 @@ def _step(inputs, audio_features, tokens, sum_logprobs): tokens, completed, sum_logprobs = self.decoder.update( tokens, logits, sum_logprobs ) + if self.decoder.source_indices is not None: + self.inference.rearrange_kv_cache(self.decoder.source_indices) return tokens, completed, sum_logprobs, pre_logits tokens, completed, sum_logprobs, pre_logits = _step( @@ -645,6 +773,7 @@ def run(self, mel: mx.array) -> List[DecodingResult]: tokens, [n_audio, self.n_group, len(self.initial_tokens)] ) tokens = tokens.reshape((n_audio * self.n_group, len(self.initial_tokens))) + audio_features = mx.repeat(audio_features, self.n_group, axis=0) # call the main sampling loop tokens, sum_logprobs, no_speech_probs = self._main_loop(audio_features, tokens) diff --git a/whisper/test.py b/whisper/test.py index f0acb3cd9..c72990562 100644 --- a/whisper/test.py +++ b/whisper/test.py @@ -185,6 +185,25 @@ def test_decode_greedy(self): self.assertAlmostEqual(result.no_speech_prob, 0.009631240740418434, places=4) self.assertAlmostEqual(result.compression_ratio, 1.2359550561797752) + def test_decode_beam_search(self): + result = decoding.decode(self.model, self.mels, beam_size=3, fp16=False) + + self.assertEqual(result.language, "en") + self.assertTrue(result.tokens) + self.assertTrue(result.text) + self.assertTrue(np.isfinite(result.avg_logprob)) + + def test_beam_search_options(self): + with self.assertRaises(ValueError): + decoding.DecodingTask( + self.model, decoding.DecodingOptions(beam_size=0, fp16=False) + ) + with self.assertRaises(ValueError): + decoding.DecodingTask( + self.model, + decoding.DecodingOptions(beam_size=2, patience=0.5, fp16=False), + ) + def test_transcribe(self): result = mlx_whisper.transcribe( TEST_AUDIO, path_or_hf_repo=MLX_FP32_MODEL_PATH, fp16=False