Structured Outputs Under the Hood: How Constrained Decoding Works

python -m venv .venv
.venv/Scripts/python -m pip install pydantic

MERGED = [
    "Sure", " Here", " is", " the", " JSON", ":\n",
    '{"', '":', '": ', '": "', '", "', '"}', ", ",
    "name", "age", "Ada", " Lovelace", "Alan", " Turing", "Grace", " Hopper",
    "thirty", "-six", "41", "85", "36",
]
SINGLE = [chr(c) for c in range(32, 127)] + ["\n"]
 
VOCAB = list(dict.fromkeys(MERGED + SINGLE))
EOS = len(VOCAB)
VOCAB.append("<eos>")
TOKEN_ID = {text: i for i, text in enumerate(VOCAB)}
LONGEST = max(len(t) for t in VOCAB[:EOS])
 
 
def encode(text: str) -> list[int]:
    """Greedy longest match, like a BPE tokenizer's preferred split."""
    ids, i = [], 0
    while i < len(text):
        for size in range(min(LONGEST, len(text) - i), 0, -1):
            if text[i : i + size] in TOKEN_ID:
                ids.append(TOKEN_ID[text[i : i + size]])
                i += size
                break
        else:
            raise ValueError(f"no token for {text[i]!r}")
    return ids
 
 
def decode(ids: list[int]) -> str:
    return "".join(VOCAB[i] for i in ids if i != EOS)

import math
from collections import Counter
 
from tokenizer import EOS, VOCAB, decode, encode
 
ANSWER = 'Sure! Here is the JSON:\n{"name": "Ada Lovelace", "age": "thirty-six"}'
HABITS = [
    'Here is the JSON:\n{"name": "Alan Turing", "age": 41}',
    "{'name': 'Grace Hopper', 'age': 85}",
]
 
 
class ToyLM:
    """A stand-in model with fixed habits. It wants to say ANSWER. When the text so far
    leaves that script, it realigns on the longest suffix it has seen in any document."""
 
    def __init__(self, answer=ANSWER, habits=HABITS):
        self.docs = [answer] + habits
        self.prior = Counter(tok for doc in self.docs for tok in encode(doc))
 
    def continuations(self, text: str) -> list[tuple[int, str]]:
        for k in range(len(text), 0, -1):
            found = [(d, doc[i + k :]) for d, doc in enumerate(self.docs)
                     for i in range(len(doc) - k + 1) if doc[i : i + k] == text[-k:]]
            if found:
                return found
        return [(d, doc) for d, doc in enumerate(self.docs)]
 
    def logits(self, prompt: str, out: list[int]) -> list[float]:
        # The prompt is ignored: this stand-in decided what it wants to say in advance.
        total = sum(self.prior.values()) + len(VOCAB)
        scores = [math.log((self.prior[tok] + 1) / total) for tok in range(len(VOCAB))]
        for doc, rest in self.continuations(decode(out)):
            bonus = 12.0 if doc == 0 else 8.0
            for tok, piece in enumerate(VOCAB[:EOS]):
                if rest.startswith(piece):
                    scores[tok] = max(scores[tok], bonus + len(piece) * 0.1)
            if rest == "":
                scores[EOS] = max(scores[EOS], bonus)
        return scores

from dataclasses import dataclass
 
from tokenizer import VOCAB
 
DIGITS = "0123456789"
 
 
@dataclass(frozen=True)
class Lit:
    text: str
    start = 0
 
    def step(self, s, ch):
        return s + 1 if s < len(self.text) and self.text[s] == ch else None
 
    def done(self, s):
        return s == len(self.text)
 
 
@dataclass(frozen=True)
class Space:
    """At most one space: an unbounded run lets a model pad with whitespace until max_tokens."""
 
    start = 0
 
    def step(self, s, ch):
        return 1 if s == 0 and ch == " " else None
 
    def done(self, s):
        return True
 
 
@dataclass(frozen=True)
class Str:
    start = "open"
 
    def step(self, s, ch):
        if s == "open":
            return "body" if ch == '"' else None
        if s == "body":
            if ch == '"':
                return "closed"
            return "body" if ch >= " " and ch != "\\" else None
        return None
 
    def done(self, s):
        return s == "closed"
 
 
