Spaces:
Sleeping
Sleeping
| """Benchmark stop rules, applied identically to AR and speculative decoding. | |
| Same criteria as scripts/bench_speed.StopRule (chantiers/06-freins-vitesse.md, T3): GSM8K | |
| stops after the '#### <answer>' line or at the next few-shot question; HumanEval at the | |
| usual function-end sequences. | |
| """ | |
| from __future__ import annotations | |
| import re | |
| class StopRule: | |
| SEQS = {"gsm8k": ("\n\nQuestion:",), "humaneval": ("\ndef ", "\nclass ", "\nif __name__", "\nprint(", "\n#")} | |
| def __init__(self, tok, bench): | |
| self.tok, self.bench, self.text = tok, bench, "" | |
| def __call__(self, t): | |
| self.text += self.tok.decode([t]) | |
| tail = self.text[-48:] | |
| if self.bench == "gsm8k" and re.search(r"####[^\n]*\d[^\n]*\n", tail): | |
| return True | |
| return any(s in tail for s in self.SEQS.get(self.bench, ())) | |
| def make_stop(tok, kind: str | None) -> StopRule | None: | |
| """A fresh stop rule (they are stateful: one per generation), or None.""" | |
| if kind is None: | |
| return None | |
| if kind not in StopRule.SEQS: | |
| raise ValueError(f"unknown stop rule {kind!r}; available: {sorted(StopRule.SEQS)} or None") | |
| return StopRule(tok, kind) | |