Struct BeamSearchEngine 

pub struct BeamSearchEngine {
    pub config: BeamSearchConfig,
}

Beam search engine.

Decoupled from the model via a get_logits closure so it can be used with any inference backend.

Fields

config: BeamSearchConfig

Search configuration.

Implementations

impl BeamSearchEngine

pub fn new(config: BeamSearchConfig) -> Self

Create a new engine with the given configuration.

pub fn search<F>( &self, initial_tokens: Vec<u32>, _vocab_size: usize, get_logits: F, ) -> BeamSearchResult

where F: FnMut(&[u32], usize) -> Vec<f32>,

Run beam search.

get_logits(beam_tokens, step) is called for every live beam at every step and must return a logit vector of length vocab_size.

pub fn apply_no_repeat_ngram( logits: &mut [f32], tokens: &[u32], ngram_size: usize, )

Zero out (set to −∞) any token that would create a repeated n-gram.

For each position in tokens where the last ngram_size - 1 tokens match the trailing ngram_size - 1 tokens of the current sequence, the following token is forbidden.

pub fn top_k_log_probs(logits: &[f32], k: usize) -> Vec<(u32, f64)>

Return the top-k (token_id, log_prob) pairs from a logit vector.

Logits are converted to log-probabilities via log-softmax.

pub fn prune_beams( beams: Vec<Beam>, beam_width: usize, length_penalty: f32, ) -> Vec<Beam>

Keep the top beam_width beams by length-normalised score.