Struct ConstrainedDecoder 

pub struct ConstrainedDecoder {
    pub index: StaticIndex,
    pub beam_width: usize,
    pub batch_size: usize,
}

Fields

index: StaticIndex

beam_width: usize

batch_size: usize

Implementations

impl ConstrainedDecoder

pub fn new(index: StaticIndex, beam_width: usize, batch_size: usize) -> Self

pub fn step(&self, logits: &[Vec<Vec<f64>>], state: &mut BeamState, step: usize)

Execute one constrained decoding step.

Arguments

• logits — raw model outputs, shape [B × M × |V|]
• state — mutable beam state (nodes, scores, partial tokens)
• step — 0-indexed decoding step t

pub fn decode<F>(&self, logit_fn: F, sid_length: usize) -> Vec<Vec<Vec<u32>>>

where F: Fn(&BeamState, usize) -> Vec<Vec<Vec<f64>>>,

Run the full constrained beam-search loop for sid_length steps.

logit_fn is called once per step; it receives the current BeamState and must return logits of shape [B × M × |V|].

Returns the top-beam_width decoded SIDs for every query in the batch.

pub fn dense_lookup( &self, state: &BeamState, step: usize, ) -> (Vec<Vec<Vec<bool>>>, Vec<Vec<Vec<u32>>>)

For steps covered by the bit-packed dense mask, look up validity in O(1) per token without touching the CSR matrix.

Returns (masks, next_nodes) shaped [B][M][|V|] and [B][M][1] respectively (one “next node” per beam; the trie node reached after the chosen token is resolved lazily in gather_state from the dense mask’s states array).

pub fn dense_first_token_valid(&self, tok: u32) -> bool

Returns true if any full-depth dense entry starts with tok.

pub fn sparse_lookup( &self, state: &BeamState, step: usize, ) -> (Vec<Vec<Vec<bool>>>, Vec<Vec<Vec<u32>>>)

For deeper steps, call VNTK on the CSR transition matrix.

Returns (masks, next_nodes) shaped [B][M][|V|] and [B][M][B_t].

pub fn resolve_next_node( &self, current_node: u32, token: u32, next_node_slots: &[u32], step: usize, ) -> u32

Resolves the next trie node for a beam that chose token at step, given the pre-computed next_node_slots from VNTK.

VNTK returns slots sorted by token ID, so we binary-search rather than doing a linear scan or a second CSR lookup.