Add Whisper candidate output and beam search

whisper/README.md

@@ -76,6 +76,13 @@ output = mlx_whisper.transcribe(speech_file, word_timestamps=True)
 print(output["segments"][0]["words"])
 ```
 
+To inspect ranked decoding candidates, set `return_candidates=True`:
+
+```python
+output = mlx_whisper.transcribe(speech_file, return_candidates=True)
+print(output["segments"][0]["candidates"])
+```
+
 To see more transcription options use:
 
 ```

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 in beam search, only applicable when temperature is zero",
+    )
     parser.add_argument(
         "--patience",
         type=float,

whisper/mlx_whisper/decoding.py

@@ -2,7 +2,7 @@
 
 import zlib
 from dataclasses import dataclass, field, replace
-from typing import Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
 
 import mlx.core as mx
 import numpy as np
@@ -114,6 +114,7 @@ class DecodingOptions:
 
     # implementation details
     fp16: bool = True  # use fp16 for most of the calculation
+    return_candidates: bool = False  # include all ranked decoding candidates
 
 
 @dataclass(frozen=True)
@@ -127,6 +128,7 @@ class DecodingResult:
     no_speech_prob: float = np.nan
     temperature: float = np.nan
     compression_ratio: float = np.nan
+    candidates: List[Dict[str, Any]] = field(default_factory=list)
 
 
 class Inference:
@@ -152,6 +154,15 @@ def reset(self):
 
 
 class SequenceRanker:
+    def scores(
+        self, tokens: List[List[mx.array]], sum_logprobs: List[List[float]]
+    ) -> List[List[float]]:
+        """
+        Given a list of groups of samples and their cumulative log probabilities,
+        return the score used to rank each sample in each group
+        """
+        raise NotImplementedError
+
     def rank(
         self, tokens: List[List[mx.array]], sum_logprobs: List[List[float]]
     ) -> List[int]:
@@ -171,7 +182,9 @@ class MaximumLikelihoodRanker(SequenceRanker):
     def __init__(self, length_penalty: Optional[float]):
         self.length_penalty = length_penalty
 
-    def rank(self, tokens: List[List[List[int]]], sum_logprobs: List[List[float]]):
+    def scores(
+        self, tokens: List[List[List[int]]], sum_logprobs: List[List[float]]
+    ) -> List[List[float]]:
         def scores(logprobs, lengths):
             result = []
             for logprob, length in zip(logprobs, lengths):
@@ -183,9 +196,12 @@ def scores(logprobs, lengths):
                 result.append(logprob / penalty)
             return result
 
-        # get the sequence with the highest score
         lengths = [[len(t) for t in s] for s in tokens]
-        return [np.argmax(scores(p, l)) for p, l in zip(sum_logprobs, lengths)]
+        return [scores(p, l) for p, l in zip(sum_logprobs, lengths)]
+
+    def rank(self, tokens: List[List[List[int]]], sum_logprobs: List[List[float]]):
+        # get the sequence with the highest score
+        return [np.argmax(scores) for scores in self.scores(tokens, sum_logprobs)]
 
 
 class TokenDecoder:
@@ -283,6 +299,147 @@ def finalize(self, tokens: mx.array, sum_logprobs: mx.array):
         return tokens, sum_logprobs
 
 
+class BeamSearchDecoder(TokenDecoder):
+    def __init__(
+        self,
+        beam_size: int,
+        eot: int,
+        inference: Inference,
+        patience: Optional[float] = None,
+    ):
+        self.beam_size = beam_size
+        self.eot = eot
+        self.inference = inference
+        self.patience = patience or 1.0
+        self.max_candidates = round(beam_size * self.patience)
+        self.finished_sequences = None
+
+        assert (
+            self.max_candidates > 0
+        ), f"Invalid beam size ({beam_size}) or patience ({patience})"
+
+    def reset(self):
+        self.finished_sequences = 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(f"{tokens.shape}[0] % {self.beam_size} != 0")
+
+        n_audio = tokens.shape[0] // self.beam_size
+        if self.finished_sequences is None:
+            self.finished_sequences = [{} for _ in range(n_audio)]
+
+        logprobs = logits.astype(mx.float32) - mx.logsumexp(
+            logits.astype(mx.float32), axis=-1, keepdims=True
+        )
+
+        mx.eval(tokens, logprobs, sum_logprobs)
+        tokens_np = np.array(tokens)
+        logprobs_np = np.array(logprobs)
+        sum_logprobs_np = np.array(sum_logprobs)
+
+        next_tokens = []
+        next_logprobs = []
+        source_indices = []
+        finished_sequences = []
+
+        for i in range(n_audio):
+            scores = {}
+            sources = {}
+            finished = {}
+
+            for j in range(self.beam_size):
+                idx = i * self.beam_size + j
+                prefix = tokens_np[idx].tolist()
+                row = logprobs_np[idx]
+                top_indices = np.argsort(row)[-(self.beam_size + 1) :][::-1]
+                for token in top_indices:
+                    score = float(sum_logprobs_np[idx] + row[token])
+                    sequence = tuple(prefix + [int(token)])
+                    if sequence not in scores or score > scores[sequence]:
+                        scores[sequence] = score
+                        sources[sequence] = idx
+
+            saved = 0
+            for sequence in sorted(scores, key=scores.get, reverse=True):
+                if sequence[-1] == self.eot:
+                    finished[sequence] = scores[sequence]
+                else:
+                    next_tokens.append(sequence)
+                    next_logprobs.append(scores[sequence])
+                    source_indices.append(sources[sequence])
+                    saved += 1
+                    if saved == self.beam_size:
+                        break
+
+            finished_sequences.append(finished)
+
+        tokens = mx.array(next_tokens, dtype=tokens.dtype)
+        sum_logprobs = mx.array(next_logprobs, dtype=sum_logprobs.dtype)
+        self.inference.rearrange_kv_cache(source_indices)
+
+        assert len(self.finished_sequences) == len(finished_sequences)
+        for previously_finished, newly_finished in zip(
+            self.finished_sequences, finished_sequences
+        ):
+            previously_finished.update(newly_finished)
+            sorted_sequences = sorted(
+                previously_finished.items(), key=lambda item: item[1], reverse=True
+            )[: self.max_candidates]
+            previously_finished.clear()
+            previously_finished.update(sorted_sequences)
+
+        completed = all(
+            len(sequences) >= self.max_candidates
+            for sequences in self.finished_sequences
+        )
+        return tokens, completed, sum_logprobs
+
+    def finalize(self, tokens: mx.array, sum_logprobs: mx.array):
+        if self.finished_sequences is None:
+            self.finished_sequences = [{} for _ in range(tokens.shape[0])]
+
+        mx.eval(tokens, sum_logprobs)
+        tokens_np = np.array(tokens)
+        sum_logprobs_np = np.array(sum_logprobs)
+
+        for i, sequences in enumerate(self.finished_sequences):
+            if len(sequences) < self.beam_size:
+                for j in np.argsort(sum_logprobs_np[i])[::-1]:
+                    sequence = tuple(tokens_np[i, j].tolist() + [self.eot])
+                    sequences[sequence] = float(sum_logprobs_np[i, j])
+                    if len(sequences) >= self.beam_size:
+                        break
+
+        n_candidates = max(len(sequences) for sequences in self.finished_sequences)
+        max_length = max(
+            len(sequence)
+            for sequences in self.finished_sequences
+            for sequence in sequences
+        )
+        padded_tokens = np.full(
+            (len(self.finished_sequences), n_candidates, max_length),
+            self.eot,
+            dtype=tokens_np.dtype,
+        )
+        padded_logprobs = np.full(
+            (len(self.finished_sequences), n_candidates),
+            -np.inf,
+            dtype=sum_logprobs_np.dtype,
+        )
+
+        for i, sequences in enumerate(self.finished_sequences):
+            for j, sequence in enumerate(sequences.keys()):
+                padded_tokens[i, j, : len(sequence)] = sequence
+                padded_logprobs[i, j] = sequences[sequence]
+
+        return mx.array(padded_tokens, dtype=tokens.dtype), mx.array(
+            padded_logprobs, dtype=sum_logprobs.dtype
+        )
+
+
 class LogitFilter:
     def apply(self, logits: mx.array, tokens: mx.array) -> mx.array:
         """Apply any filtering or masking to logits
@@ -434,7 +591,12 @@ 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,
+                self.inference,
+                options.patience,
+            )
         else:
             self.decoder = GreedyDecoder(options.temperature, tokenizer.eot)
 
@@ -623,6 +785,7 @@ def run(self, mel: mx.array) -> List[DecodingResult]:
         n_audio: int = mel.shape[0]
 
         audio_features: mx.array = self._get_audio_features(mel)  # encoder forward pass
+        original_audio_features = audio_features
         tokens: mx.array = mx.array(self.initial_tokens)
         tokens = mx.broadcast_to(tokens, (n_audio, len(self.initial_tokens)))
 
@@ -645,14 +808,30 @@ 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 = audio_features[:, None, :, :]
+            audio_features = mx.broadcast_to(
+                audio_features,
+                [
+                    n_audio,
+                    self.n_group,
+                    original_audio_features.shape[1],
+                    original_audio_features.shape[2],
+                ],
+            )
+            audio_features = audio_features.reshape(
+                (
+                    n_audio * self.n_group,
+                    original_audio_features.shape[1],
+                    original_audio_features.shape[2],
+                )
+            )
 
         # call the main sampling loop
         tokens, sum_logprobs, no_speech_probs = self._main_loop(audio_features, tokens)
 
         # reshape the tensors to have (n_audio, n_group) as the first two dimensions
-        audio_features = audio_features[:: self.n_group]
         no_speech_probs = no_speech_probs[:: self.n_group]
-        assert audio_features.shape[0] == len(no_speech_probs) == n_audio
+        assert original_audio_features.shape[0] == len(no_speech_probs) == n_audio
 
         tokens = tokens.reshape(n_audio, self.n_group, -1)
         sum_logprobs = sum_logprobs.reshape(n_audio, self.n_group)
@@ -666,10 +845,55 @@ def run(self, mel: mx.array) -> List[DecodingResult]:
         tokens = tokens.tolist()
         sum_logprobs = sum_logprobs.tolist()
         no_speech_probs = no_speech_probs.tolist()
-        tokens = [[t[: t.index(tokenizer.eot)] for t in s] for s in tokens]
+        ended_with_eot = [[tokenizer.eot in t for t in s] for s in tokens]
+        tokens = [
+            [t[: t.index(tokenizer.eot)] if tokenizer.eot in t else t for t in s]
+            for s in tokens
+        ]
 
         # select the top-ranked sample in each group
-        selected = self.sequence_ranker.rank(tokens, sum_logprobs)
+        selected = [int(i) for i in self.sequence_ranker.rank(tokens, sum_logprobs)]
+        candidate_groups: List[List[Dict[str, Any]]] = []
+        if self.options.return_candidates:
+            scores = self.sequence_ranker.scores(tokens, sum_logprobs)
+            for (
+                group_tokens,
+                group_logprobs,
+                group_scores,
+                group_eot,
+                selected_idx,
+            ) in zip(
+                tokens,
+                sum_logprobs,
+                scores,
+                ended_with_eot,
+                selected,
+            ):
+                ranked_indices = sorted(
+                    range(len(group_tokens)),
+                    key=lambda index: group_scores[index],
+                    reverse=True,
+                )
+                candidates = []
+                for rank, index in enumerate(ranked_indices):
+                    candidate_tokens = group_tokens[index]
+                    candidate_text = tokenizer.decode(candidate_tokens).strip()
+                    sum_logprob = float(group_logprobs[index])
+                    candidates.append(
+                        {
+                            "rank": rank,
+                            "index": int(index),
+                            "selected": int(index) == selected_idx,
+                            "text": candidate_text,
+                            "tokens": candidate_tokens,
+                            "sum_logprob": sum_logprob,
+                            "avg_logprob": sum_logprob / (len(candidate_tokens) + 1),
+                            "score": float(group_scores[index]),
+                            "ended_with_eot": bool(group_eot[index]),
+                            "compression_ratio": compression_ratio(candidate_text),
+                        }
+                    )
+                candidate_groups.append(candidates)
         tokens: List[List[int]] = [t[i] for i, t in zip(selected, tokens)]
         texts: List[str] = [tokenizer.decode(t).strip() for t in tokens]
 
@@ -682,12 +906,16 @@ def run(self, mel: mx.array) -> List[DecodingResult]:
             texts,
             languages,
             tokens,
-            audio_features,
+            original_audio_features,
             avg_logprobs,
             no_speech_probs,
         )
         if len(set(map(len, fields))) != 1:
             raise RuntimeError(f"inconsistent result lengths: {list(map(len, fields))}")
+        if self.options.return_candidates and len(candidate_groups) != len(texts):
+            raise RuntimeError(
+                f"inconsistent candidate lengths: {len(candidate_groups)} != {len(texts)}"
+            )
 
         return [
             DecodingResult(
@@ -699,10 +927,18 @@ def run(self, mel: mx.array) -> List[DecodingResult]:
                 no_speech_prob=no_speech_prob,
                 temperature=self.options.temperature,
                 compression_ratio=compression_ratio(text),
+                candidates=(
+                    candidate_groups[index] if self.options.return_candidates else []
+                ),
             )
-            for text, language, tokens, features, avg_logprob, no_speech_prob in zip(
-                *fields
-            )
+            for index, (
+                text,
+                language,
+                tokens,
+                features,
+                avg_logprob,
+                no_speech_prob,
+            ) in enumerate(zip(*fields))
         ]