| """Decoding (batched greedy, beam search, example-aware ranking) + CLI. |
| |
| python src/infer.py "five digit zip code, optionally dash and four more" |
| python src/infer.py -k 5 "an email address" # top 5 candidates |
| python src/infer.py --bits 8 "a uuid" # 8-bit quantized |
| python src/infer.py # interactive |
| """ |
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| import engine |
| import grammar |
| import tokenizer as tok |
| from model import GPT, Config |
|
|
| ROOT = Path(__file__).parents[1] |
| HERE = Path(__file__).parent |
| DEFAULT = HERE if (HERE / "model.safetensors").exists() else ROOT / "ckpt/best" |
| NEG = -1e9 |
|
|
|
|
| def load(path, bits=0): |
| """bits=0: as trained; 8 or 4: quantized linears (smaller, faster).""" |
| from model import quantize |
| path = Path(path) |
| c = Config(**json.loads((path / "config.json").read_text())) |
| m = GPT(c) |
| m.load_weights(str(path / "model.safetensors")) |
| if bits: |
| quantize(m, bits) |
| m.eval() |
| mx.eval(m.parameters()) |
| return m |
|
|
|
|
| def greedy(model, prompts, max_new=96): |
| """Batched greedy decode with left padding. prompts: list[list[int]].""" |
| B, L = len(prompts), max(map(len, prompts)) |
| pad = mx.array([L - len(p) for p in prompts]) |
| x = mx.array([[tok.PAD] * (L - len(p)) + p for p in prompts]) |
| pos = mx.maximum(mx.arange(L)[None] - pad[:, None], 0) |
| valid = mx.arange(L)[None] >= pad[:, None] |
| causal = mx.tril(mx.ones((L, L), dtype=mx.bool_)) |
| if model.c.prefix: |
| causal = mx.ones((L, L), dtype=mx.bool_) |
| mask = mx.where(causal[None, None] & valid[:, None, None, :], 0.0, NEG) |
| cache = [[None, None] for _ in model.blocks] |
| logits = model(x, pos, mask, cache)[:, -1] |
| out = [[] for _ in range(B)] |
| done = [False] * B |
| cur = mx.array([len(p) for p in prompts]) |
| kmask = mx.where(valid, 0.0, NEG) |
| for _ in range(max_new): |
| nxt = mx.argmax(logits, axis=-1) |
| ids = nxt.tolist() |
| for i, t in enumerate(ids): |
| if not done[i]: |
| if t == tok.EOS: |
| done[i] = True |
| else: |
| out[i].append(t) |
| if all(done): |
| break |
| kmask = mx.concatenate([kmask, mx.zeros((B, 1))], axis=1) |
| logits = model(nxt[:, None], cur[:, None], kmask[:, None, None, :], cache)[:, -1] |
| cur = cur + 1 |
| return [tok.decode(o) for o in out] |
|
|
|
|
| def beam(model, prompt, k=8, max_new=96, alpha=0.6): |
| """Beam search for one prompt -> list of (regex, score), best first.""" |
| x = mx.array([prompt]) |
| cache = [[None, None] for _ in model.blocks] |
| logp = mx.log(mx.softmax(model(x, None, None if model.c.prefix else "causal", cache)[:, -1].astype(mx.float32), axis=-1)) |
| T = len(prompt) |
| beams = [([], 0.0)] |
| finished = [] |
| for step in range(max_new): |
| lp = logp + mx.array([b[1] for b in beams])[:, None] |
| flat = lp.reshape(-1) |
| top = mx.argsort(-flat)[: 2 * k].tolist() |
| V = logp.shape[-1] |
| nb, src = [], [] |
| fl = flat.tolist() if len(top) else [] |
| for idx in top: |
| bi, t = divmod(idx, V) |
| s = fl[idx] |
| seq = beams[bi][0] |
| if t == tok.EOS: |
| finished.append((seq, s / ((5 + len(seq)) / 6) ** alpha)) |
| elif t >= tok.OFFSET: |
| nb.append((seq + [t], s)) |
| src.append(bi) |
| if len(nb) == k: |
| break |
| if not nb or (len(finished) >= k and max(f[1] for f in finished) > max(s for _, s in nb) / ((5 + step + 1) / 6) ** alpha): |
| break |
| idx = mx.array(src) |
| for c in cache: |
| c[0], c[1] = c[0][idx], c[1][idx] |
| beams = nb |
| x = mx.array([[s[-1]] for s, _ in beams]) |
| pos = mx.full((len(beams), 1), T + step) |
| logp = mx.log(mx.softmax(model(x, pos, None, cache)[:, -1].astype(mx.float32), axis=-1)) |
| finished.sort(key=lambda f: -f[1]) |
| seen, res = set(), [] |
| for seq, s in finished: |
| r = tok.decode(seq) |
| if r not in seen: |
| seen.add(r) |
| res.append((r, s)) |
| return res |
|
|
|
|
| def valid(regex): |
| try: |
| p, f = engine.parse_literal(regex) |
| except ValueError: |
| return False |
| return engine.compiles(p, f) |
|
|
|
|
| def best(model, en, k=8, use_hints=True): |
| """Beam search, then rank: compiles > matches examples in the request > in scope > model score.""" |
| import hints |
| cands = beam(model, tok.prompt(en), k=k) |
| ex = hints.examples(en) if use_hints else [] |
|
|
| def rank(c): |
| r = c[0] |
| try: |
| p, f = engine.parse_literal(r) |
| ok = engine.compiles(p, f) |
| except ValueError: |
| return (1, 0, 1, -c[1]) |
| hit = sum(engine.test(p, f, ex)) if ok and ex else 0 |
| scope = ok and grammar.in_scope(p, f) |
| return (not ok, -hit, not scope, -c[1]) |
| return sorted(cands, key=rank) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("text", nargs="*") |
| ap.add_argument("--ckpt", default=DEFAULT) |
| ap.add_argument("-k", type=int, default=1, help="show top-k candidates") |
| ap.add_argument("--beam", type=int, default=8) |
| ap.add_argument("--bits", type=int, default=0, help="quantize linears: 8 or 4") |
| a = ap.parse_args() |
| model = load(a.ckpt, a.bits) |
|
|
| def run(en): |
| for r, s in best(model, en, k=max(a.beam, a.k))[: a.k]: |
| print(r if a.k == 1 else f"{s:7.3f} {r}") |
|
|
| if a.text: |
| run(" ".join(a.text)) |
| return |
| for line in sys.stdin if not sys.stdin.isatty() else iter(lambda: input("› "), None): |
| if line.strip(): |
| run(line.strip()) |
|
|
|
|
| if __name__ == "__main__": |
| try: |
| main() |
| except (EOFError, KeyboardInterrupt): |
| pass |
|
|
