Skip to main content

generate_constrained

Function generate_constrained 

Source
pub fn generate_constrained<M, C, S>(
    model: &M,
    tokenizer: &dyn Tokenizer,
    prompt: &str,
    config: &GenerationConfig,
    constrained: &mut ConstrainedSampler<C, S>,
    on_token: impl FnMut(u32, &str),
) -> Result<String, ConstrainedGenerateError>
where M: Model, C: TokenConstraint, S: Sampler,
Expand description

Like generate_with_sampler, but every step is grammar-constrained: the ConstrainedSampler’s TokenConstraint masks disallowed tokens to -inf before its inner sampler runs, so the model physically cannot emit a token the constraint forbids. This is the keystone that makes a small model reliably produce valid JSON / valid tool names / valid structure — see [crate::constraint] for the full rationale and provenance (AID-0045).

The decode loop is otherwise identical to generate_with_sampler — prefill, one-token-at-a-time KV-cache decoding, EOS handling, streaming via on_token, the returned completion text. The only differences are that token selection goes through ConstrainedSampler::try_sample (which applies the mask at the front of the sampling path) and that the extra failure mode — the constraint masking every token — is surfaced as ConstrainedGenerateError::Constraint rather than silently emitting rubbish.

constrained carries its own generated-token history for the constraint’s crate::constraint::DecodeState; it should be freshly constructed (or ConstrainedSampler::reset) per decode so the constraint starts from an empty prefix. The constraint’s vocabulary (if it is a crate::constraint::JsonStructure) must be sized to the same vocab as the model’s logits row, or its mask can never allow the higher token ids.

§Errors

ConstrainedGenerateError::Runtime for any tokenize/forward/decode failure (exactly generate’s error set), or ConstrainedGenerateError::Constraint if the constraint leaves no valid token at some step. On a constraint error the partial completion is discarded — the loop stops and returns the error rather than a truncated string, so a caller cannot mistake a masked-out dead end for a finished completion.