Constrained decoding

Grammars let you force generation to match a pattern — digits only, fixed choices, JSON shapes, regex-style constraints. The mask is applied in logit space at every step, so it composes cleanly with greedy sampling, temperature sampling, beam search, and every early-exit setting.

The public surface is the tl.Grammar class plus the grammar= kwarg on Model.generate.

Quickstart

import numpy as np
import tinyloop_py as tl
from transformers import GPT2TokenizerFast

# 1. Load model + tokenizer.
model = tl.Model("model.tinyloop", max_seq_len=512)
tok = GPT2TokenizerFast.from_pretrained("gpt2")

# 2. Build the token-id → raw-bytes map the grammar needs.
token_bytes = _build_token_bytes(tok)  # see "Building token_bytes" below

# 3. Compile a grammar (regex).
g = tl.Grammar(r"[0-9]+")
g.set_token_vocabulary(token_bytes)

# 4. Generate with the grammar — output is guaranteed to be digits only.
prompt = np.asarray(tok.encode("Answer:"), dtype=np.int32)
out = model.generate(prompt, grammar=g, max_tokens=16, temperature=0.0)
text = tok.decode(out[len(prompt):])
print(text)  # e.g. "112345678974467"

tl.Grammar(pattern)

Compile a byte-level regex into an NFA.

g = tl.Grammar(r"\{\"name\":\"[a-z]+\"\}")

Parameters

Name	Type	Description
pattern	str	Regex pattern. Matching is on raw bytes, not Unicode codepoints — see Supported syntax below.

Returns

A compiled tl.Grammar. Raises ValueError on parse error (unbalanced parens, unterminated class, trailing backslash, etc.).

Supported syntax

Thompson's NFA construction with a deliberately small subset of classic regex. Designed for constrained decoding, not full PCRE compatibility.

Construct	Example	Meaning
Literal byte	a	Match the byte a.
Character class	[a-z0-9_]	Match any single byte in the set.
Negated class	[^abc]	Match any single byte not in the set.
Any byte	.	Match any single byte (including \n and \0).
Alternation	a|b	Match either.
Grouping	(ab|cd)	Group a subexpression.
Kleene star	a*	Zero or more.
Plus	a+	One or more.
Optional	a?	Zero or one.
Digit class	\d / \D	One digit / non-digit.
Word class	\w / \W	One word character / non-word.
Whitespace	\s / \S	One whitespace / non-whitespace.
Escape char	\. \\ \[ \] \( \) | \* \+ \?	Literal punctuation.
Escape newline	\n \t \r	Control bytes.

Multi-byte UTF-8 is handled correctly because matching is byte-level: . matches any single byte, not a codepoint. For Unicode-aware matching, express the expected encoding directly:

# Match any Han character (3-byte UTF-8 range 0xE4xx...0xE9xx)
g = tl.Grammar(r"([\xE4-\xE9][\x80-\xBF][\x80-\xBF])+")

Raises

• ValueError — pattern has a parse error. The error message names the character position.

grammar.set_token_vocabulary(token_bytes)

Install the tokenizer's byte-sequence-per-id map. Must be called before passing the grammar to Model.generate.

Parameters

Name	Type	Description
token_bytes	List[bytes]	One bytes object per token id. token_bytes[i] is the raw byte sequence the sampler would emit when token id i is chosen. Length becomes the grammar's vocab_size.

What "raw bytes" means for BPE

Most BPE tokenizers (GPT-2, Llama, etc.) don't store raw bytes directly — they store a "bytes-to-unicode" shifted-string representation internally. The grammar works on the real underlying bytes, so the caller has to reverse that trick.

The canonical GPT-2 function is bytes_to_unicode() from transformers.models.gpt2.tokenization_gpt2; reversing it gives a unicode_char → byte map. Inline it to avoid coupling to the slow tokenizer class:

def gpt2_bytes_to_unicode():
    bs = (list(range(ord("!"), ord("~") + 1)) +
          list(range(ord("¡"), ord("¬") + 1)) +
          list(range(ord("®"), ord("ÿ") + 1)))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b); cs.append(256 + n); n += 1
    return dict(zip(bs, (chr(c) for c in cs)))

