Spaces:
Sleeping
Sleeping
| """Grammar-constrained LogitsProcessor (Python port of src/logits.js). | |
| At each generation step: | |
| 1. Decode the generated suffix back to text. | |
| 2. Advance the grammar NFA by that text. | |
| 3. For every candidate token id, check whether appending its decoded text | |
| keeps the NFA alive; mask losers to -inf (via the shared LegalCache). | |
| 4. EOS is allowed only once the NFA has reached an accept state. | |
| Per-token decoding can disagree with BPE sequence-decoding in edge cases | |
| (merged punctuation, etc.); for the acrostic patterns we care about this | |
| approximation is fine. | |
| """ | |
| from __future__ import annotations | |
| import time | |
| from transformers import LogitsProcessor | |
| from masking import LegalCache | |
| def build_token_text_table(tokenizer, vocab_size): | |
| """One-shot build of tokenId -> text, using per-token decode. Special tokens | |
| decode to '' under skip_special_tokens=True, which we treat as | |
| "disallowed" (empty string).""" | |
| texts = tokenizer.batch_decode( | |
| [[i] for i in range(vocab_size)], skip_special_tokens=True | |
| ) | |
| return [t if isinstance(t, str) else "" for t in texts] | |
| class GrammarLogitsProcessor(LogitsProcessor): | |
| def __init__(self, grammar, tokenizer, token_text, eos_token_ids=(), legal_cache=None): | |
| super().__init__() | |
| self.grammar = grammar | |
| self.tokenizer = tokenizer | |
| self.token_text = token_text | |
| self.cache = legal_cache or LegalCache(grammar, token_text, eos_token_ids) | |
| self.prompt_length = None | |
| self.stats = _fresh_stats() | |
| def reset(self): | |
| self.prompt_length = None | |
| self.stats = _fresh_stats() | |
| def __call__(self, input_ids, scores): | |
| t_entry = time.perf_counter() | |
| ids = input_ids[0] | |
| if self.prompt_length is None: | |
| self.prompt_length = ids.shape[0] | |
| generated = ids[self.prompt_length:].tolist() | |
| text = ( | |
| self.tokenizer.decode(generated, skip_special_tokens=True) | |
| if generated | |
| else "" | |
| ) | |
| state = self.grammar.advance(self.grammar.initial, text) | |
| data = scores[0] | |
| if state == -1: | |
| # Already violated; nothing useful to do without rewinding. Let the | |
| # original logits through so generation at least terminates. | |
| self._record(time.perf_counter() - t_entry, -1) | |
| return scores | |
| illegal = self.cache.illegal_tensor(state) | |
| data[illegal.to(data.device)] = float("-inf") | |
| self._record(time.perf_counter() - t_entry, int((~illegal).sum().item())) | |
| return scores | |
| def _record(self, dt, survivors): | |
| st = self.stats | |
| st["calls"] += 1 | |
| st["total_ms"] += dt * 1000.0 | |
| st["per_step"].append({"ms": dt * 1000.0, "survivors": survivors}) | |
| def _fresh_stats(): | |
| return {"calls": 0, "total_ms": 0.0, "per_step": []} | |