| """Prefix grammar for generating one ActionCodec2Sparse chunk.""" | |
| from __future__ import annotations | |
| from collections.abc import Sequence | |
| import numpy as np | |
| from .vocabulary import SparseVocabulary | |
| SparseState = tuple[int, tuple[int, ...]] | |
| class SparseGrammar: | |
| """Legal next tokens of one ``horizon``-frame chunk over the active slots.""" | |
| def __init__(self, vocabulary: SparseVocabulary, horizon: int, active: Sequence[bool]): | |
| if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: | |
| raise ValueError("horizon must be a positive integer") | |
| layout, tables = vocabulary.layout, vocabulary.tables | |
| self.horizon, self.lookahead = horizon, vocabulary.lookahead | |
| self.vocab_size = vocabulary.vocab_size | |
| self.masks = np.zeros((len(tables), self.lookahead), dtype=np.int64) | |
| for token, masks in enumerate(tables.masks): | |
| self.masks[token, : len(masks)] = masks | |
| self.e_bit = 1 << layout.e_dim | |
| allowed = self.e_bit | sum(1 << d for d, flag in enumerate(active) if flag) | |
| self.in_scope = np.all((self.masks & ~allowed) == 0, axis=1) | |
| self.span = np.asarray(tables.span, dtype=np.int64) | |
| self.holds_e = (self.masks & self.e_bit) != 0 | |
| self.closes_first = self.holds_e[:, 0] | |
| self.absolute = sum( | |
| 1 << d for d, zero in enumerate(layout.zero) if zero is None and active[d] | |
| ) | |
| def initial(self) -> SparseState: | |
| return 0, (0,) * (self.horizon + self.lookahead) | |
| def front(self, state: SparseState) -> int: | |
| return state[0] | |
| def mask(self, state: SparseState) -> np.ndarray: | |
| """Boolean ``(V,)`` mask of tokens legal after ``state``.""" | |
| front, cover = state | |
| if front >= self.horizon: | |
| return np.zeros(len(self.masks), dtype=bool) | |
| window = np.asarray(cover[front : front + self.lookahead], dtype=np.int64) | |
| legal = self.in_scope & (front + self.span <= self.horizon) | |
| legal &= np.all((self.masks & window[None, :]) == 0, axis=1) | |
| closed = (window & self.e_bit) != 0 | |
| legal &= ~np.any((self.masks != 0) & closed[None, :], axis=1) | |
| last = self.horizon - 1 - front | |
| if last < self.lookahead: | |
| legal &= ~self.holds_e[:, last] | |
| if front == 0 and (cover[0] & self.absolute) != self.absolute: | |
| stated = (cover[0] | self.masks[:, 0]) & self.absolute | |
| legal &= ~self.closes_first | (stated == self.absolute) | |
| return legal | |
| def advance(self, state: SparseState, token: int) -> SparseState | None: | |
| """State after ``token``, or ``None`` when it is illegal.""" | |
| if not 0 <= token < len(self.masks) or not self.mask(state)[token]: | |
| return None | |
| front, cover = state | |
| updated = list(cover) | |
| for offset in range(self.lookahead): | |
| updated[front + offset] |= int(self.masks[token, offset]) | |
| while front < self.horizon and updated[front] & self.e_bit: | |
| front += 1 | |
| return front, tuple(updated) | |
| def is_complete(self, state: SparseState) -> bool: | |
| """Whether the chunk may end here (every active absolute slot stated).""" | |
| return (state[1][0] & self.absolute) == self.absolute | |
| def accepts(self, tokens: Sequence[int]) -> bool: | |
| state = self.initial() | |
| for token in tokens: | |
| state = self.advance(state, int(token)) | |
| if state is None: | |
| return False | |
| return self.is_complete(state) | |