@dataclass(frozen=True)
class Int:
    start = "sign"
 
    def step(self, s, ch):
        if s == "sign" and ch == "-":
            return "minus"
        if s in ("sign", "minus"):
            return "zero" if ch == "0" else "digits" if ch in DIGITS else None
        return "digits" if s == "digits" and ch in DIGITS else None
 
    def done(self, s):
        return s in ("zero", "digits")
 
 
SEGMENTS = {"string": Str, "integer": Int}
 
 
def compile_schema(schema: dict) -> tuple[list, list[str]]:
    """A flat object schema to a list of segments, plus the keywords it could not enforce."""
    segs, ignored = [Lit("{"), Space()], []
    for i, (key, prop) in enumerate(schema["properties"].items()):
        if i:
            segs += [Lit(","), Space()]
        segs += [Lit(f'"{key}"'), Space(), Lit(":"), Space(), SEGMENTS[prop["type"]]()]
        ignored += [f"{key}.{kw}" for kw in prop if kw not in ("type", "title")]
    return segs + [Space(), Lit("}")], ignored
 
 
class JsonMachine:
    """Character-level automaton. A state is (segment index, position inside it)."""
 
    def __init__(self, schema: dict):
        self.segs, self.ignored = compile_schema(schema)
        self.start = (0, self.segs[0].start)
 
    def step(self, state, ch):
        i, s = state
        while i < len(self.segs):
            nxt = self.segs[i].step(s, ch)
            if nxt is not None:
                return (i, nxt)
            if not self.segs[i].done(s):
                return None
            i += 1
            s = self.segs[i].start if i < len(self.segs) else None
        return None
 
    def accepting(self, state):
        i, s = state
        return self.segs[i].done(s) and all(seg.done(seg.start) for seg in self.segs[i + 1 :])
 
    def walk(self, state, text):
        for ch in text:
            state = self.step(state, ch)
            if state is None:
                return None
        return state
 
    def allowed(self, state, vocab=VOCAB) -> dict[int, tuple]:
        """Token id to the state after it. Scans the whole vocabulary: the naive way.
        The last vocabulary entry is end of sequence."""
        nxt = {}
        for tok, text in enumerate(vocab[:-1]):
            after = self.walk(state, text)
            if after is not None:
                nxt[tok] = after
        if self.accepting(state):
            nxt[len(vocab) - 1] = state
        return nxt

import math
 
from tokenizer import EOS, VOCAB, decode
 
 
def generate(lm, prompt, machine=None, index=None, max_tokens=40, trace=False):
    """Greedy decoding. With a machine, every step masks the tokens the grammar rejects."""
    out, state = [], machine.start if machine else None
    for _ in range(max_tokens):
        raw = lm.logits(prompt, out)
        logits = raw
        if machine:
            allowed = index[state] if index is not None else machine.allowed(state)
            logits = [x if tok in allowed else -math.inf for tok, x in enumerate(raw)]
        tok = max(range(len(logits)), key=logits.__getitem__)
        if trace:
            wanted = max(range(len(raw)), key=raw.__getitem__)
            note = f"  wanted {VOCAB[wanted]!r}" if wanted != tok else ""
            count = len(allowed) if machine else "-"
            print(f"{len(out):>2}  {VOCAB[tok]!r:<13} allowed={count:<4}{note}")
        if tok == EOS:
            return decode(out), "stop", out
        out.append(tok)
        if machine:
            state = allowed[tok]
    return decode(out), "length", out

import json
 
from pydantic import BaseModel
 
from decode import generate
from grammar import JsonMachine
from index import build_index
from lm import ToyLM
from tokenizer import VOCAB, encode
 
 
class Person(BaseModel):
    name: str
    age: int
 
 
lm = ToyLM()
prompt = "Extract name and age as JSON: Ada Lovelace was thirty-six when she died."
 
text, finish, _ = generate(lm, prompt)
print(f"unconstrained, finish={finish}\n{text}")
try:
    json.loads(text)
except json.JSONDecodeError as err:
    print(f"json.loads: {err}\n")
 
