Skip to content

n best beam search - #48

Open
Niduank wants to merge 2 commits into
senstella:masterfrom
synth-inc:feat/n-best-beam-search
Open

n best beam search#48
Niduank wants to merge 2 commits into
senstella:masterfrom
synth-inc:feat/n-best-beam-search

Conversation

@Niduank

@Niduank Niduank commented Jan 21, 2026

Copy link
Copy Markdown

Summary

  • Add n_best parameter to Beam dataclass to return multiple hypotheses from beam search decoding
  • Add NBestHypothesis dataclass with text, score, and confidence fields
  • Add hypotheses attribute to AlignedResult to store N-best results
  • Update documentation with N-best usage examples

Test plan

  • Verify greedy decoding still works (returns hypotheses=None)
  • Verify beam search with n_best=1 returns single hypothesis
  • Verify beam search with n_best=3 returns 3 hypotheses sorted by score
  • Verify streaming transcription works with new decode signature
  • Verify chunked transcription works
  • Test with v3 model (mlx-community/parakeet-tdt-0.6b-v3)

Usage

  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.wav", decoding_config=config)                                                        
                                                                                                                        
  for hyp in result.hypotheses:                                                                                         
      print(f"Text: {hyp.text}, Score: {hyp.score:.2f}, Confidence: {hyp.confidence:.2f}")

- Add n_best parameter to Beam dataclass (default: 1)
- Add NBestHypothesis dataclass to store hypothesis info (text, score, confidence)
- Modify decode_beam() to return N best hypotheses sorted by normalized score
- Add hypotheses attribute to AlignedResult to store N-best results
- Update ParakeetTDT and ParakeetRNNT generate() methods to populate hypotheses
- Update streaming code to handle new decode() return signature
- Fix typo in beam decoding example (decoding -> Beam)
- Add N-best hypotheses usage example
- Document NBestHypothesis in Timestamp Result section
@Niduank Niduank changed the title Feat/n best beam search n best beam search Jan 21, 2026
@senstella
senstella self-requested a review January 25, 2026 04:14

@senstella senstella left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for taking the time to put this PR! I really appreciate the contribution. I had a few thoughts on the interface design that might help make things easier to maintain in future. Please let me know if you have other ideas!

Comment thread parakeet_mlx/alignment.py
class NBestHypothesis:
"""Represents a single hypothesis from N-best beam search decoding."""

text: str

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think implementing this as other dataclass is a bit redundant. Perhaps we can have list[AlignedSentence] instead?

Comment thread parakeet_mlx/parakeet.py
)
results_nbest.append(n_best_hyps)

return results, results_nbest, results_hidden

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we could have an interface other than result, results_nbest, results_hidden. Perhaps we need to make something like DecodingOutput / DecodingResult so that we can extend that if there's any other modifications or improvements would happen.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants