diff --git a/README.md b/README.md index c6c3864..1d31a64 100644 --- a/README.md +++ b/README.md @@ -161,9 +161,9 @@ from parakeet_mlx import from_pretrained, DecodingConfig, Beam model = from_pretrained("mlx-community/parakeet-tdt-0.6b-v3") config = DecodingConfig( - decoding = decoding( + decoding = Beam( beam_size=5, length_penalty=0.013, patience=3.5, duration_reward=0.67 - # Refer to CLI options for each parameters + # Refer to CLI options for each parameters ) ) @@ -172,6 +172,24 @@ result = model.transcribe("audio_file.wav", decoding_config=config) print(result.sentences) ``` +Get N-best hypotheses with beam decoding: + +```py +from parakeet_mlx import from_pretrained, DecodingConfig, Beam + +model = from_pretrained("mlx-community/parakeet-tdt-0.6b-v3") + +config = DecodingConfig( + decoding = Beam(beam_size=5, n_best=3) +) + +result = model.transcribe("audio_file.wav", decoding_config=config) + +# Access N-best hypotheses +for hyp in result.hypotheses: + print(f"Text: {hyp.text}, Score: {hyp.score:.2f}, Confidence: {hyp.confidence:.2f}") +``` + Use local attention: ```py @@ -217,6 +235,11 @@ Using `from_pretrained` downloads a model from Hugging Face and stores the downl - `AlignedResult`: Top-level result containing the full text and sentences - `text`: Full transcribed text - `sentences`: List of `AlignedSentence` + - `hypotheses`: List of `NBestHypothesis` (only with beam decoding and `n_best > 1`) +- `NBestHypothesis`: Alternative hypotheses from beam search + - `text`: Hypothesis text + - `score`: Log probability score + - `confidence`: Confidence score (0.0 to 1.0) - `AlignedSentence`: Sentence-level alignments with start/end times - `text`: Sentence text - `start`: Start time in seconds diff --git a/parakeet_mlx/__init__.py b/parakeet_mlx/__init__.py index 738217c..e1d1c99 100644 --- a/parakeet_mlx/__init__.py +++ b/parakeet_mlx/__init__.py @@ -2,6 +2,7 @@ AlignedResult, AlignedSentence, AlignedToken, + NBestHypothesis, SentenceConfig, ) from parakeet_mlx.parakeet import ( @@ -40,4 +41,5 @@ "AlignedResult", "AlignedSentence", "AlignedToken", + "NBestHypothesis", ] diff --git a/parakeet_mlx/alignment.py b/parakeet_mlx/alignment.py index c6b0714..1d72125 100644 --- a/parakeet_mlx/alignment.py +++ b/parakeet_mlx/alignment.py @@ -35,10 +35,20 @@ def __post_init__(self) -> None: self.confidence = float(np.exp(np.mean(np.log(confidences + 1e-10)))) +@dataclass +class NBestHypothesis: + """Represents a single hypothesis from N-best beam search decoding.""" + + text: str + score: float + confidence: float + + @dataclass class AlignedResult: text: str sentences: list[AlignedSentence] + hypotheses: list[NBestHypothesis] | None = None # N-best hypotheses from beam search def __post_init__(self) -> None: self.text = self.text.strip() diff --git a/parakeet_mlx/parakeet.py b/parakeet_mlx/parakeet.py index 12a6b45..7f21015 100644 --- a/parakeet_mlx/parakeet.py +++ b/parakeet_mlx/parakeet.py @@ -11,6 +11,7 @@ from parakeet_mlx.alignment import ( AlignedResult, AlignedToken, + NBestHypothesis, SentenceConfig, merge_longest_common_subsequence, merge_longest_contiguous, @@ -81,6 +82,7 @@ class Greedy: @dataclass class Beam: beam_size: int = 5 + n_best: int = 1 # Number of hypotheses to return length_penalty: float = 1.0 patience: float = 1.0 duration_reward: float = 0.7 # TDT-only @@ -292,15 +294,20 @@ def decode( hidden_state: Optional[list[Optional[tuple[mx.array, mx.array]]]] = None, *, config: DecodingConfig = DecodingConfig(), - ) -> tuple[list[list[AlignedToken]], list[Optional[tuple[mx.array, mx.array]]]]: + ) -> tuple[ + list[list[AlignedToken]], + list[list[NBestHypothesis]] | None, + list[Optional[tuple[mx.array, mx.array]]], + ]: """Run TDT decoder with features, optional length and decoder state. Outputs list[list[AlignedToken]] and updated hidden state""" mx.eval(features) match config.decoding: case Greedy(): - return self.decode_greedy( + tokens, hidden = self.decode_greedy( features, lengths, last_token, hidden_state, config=config ) + return tokens, None, hidden case Beam(): return self.decode_beam( features, lengths, last_token, hidden_state, config=config @@ -318,7 +325,11 @@ def decode_beam( hidden_state: Optional[list[Optional[tuple[mx.array, mx.array]]]] = None, *, config: DecodingConfig = DecodingConfig(), - ) -> tuple[list[list[AlignedToken]], list[Optional[tuple[mx.array, mx.array]]]]: + ) -> tuple[ + list[list[AlignedToken]], + list[list[NBestHypothesis]], + list[Optional[tuple[mx.array, mx.array]]], + ]: assert isinstance(config.decoding, Beam) # type guarntee beam_token = min(config.decoding.beam_size, len(self.vocabulary) + 1) @@ -348,7 +359,10 @@ def __hash__(self) -> int: if last_token is None: last_token = list([None] * B) + n_best = config.decoding.n_best + results = [] + results_nbest = [] results_hidden = [] for batch in range(B): feature = features[batch : batch + 1] @@ -507,21 +521,47 @@ def __hash__(self) -> int: if not finished_hypothesis: results.append([]) + results_nbest.append([]) results_hidden.append(hidden_state[batch]) else: length_penalty = ( config.decoding.length_penalty ) # mypy assumes weirdly so we go in safe way - best = max( + # Sort hypotheses by normalized score + sorted_hyps = sorted( finished_hypothesis, key=lambda x: x.score / (max(1, len(x.hypothesis)) ** length_penalty), + reverse=True, ) + + # Get the best hypothesis for backward compatibility + best = sorted_hyps[0] results.append(best.hypothesis) results_hidden.append(best.hidden_state) - return results, results_hidden + # Build N-best hypotheses list + n_best_hyps = [] + for hyp in sorted_hyps[:n_best]: + hyp_text = "".join( + token.text for token in hyp.hypothesis + ).strip() + hyp_len = max(1, len(hyp.hypothesis)) + normalized_score = hyp.score / hyp_len + confidence = math.exp(normalized_score) + # Clamp confidence to [0, 1] range + confidence = min(1.0, max(0.0, confidence)) + n_best_hyps.append( + NBestHypothesis( + text=hyp_text, + score=hyp.score, + confidence=confidence, + ) + ) + results_nbest.append(n_best_hyps) + + return results, results_nbest, results_hidden def decode_greedy( self, @@ -626,14 +666,21 @@ def generate( features, lengths = self.encoder(mel) mx.eval(features, lengths) - result, _ = self.decode(features, lengths, config=decoding_config) + tokens_result, nbest_result, _ = self.decode( + features, lengths, config=decoding_config + ) - return [ - sentences_to_result( + results = [] + for i, hypothesis in enumerate(tokens_result): + aligned_result = sentences_to_result( tokens_to_sentences(hypothesis, decoding_config.sentence) ) - for hypothesis in result - ] + # Add N-best hypotheses if available (beam search) + if nbest_result is not None and i < len(nbest_result): + aligned_result.hypotheses = nbest_result[i] + results.append(aligned_result) + + return results class ParakeetRNNT(BaseParakeet): @@ -660,8 +707,12 @@ def decode( hidden_state: Optional[list[Optional[tuple[mx.array, mx.array]]]] = None, *, config: DecodingConfig = DecodingConfig(), - ) -> tuple[list[list[AlignedToken]], list[Optional[tuple[mx.array, mx.array]]]]: - """Run TDT decoder with features, optional length and decoder state. Outputs list[list[AlignedToken]] and updated hidden state""" + ) -> tuple[ + list[list[AlignedToken]], + list[list[NBestHypothesis]] | None, + list[Optional[tuple[mx.array, mx.array]]], + ]: + """Run RNNT decoder with features, optional length and decoder state. Outputs list[list[AlignedToken]] and updated hidden state""" assert isinstance(config.decoding, Greedy), ( "Only greedy decoding is supported for RNNT decoder now" ) @@ -740,7 +791,7 @@ def decode( results.append(hypothesis) - return results, hidden_state + return results, None, hidden_state def generate( self, mel: mx.array, *, decoding_config: DecodingConfig = DecodingConfig() @@ -751,14 +802,21 @@ def generate( features, lengths = self.encoder(mel) mx.eval(features, lengths) - result, _ = self.decode(features, lengths, config=decoding_config) + tokens_result, nbest_result, _ = self.decode( + features, lengths, config=decoding_config + ) - return [ - sentences_to_result( + results = [] + for i, hypothesis in enumerate(tokens_result): + aligned_result = sentences_to_result( tokens_to_sentences(hypothesis, decoding_config.sentence) ) - for hypothesis in result - ] + # Add N-best hypotheses if available + if nbest_result is not None and i < len(nbest_result): + aligned_result.hypotheses = nbest_result[i] + results.append(aligned_result) + + return results class ParakeetCTC(BaseParakeet): @@ -1059,7 +1117,7 @@ def add_audio(self, audio: mx.array) -> None: finalized_length = max(0, length - self.drop_size) if isinstance(self.model, ParakeetTDT) or isinstance(self.model, ParakeetRNNT): - finalized_tokens, finalized_state = self.model.decode( + finalized_tokens, _, finalized_state = self.model.decode( features, mx.array([finalized_length]), [self.last_token], @@ -1072,7 +1130,7 @@ def add_audio(self, audio: mx.array) -> None: finalized_tokens[0][-1].id if len(finalized_tokens[0]) > 0 else None ) - draft_tokens, _ = self.model.decode( + draft_tokens, _, _ = self.model.decode( features[:, finalized_length:], mx.array( [