Structured Decoding: Token Masks From Schema Automata

This code implements the same automaton as the demo: a token-level automaton for a tiny retrieval-call schema with fields in a fixed order. The unconstrained decoder can choose an invalid token immediately. The constrained decoder produces schema-valid output. A semantic-mismatch profile stays schema-valid while choosing the wrong source for the task.

from collections import defaultdict
from math import exp

V = [
    "{", "}", ":", ",",
    '"tool"', '"source"', '"k"',
    '"retrieve"', '"lookup"', '"docs"', '"tickets"',
    "1", "2",
    '"DROP"', '"extra"', "true", "<eos>",
]

D = {
    0: {"{": 1},
    1: {'"tool"': 2},
    2: {":": 3},
    3: {'"retrieve"': 4, '"lookup"': 4},
    4: {",": 5},
    5: {'"source"': 6},
    6: {":": 7},
    7: {'"docs"': 8, '"tickets"': 8},
    8: {",": 9},
    9: {'"k"': 10},
    10: {":": 11},
    11: {"1": 12, "2": 12},
    12: {"}": 13},
    13: {"<eos>": 14},
}

ACCEPT = {14}

VALID_PROFILES = {"schema-friendly", "format-confused", "semantic-mismatch"}

def good_states():
    rev = defaultdict(set)
    for q, arcs in D.items():
        for tok, r in arcs.items():
            rev[r].add(q)
    good = set(ACCEPT)
    stack = list(ACCEPT)
    while stack:
        r = stack.pop()
        for q in rev[r]:
            if q not in good:
                good.add(q)
                stack.append(q)
    return good

GOOD = good_states()

def allowed(q):
    return [tok for tok in V if D.get(q, {}).get(tok) in GOOD]

def step(q, tok):
    return D.get(q, {}).get(tok)

def base_logits(q, profile):
    if profile not in VALID_PROFILES:
        raise ValueError(f"unknown profile: {profile}")

    z = {tok: -2.0 for tok in V}
    for tok in allowed(q):
        z[tok] = 1.0

    preferred = {3: '"retrieve"', 7: '"docs"', 11: "2"}
    if q in preferred:
        z[preferred[q]] = 2.0

    if profile == "format-confused":
        z['"DROP"'] = 5.0
        z['"extra"'] = 4.0

    if profile == "semantic-mismatch" and q == 7:
        z['"tickets"'] = 5.0
        z['"docs"'] = 1.0

    return z

def softmax(z, toks):
    m = max(z[t] for t in toks)
    weights = {t: exp(z[t] - m) for t in toks}
    total = sum(weights.values())
    return {t: weights.get(t, 0.0) / total for t in V}

def masked_probs(z, q):
    return softmax(z, allowed(q))

def rejection_mass(z, q):
    raw = softmax(z, V)
    return 1.0 - sum(raw[t] for t in allowed(q))

def decode(profile="format-confused", constrained=True, max_steps=32):
    q = 0
    toks = []
    states = [q]

    for _ in range(max_steps):
        z = base_logits(q, profile)
        p = masked_probs(z, q) if constrained else softmax(z, V)
        tok = max(V, key=lambda t: (p[t], -V.index(t)))
        toks.append(tok)
        q_next = step(q, tok)

        if q_next is None:
            return {"tokens": toks, "states": states, "accepted": False}

        q = q_next
        states.append(q)

        if q in ACCEPT:
            return {"tokens": toks, "states": states, "accepted": True}

    return {"tokens": toks, "states": states, "accepted": False}

def parse_source(tokens):
    i = tokens.index('"source"')
    return tokens[i + 2].strip('"')

assert allowed(0) == ["{"]
assert set(allowed(3)) == {'"retrieve"', '"lookup"'}
assert allowed(13) == ["<eos>"]

bad = decode("format-confused", constrained=False)
assert bad["accepted"] is False
assert bad["tokens"][0] == '"DROP"'

good = decode("format-confused", constrained=True)
assert good["accepted"] is True
assert good["tokens"] == [
    "{", '"tool"', ":", '"retrieve"', ",",
    '"source"', ":", '"docs"', ",",
    '"k"', ":", "2", "}", "<eos>",
]

z = base_logits(3, "format-confused")
p = masked_probs(z, 3)
assert p['"DROP"'] == 0.0
assert rejection_mass(z, 3) > 0.90

sem = decode("semantic-mismatch", constrained=True)
assert sem["accepted"] is True
assert parse_source(sem["tokens"]) == "tickets"

Mask, renormalize, then top-p

This standard-library Python code does the worked example above: the masked and renormalized distribution, both orders of the mask and top-$p$, and the second state.

def mask_and_renormalize(raw, mask):
    kept = [p * m for p, m in zip(raw, mask)]
    surviving = sum(kept)
    if surviving == 0:                       # any positive mass is renormalized
        raise ValueError("no allowed token has positive probability")
    return [p / surviving for p in kept], surviving

def top_p(dist, threshold):
    order = sorted(range(len(dist)), key=lambda i: (-dist[i], i))
    keep, mass = set(), 0.0
    for i in order:
        keep.add(i)
        mass += dist[i]
        if mass >= threshold:                # stop once the kept mass reaches p
            break
    return [x / mass if i in keep else 0.0 for i, x in enumerate(dist)]

TOKENS = ['"retrieve"', '"lookup"', '"DROP"', '"extra"', "true"]
RAW = [0.20, 0.10, 0.50, 0.15, 0.05]
MASK = [1, 1, 0, 0, 0]

constrained, surviving = mask_and_renormalize(RAW, MASK)
print("surviving mass:", round(surviving, 12), "constrained:", [round(x, 6) for x in constrained])

raw_first = top_p(RAW, 0.4)                  # keeps "DROP" alone
print("top-p first, then mask:", sum(x * m for x, m in zip(raw_first, MASK)))
print("mask first, then top-p:", top_p(constrained, 0.4))

source, _ = mask_and_renormalize([0.10, 0.60, 0.15, 0.10, 0.05], MASK)
print('"docs", "tickets":', [round(x, 6) for x in source[:2]])

assert mask_and_renormalize([1e-10, 1 - 1e-10], [1, 0])[0] == [1.0, 0.0]