Grammar-Constrained Decoding

The problem, from first principles

Applications often need the model's output to be valid JSON, a number or a fixed set of labels. Asking nicely in the prompt fails some fraction of the time, and parsing failures break the pipeline. Constrained decoding guarantees validity by editing the logits at every step: tokens that would take the text outside the allowed language are set to -inf so they can never be sampled. The allowed language is described by a small state machine. At each step you ask, for every vocabulary token, whether feeding its characters to the machine from the current state stays valid, and mask the rest.

From theory to code

Implement advance, allowed_token_ids and mask_logits.

Constraints

• The state machine is a dict transitions mapping (state, char) to the next state. A character with no entry is invalid.
• advance(transitions, state, token) feeds the characters of token one at a time and returns the final state, or None if any character is invalid. An empty token returns state unchanged.
• allowed_token_ids(transitions, state, vocab) returns the sorted list of ids i such that advance(transitions, state, vocab[i]) is not None and vocab[i] is non-empty.
• mask_logits(logits, allowed) returns a copy of logits with -inf at every id not in allowed.

Hints