Spaces:
Sleeping
Sleeping
File size: 3,981 Bytes
371d90c | 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 | """Guided decoding: lambda = 0 is greedy, lookahead is exact, regex outputs always match."""
from __future__ import annotations
import random
import pytest
import regex
import torch
from conftest import ARCHS
from decoding.common import KV
from decoding.guided import Constraint, rollout, run_classifier, run_regex
BASE = {"max_new_tokens": 16, "mode": "greedy", "temperature": 1.0, "seed": 0, "rep": 1.0}
class FakeScorer:
"""Scores texts by how many times a marker substring occurs (deterministic, no model)."""
def __init__(self, marker: str):
self.marker = marker
def __call__(self, texts, _target):
return [float(t.count(self.marker)) for t in texts]
def scorers(marker="e"):
return {"sentiment": FakeScorer(marker), "topic": FakeScorer(marker)}
@pytest.mark.parametrize("arch", ARCHS)
@pytest.mark.parametrize("lookahead", [0, 2])
def test_lambda_zero_is_greedy(tiny_family, arch, lookahead):
fam = tiny_family(arch)
P = {**BASE, "prompt": "Write a sentence.", "attribute": "sentiment", "target": "positive", "lam": 0.0,
"k": 8, "lookahead": lookahead}
res = run_classifier(fam, P, scorers())
assert res["runs"]["method"]["ids"] == res["runs"]["baseline"]["ids"]
assert res["summary"]["changed"] == 0
def test_large_lambda_steers_toward_marker(tiny_family):
fam = tiny_family("llama")
P = {**BASE, "prompt": "Write a sentence.", "attribute": "topic", "topic": "x", "lam": 50.0, "k": 20,
"lookahead": 0}
res = run_classifier(fam, P, scorers("e"))
traj = res["trajectory"]
assert res["summary"]["changed"] > 0
assert traj["method"][-1] >= traj["baseline"][-1]
@pytest.mark.parametrize("arch", ARCHS)
def test_batched_rollout_equals_sequential(tiny_family, arch):
fam = tiny_family(arch)
prefix = [random.Random(1).randrange(3, 1000) for _ in range(9)]
kv = KV(fam.main)
kv.run(prefix)
cand = torch.tensor([5, 17, 99, 250])
batched = rollout(kv, cand, 4)
assert kv.toks == prefix # the original cache is untouched
for c, cont in zip(cand.tolist(), batched):
seq, ref = prefix + [c], KV(fam.main)
expect = []
for _ in range(4):
t = int(ref.run(seq).logits[0, -1].argmax())
expect.append(t)
seq = seq + [t]
assert cont == expect
PATTERNS = [
r"\d{4}-\d{2}-\d{2}",
r"(A|B|C|D)",
r'\{"name": "[a-z]{1,8}", "age": [0-9]{1,3}\}',
r"[a-z]{1,8}@[a-z]{1,8}\.(com|org)",
]
@pytest.mark.parametrize("arch", ARCHS)
@pytest.mark.parametrize("pattern", PATTERNS)
def test_regex_output_fullmatches(tiny_family, arch, pattern):
fam = tiny_family(arch)
P = {**BASE, "max_new_tokens": 40, "prompt": "Answer:", "pattern": pattern, "top_k": 32}
res = run_regex(fam, P)
meth = res["runs"]["method"]
assert res["summary"]["valid"]["method"], (pattern, meth["text"], meth["finish"])
assert regex.fullmatch(pattern, meth["text"])
for st in meth["steps"]:
if st["id"] in fam.stop_ids:
assert st["x"]["complete"] # end-of-sequence only after a complete match
def test_unbounded_pattern_output_stays_a_viable_prefix(tiny_family):
"""With `+` a random model may never close the pattern; the output must still be a valid prefix."""
fam = tiny_family("llama")
pattern = r"[a-z]+@[a-z]+\.(com|org)"
res = run_regex(fam, {**BASE, "max_new_tokens": 12, "prompt": "Answer:", "pattern": pattern, "top_k": 32})
text = res["runs"]["method"]["text"]
assert regex.fullmatch(pattern, text, partial=True)
assert res["summary"]["valid"]["method"] or res["runs"]["method"]["finish"] == "length"
def test_constraint_rejects_bad_patterns():
with pytest.raises(ValueError):
Constraint("(unclosed")
with pytest.raises(ValueError):
Constraint("a" * 400)
con = Constraint(r"\d{3}")
assert con.viable("12") and not con.viable("1a") and con.complete("123") and not con.complete("12")
|