machine = JsonMachine(Person.model_json_schema())
index = build_index(machine)
print(f"index: {len(index)} states, vocabulary of {len(VOCAB)}\n")
text, finish, ids = generate(lm, prompt, machine, index, trace=True)
print(f"\nconstrained, finish={finish}\n{text}")
print("json.loads:", json.loads(text))
print("generated:", [VOCAB[i] for i in ids[:4]])
print("canonical:", [VOCAB[i] for i in encode(text)[:4]])

$ python demo.py
unconstrained, finish=stop
Sure! Here is the JSON:
{"name": "Ada Lovelace", "age": "thirty-six"}
json.loads: Expecting value: line 1 column 1 (char 0)
 
index: 29 states, vocabulary of 123
 
 0  '{'           allowed=2     wanted 'Sure'
 1  '"'           allowed=2
 2  'name'        allowed=2
 3  '": "'        allowed=4
 4  'Ada'         allowed=115
 5  ' Lovelace'   allowed=115
 6  '", "'        allowed=115
 7  'age'         allowed=2
 8  '": '         allowed=3     wanted '": "'
 9  '41'          allowed=14    wanted '"'
10  '}'           allowed=15
11  '<eos>'       allowed=1
 
constrained, finish=stop
{"name": "Ada Lovelace", "age": 41}
json.loads: {'name': 'Ada Lovelace', 'age': 41}
generated: ['{', '"', 'name', '": "']
canonical: ['{"', 'name', '": "', 'Ada']

from grammar import JsonMachine
from tokenizer import VOCAB
 
 
def build_index(machine: JsonMachine, vocab=VOCAB) -> dict[tuple, dict[int, tuple]]:
    """Every state reachable at a token boundary, mapped to {allowed token: next state}."""
    index, todo = {}, [machine.start]
    while todo:
        state = todo.pop()
        if state not in index:
            index[state] = machine.allowed(state, vocab)
            todo += [s for s in index[state].values() if s not in index]
    return index

import json
import sys
import time
 
from pydantic import BaseModel
 
from grammar import JsonMachine
from index import build_index
 
 
class Person(BaseModel):
    name: str
    age: int
 
 
vocab = json.load(open(sys.argv[1], encoding="utf8")) + ["<eos>"]
machine = JsonMachine(Person.model_json_schema())
 
t0 = time.perf_counter()
index = build_index(machine, vocab)
build = time.perf_counter() - t0
 
states = list(index)
t0 = time.perf_counter()
for state in states:
    machine.allowed(state, vocab)
scan = (time.perf_counter() - t0) / len(states)
 
t0 = time.perf_counter()
for _ in range(1000):
    for state in states:
        index[state]
lookup = (time.perf_counter() - t0) / (1000 * len(states))
 
sizes = sorted((len(v), s) for s, v in index.items())
print(f"vocabulary: {len(vocab):,} tokens")
print(f"states at token boundaries: {len(index)}")
print(f"index build: {build:.2f} s, {sum(len(v) for v in index.values()):,} entries")
print(f"naive mask, one step: {scan * 1000:.1f} ms")
print(f"index lookup, one step: {lookup * 1e9:.0f} ns")
print(f"smallest allowed set: {sizes[0][0]} tokens at {sizes[0][1]}")
print(f"largest allowed set: {sizes[-1][0]:,} tokens at {sizes[-1][1]}")

$ python bench.py gpt2_vocab.json
vocabulary: 50,258 tokens
states at token boundaries: 29
index build: 1.51 s, 54,700 entries
naive mask, one step: 47.4 ms
index lookup, one step: 97 ns
smallest allowed set: 1 tokens at (1, 1)
largest allowed set: 50,122 tokens at (6, 'body')

import re
from dataclasses import dataclass
 
from pydantic import BaseModel, ConfigDict, Field, ValidationError, ValidationInfo, model_validator
 
UNITS = "zero one two three four five six seven eight nine".split()
TENS = {"twenty": 20, "thirty": 30, "forty": 40, "fifty": 50, "sixty": 60, "seventy": 70, "eighty": 80, "ninety": 90}
 
 
def numbers_in(text: str) -> set[int]:
    """Digits and simple English number words (thirty-six) that appear in the text."""
    found = {int(d) for d in re.findall(r"\d+", text)}
    for tens, units in re.findall(r"\b(" + "|".join(TENS) + r")(?:-(\w+))?", text.lower()):
        found.add(TENS[tens] + (UNITS.index(units) if units in UNITS else 0))
    return found
 
 
