Skip to main content

beam_search

Function beam_search 

Source
pub fn beam_search(
    masked_log_probs: &[Vec<Vec<f64>>],
    parent_scores: &[Vec<f64>],
    beam_width: usize,
) -> (Vec<Vec<u32>>, Vec<Vec<f64>>, Vec<Vec<usize>>)
Expand description

Beam search selection over masked log-probabilities.

Scores are accumulated as parent_score + log_prob(token).

Returns (new_tokens, new_scores, src_beams), all shaped [B][M].