Spaces:
Sleeping
Sleeping
| """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)} | |
| 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] | |
| 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)", | |
| ] | |
| 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") | |