class Person(BaseModel):
    model_config = ConfigDict(extra="forbid")
 
    name: str = Field(min_length=1)
    age: int = Field(ge=0, le=130)
 
    @model_validator(mode="after")
    def grounded_in_source(self, info: ValidationInfo) -> "Person":
        source = (info.context or {}).get("source", "")
        if self.name not in source:
            raise ValueError(f"name {self.name!r} does not appear in the source text")
        if self.age not in numbers_in(source):
            raise ValueError(f"age {self.age} is not stated in the source text")
        return self
 
 
@dataclass
class Failed:
    reason: str
    attempts: list[str]
 
 
def extract(source: str, generate, max_attempts: int = 4, budget: int = 8):
    prompt, seen, log = f"Extract name and age as JSON: {source}", set(), []
    for attempt in range(1, max_attempts + 1):
        raw, finish = generate(prompt, max_tokens=budget)
        log.append(raw)
        if finish == "length":
            print(f"attempt {attempt}: cut off at {budget} tokens: {raw!r}")
            budget *= 2
            continue
        if raw in seen:
            print(f"attempt {attempt}: same output as before, giving up")
            return Failed("repeated output", log)
        seen.add(raw)
        try:
            person = Person.model_validate_json(raw, context={"source": source})
        except ValidationError as err:
            problems = "; ".join(e["msg"] for e in err.errors())
            print(f"attempt {attempt}: {raw} rejected: {problems}")
            prompt += f"\nYour previous answer {raw} was rejected: {problems}. Fix it."
            continue
        print(f"attempt {attempt}: accepted {person!r}")
        return person
    return Failed("out of attempts", log)

from pydantic import ValidationError
 
from decode import generate
from extract import Person, extract
from grammar import JsonMachine
from index import build_index
from lm import ToyLM
 
source = "Ada Lovelace was thirty-six when she died."
machine = JsonMachine(Person.model_json_schema())
index = build_index(machine)
print("grammar could not enforce:", machine.ignored)
 
for raw in [
    'Sure! Here is the JSON:\n{"name": "Ada Lovelace", "age": "thirty-six"}',
    '{"name": "Ada Lovelace", "age": "thirty-six"}',
    '{"name": "", "age": 999}',
]:
    try:
        Person.model_validate_json(raw, context={"source": source})
    except ValidationError as err:
        print(" | ".join(e["msg"] for e in err.errors()))
 
 
def constrained(lm):
    def call(prompt, max_tokens):
        text, finish, _ = generate(lm, prompt, machine, index, max_tokens=max_tokens)
        return text, finish
 
    return call
 
 
print("\n# the toy model")
print(extract(source, constrained(ToyLM())))
 
print("\n# a toy that wants to write the age as digits")
print(repr(extract(source, constrained(ToyLM(answer='{"name": "Ada Lovelace", "age": 36}')))))

$ python run_extract.py
grammar could not enforce: ['name.minLength', 'age.maximum', 'age.minimum']
Invalid JSON: expected value at line 1 column 1
Input should be a valid integer, unable to parse string as an integer
String should have at least 1 character | Input should be less than or equal to 130
 
# the toy model
attempt 1: cut off at 8 tokens: '{"name": "Ada Lovelace", "age'
attempt 2: {"name": "Ada Lovelace", "age": 41} rejected: Value error, age 41 is not stated in the source text
attempt 3: same output as before, giving up
Failed(reason='repeated output', attempts=['{"name": "Ada Lovelace", "age', '{"name": "Ada Lovelace", "age": 41}', '{"name": "Ada Lovelace", "age": 41}'])
 
# a toy that wants to write the age as digits
attempt 1: cut off at 8 tokens: '{"name": "Ada Lovelace", "age": '
attempt 2: accepted Person(name='Ada Lovelace', age=36)
Person(name='Ada Lovelace', age=36)

{"name": "Marie Curie', 'age': 'Marie Curie', 'age': 'Marie Curie', 'age': 'Marie Curie', 'age'