Structured Decoding: Token Masks From Schema Automata

This witness implements the same finite object as the demo: a token-level automaton for a tiny ordered retrieval-call schema. The unconstrained decoder can choose an invalid token immediately. The constrained decoder is schema-valid. A semantic-mismatch profile stays schema-valid while choosing the wrong source for the hidden 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"