| """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): |
| |
| 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 |
| 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 |
|
|
