| """Constrained-generation grammars for the two ActionCodec2 BPE schemes.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
| from numbers import Integral |
|
|
| import numpy as np |
|
|
| from ._validation import active_dimensions |
| from .bpe import BPESerializer, SetBPESerializer |
| from .cache import BoundedCache, cached_prefix_state |
| from .native_backend import native_kernel |
|
|
|
|
| def _coverage_metadata(serializer, dims): |
| |
| cache = getattr(serializer, "_coverage_metadata", None) |
| if cache is None: |
| cache = serializer._coverage_metadata = BoundedCache(32) |
| metadata = cache.get(dims) |
| if metadata is not None: |
| return metadata |
| metadata = {"compiled": None} |
| if serializer.scheme == "set_bpe": |
| full = sum(1 << d for d in dims) |
| metadata["roots"] = np.asarray( |
| [serializer.spatial.anchor[serializer.root[t]] for t in range(serializer.vocab_size)], |
| dtype=np.int64, |
| ) |
| patterns = tuple( |
| tuple((offset, int(serializer.spatial.support[token])) for offset, token in pattern) |
| for pattern in serializer.patterns |
| ) |
| metadata["patterns"] = patterns |
| metadata["scope"] = np.asarray( |
| [all(support & ~full == 0 for _, support in pattern) for pattern in patterns], |
| dtype=bool, |
| ) |
| else: |
| phase = np.zeros((len(dims), serializer.vocab_size), dtype=bool) |
| for index in range(len(dims)): |
| for token, expansion in enumerate(serializer.vocabulary.token_dims): |
| phase[index, token] = all( |
| d == dims[(index + offset) % len(dims)] for offset, d in enumerate(expansion) |
| ) |
| metadata["phase"] = phase |
| metadata["lengths"] = np.asarray(serializer.token_cell_counts, dtype=np.int64) |
| cache[dims] = metadata |
| return metadata |
|
|
|
|
| class _PrefixGrammar: |
| |
|
|
| def _state(self, tokens: Sequence[int]) -> tuple[int, ...] | None: |
| checked = [] |
| for token in tokens: |
| if isinstance(token, bool) or not isinstance(token, Integral): |
| return None |
| checked.append(int(token)) |
| return cached_prefix_state( |
| tuple(checked), self._state_cache, self._empty_state, self._advance |
| ) |
|
|
| def _next_token_mask_from_state(self, state: tuple[int, ...]) -> np.ndarray: |
| |
| cached = self._mask_cache.get(state) |
| if cached is None: |
| compiled = self._compiled() |
| cached = compiled.mask(state) if compiled is not None else self._python_mask(state) |
| self._mask_cache[state] = cached |
| return cached.copy() |
|
|
| def next_token_mask(self, tokens: Sequence[int]) -> np.ndarray: |
| """Return a boolean vocabulary mask for the next generated token.""" |
| state = self._state(tokens) |
| if state is None: |
| raise ValueError("token prefix is not a legal ActionCodec2 stream") |
| return self._next_token_mask_from_state(state) |
|
|
| def is_complete(self, tokens: Sequence[int]) -> bool: |
| """Whether ``tokens`` is exactly one complete trajectory of this horizon.""" |
| state = self._state(tokens) |
| return state is not None and self._state_is_complete(state) |
|
|
|
|
| class ActionCodec2Grammar(_PrefixGrammar): |
| 'Prefix grammar for one exact, fixed-horizon rooted token stream.' |
|
|
| def __init__( |
| self, |
| serializer: SetBPESerializer, |
| *, |
| horizon: int, |
| active_dims: Sequence[int] | None = None, |
| ) -> None: |
| if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: |
| raise ValueError("horizon must be a positive integer") |
| self.horizon = horizon |
| self.vocab_size = serializer.vocab_size |
| dims = active_dimensions(active_dims, serializer.spatial.dimension) |
| self._full = sum(1 << dimension for dimension in dims) |
| self._metadata = _coverage_metadata(serializer, dims) |
| self._root_anchor = self._metadata["roots"] |
| self._patterns = self._metadata["patterns"] |
| self._in_scope = self._metadata["scope"] |
| self._mask_cache = BoundedCache(64) |
| self._state_cache = BoundedCache(256) |
| self._empty_state = self._initial_state() |
| self._native = None |
| self._compiled() |
|
|
| def _compiled(self): |
| constructor = native_kernel("SetGrammar") |
| if constructor is None: |
| return None |
| if self._native is None: |
| prototype = self._metadata["compiled"] |
| if prototype is None: |
| prototype = self._metadata["compiled"] = constructor( |
| 1, self._full, self._root_anchor.tolist(), self._patterns |
| ) |
| self._native = prototype.clone(self.horizon) |
| return self._native |
|
|
| def _advance(self, state: tuple[int, ...], token: int) -> tuple[int, ...] | None: |
| compiled = self._compiled() |
| if compiled is not None: |
| return compiled.advance(state, token) |
| if not 0 <= token < self.vocab_size or not self._in_scope[token]: |
| return None |
| coverage = list(state) |
| front = next( |
| (index for index, mask in enumerate(coverage) if mask != self._full), self.horizon |
| ) |
| if front == self.horizon: |
| return None |
| uncovered = self._full & ~coverage[front] |
| expected_anchor = int(uncovered & -uncovered).bit_length() - 1 |
| if int(self._root_anchor[token]) != expected_anchor: |
| return None |
| for offset, support in self._patterns[token]: |
| step = front + offset |
| if step >= self.horizon or coverage[step] & support: |
| return None |
| coverage[step] |= support |
| return tuple(coverage) |
|
|
| def _initial_state(self) -> tuple[int, ...]: |
| |
| return (0,) * self.horizon |
|
|
| def _covered_steps(self, state: tuple[int, ...]) -> int: |
| return next((i for i, mask in enumerate(state) if mask != self._full), self.horizon) |
|
|
| def _state_is_complete(self, state: tuple[int, ...]) -> bool: |
| |
| return all(mask == self._full for mask in state) |
|
|
| def _python_mask(self, state: tuple[int, ...]) -> np.ndarray: |
| if self._state_is_complete(state): |
| return np.zeros(self.vocab_size, dtype=bool) |
| front = next(index for index, mask in enumerate(state) if mask != self._full) |
| uncovered = self._full & ~state[front] |
| expected_anchor = int(uncovered & -uncovered).bit_length() - 1 |
| mask = (self._root_anchor == expected_anchor) & self._in_scope |
| for token in np.flatnonzero(mask): |
| for offset, support in self._patterns[int(token)]: |
| step = front + offset |
| if step >= self.horizon or state[step] & support: |
| mask[token] = False |
| break |
| return mask |
|
|
|
|
| class ActionCodec2BPEGrammar(_PrefixGrammar): |
| 'Prefix grammar for one exact, fixed-horizon plain-BPE token stream.' |
|
|
| def __init__( |
| self, |
| serializer: BPESerializer, |
| *, |
| horizon: int, |
| active_dims: Sequence[int] | None = None, |
| ) -> None: |
| if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: |
| raise ValueError("horizon must be a positive integer") |
| self.horizon = horizon |
| self.vocab_size = serializer.vocab_size |
| dims = active_dimensions(active_dims, serializer.vocabulary.dimension) |
| self._width = len(dims) |
| self._cells = horizon * self._width |
| self._metadata = _coverage_metadata(serializer, dims) |
| self._token_cells = self._metadata["lengths"] |
| self._phase = self._metadata["phase"] |
| self._mask_cache = BoundedCache(64) |
| self._state_cache = BoundedCache(256) |
| self._empty_state = self._initial_state() |
|
|
| self._active_dims = dims |
| self._native = None |
| self._compiled() |
|
|
| def _compiled(self): |
| constructor = native_kernel("PlainGrammar") |
| if constructor is None: |
| return None |
| if self._native is None: |
| prototype = self._metadata["compiled"] |
| if prototype is None: |
| prototype = self._metadata["compiled"] = constructor( |
| 1, self._token_cells.tolist(), self._active_dims, self._phase |
| ) |
| self._native = prototype.clone(self.horizon) |
| return self._native |
|
|
| def _initial_state(self) -> tuple[int, ...]: |
| return (0,) |
|
|
| def _covered_steps(self, state: tuple[int, ...]) -> int: |
| return state[0] // self._width |
|
|
| def _state_is_complete(self, state: tuple[int, ...]) -> bool: |
| return state[0] == self._cells |
|
|
| def _advance(self, state: tuple[int, ...], token: int) -> tuple[int, ...] | None: |
| compiled = self._compiled() |
| if compiled is not None: |
| return compiled.advance(state, token) |
| if not 0 <= token < self.vocab_size: |
| return None |
| covered = state[0] |
| length = int(self._token_cells[token]) |
| if covered + length > self._cells: |
| return None |
| if not self._phase[covered % self._width, token]: |
| return None |
| return (covered + length,) |
|
|
| def _python_mask(self, state: tuple[int, ...]) -> np.ndarray: |
| covered = state[0] |
| if covered == self._cells: |
| return np.zeros(self.vocab_size, dtype=bool) |
| return self._phase[covered % self._width] & (self._token_cells <= self._cells - covered) |
|
|
