pub struct ConstrainedSampler<C: TokenConstraint, S: Sampler> { /* private fields */ }Expand description
Wraps any Sampler so the constraint’s mask is applied at the front
of the sampling path — the drop-in integration seam for real decoding.
§How it plugs in
ConstrainedSampler::try_sample does, in order: read the generated-so-far
prefix -> ask the TokenConstraint for the AllowedSet -> mask_logits
the raw logits (-inf on disallowed) -> hand the masked row to the inner
Sampler (temperature/top-k/top-p/… all run after the mask, exactly
as required) -> record the chosen token so the next step’s DecodeState
is correct. Because masking happens before the inner sampler, the sampler
can only ever pick an allowed token — greedy’s argmax skips -inf, and
stochastic softmax gives -inf zero probability.
§Why try_sample is fallible instead of an infallible Sampler impl
Sampler::sample returns a bare u32 — it cannot report “the constraint
left nothing to sample”. That case (ConstraintError::NoTokenAllowed)
must be an honest error, not a panic and not a silently-wrong token, so the
constrained path is deliberately its own fallible method rather than an
impl Sampler that would have to swallow the error. Callers drive it
through crate::generate::generate_constrained, which threads the error
out cleanly.
Implementations§
Source§impl<C: TokenConstraint, S: Sampler> ConstrainedSampler<C, S>
impl<C: TokenConstraint, S: Sampler> ConstrainedSampler<C, S>
Sourcepub fn new(constraint: C, inner: S) -> Self
pub fn new(constraint: C, inner: S) -> Self
Wrap inner so constraint masks its logits every step.
Sourcepub fn reset(&mut self)
pub fn reset(&mut self)
Forget the generated history so the same sampler can drive a fresh decode (the constraint restarts from an empty prefix).
Sourcepub fn try_sample(&mut self, logits: &[f32]) -> Result<u32, ConstraintError>
pub fn try_sample(&mut self, logits: &[f32]) -> Result<u32, ConstraintError>
Mask, then sample — the constrained step. See the type’s docs for the exact ordering it guarantees.
§Errors
ConstraintError::NoTokenAllowed if the constraint masked every
in-range token this step (see mask_logits). The generated history is
left unchanged on error, so a caller may retry with a widened
constraint without corrupting the decode state.