def _build_token_bytes(tok) -> list[bytes]:
    encoder = gpt2_bytes_to_unicode()
    byte_decoder = {c: b for b, c in encoder.items()}
    vocab_size = len(tok.get_vocab())
    out = []
    for i in range(vocab_size):
        s = tok.convert_ids_to_tokens(i, skip_special_tokens=False)
        try:
            raw = bytes(byte_decoder[c] for c in s)
        except KeyError:
            # Special tokens (<|endoftext|>, etc.) aren't in the byte map;
            # give them an impossible byte so the grammar rejects them.
            raw = b"\x00"
        out.append(raw)
    return out

For other tokenizers, the equivalent inversion depends on the vocabulary layout — SentencePiece uses sp.id_to_piece(i).replace("▁", " "), Llama's BPE uses a different byte-to-unicode table, etc. The grammar only needs the final bytes per id; how you recover them is tokenizer-specific.

grammar.initial_legal_tokens()

Debug helper returning the list of token ids legal at the initial grammar state. Useful for checking that set_token_vocabulary was set up correctly:

g = tl.Grammar(r"[0-9]+")
g.set_token_vocabulary(token_bytes)
legal = g.initial_legal_tokens()
print(f"{len(legal)} digit-starting tokens in the vocab")
# Should be 10+ for any reasonable GPT-2-style tokenizer (0..9 + multi-digit).

Properties

Property	Meaning
grammar.vocab_size	Number of tokens the grammar was built for (0 before set_token_vocabulary runs).
grammar.num_states	NFA node count — roughly proportional to regex complexity. Useful for sanity: [0-9]+ compiles to ~4 states, a full JSON-shape pattern compiles to ~15.

Using with Model.generate

Pass the grammar as the grammar= kwarg. All other kwargs (sampling, beam, stop sequences) compose normally:

# Greedy + grammar
out = model.generate(prompt, grammar=g, max_tokens=32, temperature=0.0)

# Sampling + grammar
out = model.generate(prompt, grammar=g, max_tokens=32,
                     temperature=0.8, top_k=50)

# Beam search + grammar (per-beam parse state tracked internally)
out = model.generate(prompt, grammar=g, max_tokens=32,
                     beam_size=4, length_penalty=0.7)

When generation stops

Generation under a grammar ends on any of:

1. max_tokens reached.
2. A stop sequence matches (if stop_sequences is set).
3. The grammar reaches a terminal state where no token is legal — i.e. the grammar has fully consumed and can't accept any continuation. This is the normal way constrained outputs end: (yes|no|maybe) emits "maybe" and the grammar rejects every next byte.
4. EOS token is sampled (if it passes the grammar mask).

The grammar does not require an accepting state to have been reached — partial outputs that are mid-derivation when max_tokens hits are returned as-is. If you need a guaranteed-complete output, make the grammar require a terminator ("}" for JSON, $ via \n for line-structured output) and set max_tokens comfortably larger than the longest valid derivation.

Per-beam grammar state

