File size: 6,071 Bytes
0cf8154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """Schema-constrained decoding of the `<|spec|>` JSON.
The grammar accepts exactly the canonical compact form written by `zenith.format.spec_json`:
`{"sim":"<sim>"}` for meta sims, otherwise `{"sim":"<sim>","op":"<op>"` followed by zero or more
`,"<param>":<number>` (params of that op, no repeats, numbers of at most 12 characters) and `}`.
Ranges and exclusive parameters are left to engine validation.
"""
from __future__ import annotations
import re
import torch
from zenith.engines.base import Param, SpecError
from zenith.engines.registry import META_SIMS
from zenith.model.generate import generate
_NUM_FULL = re.compile(r"-?(0|[1-9]\d*)(\.\d+)?(e-?\d+)?")
_NUM_PARTIAL = re.compile(r"-?|-?(0|[1-9]\d*)(\.\d*)?|-?(0|[1-9]\d*)(\.\d+)?e-?\d*")
_NUM_CHARS = set("-0123456789.e")
MAX_NUMBER_CHARS = 12
PREFIX, COMPLETE = "prefix", "complete"
class _Partial(Exception):
"""The text ended inside a valid construction."""
class _Invalid(Exception):
"""The text cannot be extended into a valid spec."""
class SpecGrammar:
def __init__(self, schemas: dict[str, dict[str, list[Param]]]):
self.ops = {sim: {op: [p.name for p in params] for op, params in ops.items()} for sim, ops in schemas.items()}
self.sims = [*self.ops, *META_SIMS]
def state(self, text: str) -> str | None:
"""PREFIX, COMPLETE, or None if `text` is not the start of a valid spec."""
pos = 0
def lit(s: str) -> None:
nonlocal pos
rest = text[pos:]
if len(rest) < len(s):
if s.startswith(rest):
raise _Partial
raise _Invalid
if not rest.startswith(s):
raise _Invalid
pos += len(s)
def name(options) -> str:
nonlocal pos
rest = text[pos:]
end = rest.find('"')
if end < 0:
if any(o.startswith(rest) for o in options):
raise _Partial
raise _Invalid
if rest[:end] not in options:
raise _Invalid
pos += end + 1
return rest[:end]
def number() -> None:
nonlocal pos
start = pos
while pos < len(text) and text[pos] in _NUM_CHARS:
pos += 1
run = text[start:pos]
if len(run) > MAX_NUMBER_CHARS:
raise _Invalid
if pos == len(text):
# A run that is not yet a full number needs room to grow, or it could never be closed.
if _NUM_FULL.fullmatch(run) or (_NUM_PARTIAL.fullmatch(run) and len(run) < MAX_NUMBER_CHARS):
raise _Partial
raise _Invalid
if not _NUM_FULL.fullmatch(run):
raise _Invalid
try:
lit('{"sim":"')
sim = name(self.sims)
if sim in META_SIMS:
lit("}")
else:
lit(',"op":"')
params = self.ops[sim][name(list(self.ops[sim]))]
used: set[str] = set()
while True:
if pos == len(text):
raise _Partial
if text[pos] == "}":
pos += 1
break
remaining = [n for n in params if n not in used]
if not remaining:
raise _Invalid # only "}" can follow once every parameter is used
lit(',"')
p = name(remaining)
used.add(p)
lit(":")
number()
if pos != len(text):
raise _Invalid
return COMPLETE
except _Partial:
return PREFIX
except _Invalid:
return None
def is_prefix(self, text: str) -> bool:
return self.state(text) is not None
def is_complete(self, text: str) -> bool:
return self.state(text) == COMPLETE
class SpecConstraint:
"""`logits_hook` keeping only the `k` best-scoring tokens that extend a valid spec prefix."""
def __init__(self, tok, grammar: SpecGrammar, prompt_len: int, k: int):
self.tok, self.grammar, self.prompt_len, self.k = tok, grammar, prompt_len, k
self.pieces = [tok.decode([i]) for i in range(tok.vocab_size)]
def __call__(self, ids: list[int], logits: torch.Tensor) -> torch.Tensor:
text = self.tok.decode(ids[self.prompt_len :])
keep = []
for i in torch.argsort(logits[0], descending=True).tolist():
if logits[0, i] == float("-inf"):
break
if self.grammar.is_prefix(text + self.pieces[i]):
keep.append(i)
if len(keep) == self.k:
break
if not keep:
raise SpecError("no token can continue the spec")
masked = torch.full_like(logits, float("-inf"))
masked[0, keep] = logits[0, keep]
return masked
_GRAMMARS: dict[int, SpecGrammar] = {}
def _grammar() -> SpecGrammar:
from zenith.engines.registry import schemas
if 0 not in _GRAMMARS:
_GRAMMARS[0] = SpecGrammar(schemas())
return _GRAMMARS[0]
def decode_spec(model, tok, prompt_ids: list[int], temperature: float, top_k: int | None, max_tokens: int = 96) -> str:
"""Generate a schema-valid spec after `prompt_ids` (which ends with `<|spec|>`)."""
grammar = _grammar()
n = len(prompt_ids)
k = 1 if temperature == 0 else (top_k or 50)
device = next(model.parameters()).device
out = generate(model, torch.tensor([prompt_ids], device=device), max_tokens, temperature=temperature, top_k=top_k,
logits_hook=SpecConstraint(tok, grammar, n, k),
stop_fn=lambda ids: grammar.is_complete(tok.decode(ids[n:])))
text = tok.decode(out[0, n:].tolist())
if not grammar.is_complete(text):
raise SpecError(f"spec not finished within {max_tokens} tokens: {text!r}")
return text
|