Struct BeamSearch 

pub struct BeamSearch { /* private fields */ }

Beam search decoder for sequence generation.

Maintains the top-k (beam_size) hypotheses at each decoding step.

Example

let beam_search = BeamSearch::new(5);  // beam_size=5
let hypotheses = beam_search.search(&logits_fn, start_token, max_len);

Implementations

impl BeamSearch

pub fn new(beam_size: usize) -> BeamSearch

Create a new beam search decoder.

Arguments

• beam_size - Number of beams to maintain

pub fn with_length_penalty(self, penalty: f32) -> BeamSearch

Set length penalty (>1 favors longer sequences).

pub fn with_early_stopping(self) -> BeamSearch

Enable early stopping when all beams reach EOS.

pub fn with_eos_token_id(self, eos_token_id: usize) -> BeamSearch

Set EOS token ID.

pub fn step( &self, log_probs: &Tensor, current_beams: &[BeamHypothesis], ) -> Vec<BeamHypothesis>

Perform beam search given log probabilities.

Arguments

• log_probs - Log probabilities for each token [vocab_size]
• current_beams - Current beam hypotheses

Returns

Updated beam hypotheses after one step.

pub fn init(&self, start_token: usize) -> Vec<BeamHypothesis>

Initialize beam search with a start token.

pub fn all_done(&self, beams: &[BeamHypothesis]) -> bool

Check if all beams are done.

pub fn best(&self, beams: &[BeamHypothesis]) -> Option<BeamHypothesis>

Get the best hypothesis.

pub fn beam_size(&self) -> usize

Get beam_size.

pub fn length_penalty(&self) -> f32

Get length_penalty.

Trait Implementations

impl Debug for BeamSearch

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.