Under beam search, each beam maintains its own parse state that diverges as beams pick different next tokens. This is exactly what you want — a beam that committed to {"a and another that committed to {"b are tracking different grammar positions and should be masked accordingly.

Costs

• Grammar compile time — O(pattern length). Happens once per tl.Grammar(pattern). Complex patterns (JSON schema) compile in under a millisecond.
• Per-step mask cost — O(vocab_size × avg_token_bytes) worst case, amortised by a per-state mask cache. Repeat steps at the same parse state are O(1) after the first.
• Logit mask application — O(vocab_size) on the host (logits are already host-side after scoring).

No GPU kernel is added — the mask is pure host code. For vocab=50257, the mask cost is microseconds per step, negligible next to attention.

Common patterns

Digits only

g = tl.Grammar(r"[0-9]+")

Multiple choice (function-calling, classification)

Use the grammar_choice-style pattern:

g = tl.Grammar(r"(search|reply|escalate|refuse)")

Or programmatically:

choices = ["search", "reply", "escalate", "refuse"]
pattern = "(" + "|".join(re.escape(c) for c in choices) + ")"
g = tl.Grammar(pattern)

JSON with a fixed schema

# Minimal: object with one string field.
g = tl.Grammar(r'\{"name":"[a-z]+"\}')

# More flexible: allow spaces around punctuation.
g = tl.Grammar(r'\{\s*"name"\s*:\s*"[a-z ]+"\s*\}')

# With multiple fields — verbose, but a proper GBNF→FSA parser that
# handles recursive schema is tracked as follow-up work.
g = tl.Grammar(
    r'\{"name":"[a-z]+","age":[0-9]+,"city":"[a-z ]+"\}')

For full JSON-schema compliance (recursive objects, arrays, nested structures), a dedicated grammar_json_schema(schema_dict) helper is on the roadmap. Until then, hand-write the regex for your specific shape.

Constrained with a leading space (BPE gotcha)

Most GPT-2-family tokens start with a Ġ (encoded space) because the tokenizer prefers space-prefixed tokens. When you're generating at the start of a line after a prompt like "Answer:", the natural continuation is typically a space-prefixed token.

If the grammar disallows space at the start (e.g. [0-9]+), the mask rejects every space-prefixed token and the model has to fall through to non-space-prefixed digit tokens — which GPT-2 does have (single-digit tokens, short-number tokens) but they're lower-probability. The output might be less natural-sounding even though it satisfies the grammar.

Workaround: allow an optional leading space in the grammar:

g = tl.Grammar(r" ?[0-9]+")

Grammar + beam search

g = tl.Grammar(r"(Yes|No|Maybe)\.")
out = model.generate(
    prompt, grammar=g, max_tokens=8,
    beam_size=4, length_penalty=1.0,
)

Each beam maintains its own grammar state. Dead-end beams carry forward as "finished" so the global top-k can still rank them.

Error conditions

Error	Typical cause	Fix
ValueError: Grammar: unexpected trailing character at position X	Unbalanced parens or unterminated class	Count your (, ), [, ] and escape literals with \.
No tokens generated, output empty	The grammar rejects every token at the initial state	Check with grammar.initial_legal_tokens() — if empty, the vocab doesn't contain any byte sequence that can start a valid derivation. Typically fixed by allowing a leading space (see gotcha above).
Greedy under grammar produces a short output that doesn't reach the accepting state	Grammar reached a terminal dead-end before max_tokens	Normal behaviour — the grammar stopped generation because no next token is legal. Output is still a valid prefix.
vocab_size on the grammar doesn't match your model	Forgot to call set_token_vocabulary, or passed the wrong tokenizer's list	Ensure len(token_bytes) == model.config()["vocab_size"].

Implementation notes

NFA representation

The grammar compiles to a classical Thompson's NFA:

• MATCH nodes have a 256-bit byte-class bitmask and one out-edge.
• SPLIT nodes have two epsilon out-edges (used for |, *, +, ?).
• ACCEPT is the terminal state.

Per-step masking simulates each candidate token's bytes against the current active set of NFA nodes; tokens whose simulation leaves a non-empty active set are legal.

Mask caching

The per-state legal-token mask is cached keyed on the serialized active-node vector. Under greedy sampling the parse state usually oscillates between a small set of recurrent states (especially for simple grammars), so the cache hit rate is near 100 % after the first couple of steps.

Under beam search, each beam has its own state; caching is shared across beams because the key is the active-set content, not the beam identity.

What this isn't (yet)

• Full GBNF grammar file support — llama.cpp's GBNF dialect with recursive rule definitions isn't parsed yet. A grammar_gbnf(gbnf_str) compiler building on the same internal API is tracked in pending.md.
• JSON schema compiler — pass a Python schema dict and get a grammar back. Tracked as a follow-up building on the GBNF parser.
• Pushdown-automaton — the current NFA is regular (finite-state). True CFGs with unbounded nesting (deeply-nested JSON, balanced brackets) need a stack. The regex engine approximates by demanding an explicit nesting depth in the pattern.

See also

• Generation — the host API that accepts grammar=.
• tests/test_constrained_decoding.py — ready-made examples for digits / choice / JSON shapes.