| """Fitted ActionCodec2Sparse token layer: spatial Set-BPE per sparse frame plus temporal rules.""" |
|
|
| from __future__ import annotations |
|
|
| import heapq |
| import json |
| from collections.abc import Sequence |
| from numbers import Integral |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from ..native_backend import native_kernel, resolve_backend |
| from ..configuration import _integer |
| from ..setbpe.spatial import SpatialVocabulary |
| from .layout import SparseLayout |
|
|
| SPARSE_BPE_NAME = "sparse_bpe.json" |
| FORMAT = "actioncodec2_sparse_bpe" |
| FORMAT_VERSION = 1 |
| MAX_LOOKAHEAD = 32 |
|
|
|
|
| def _ctz(mask: int) -> int: |
| return (mask & -mask).bit_length() - 1 |
|
|
|
|
| class TokenTables: |
| """Per-token geometry shared by the encoder, decoder, trainer and grammar.""" |
|
|
| def __init__(self, support: Sequence[int], e_dim: int, lookahead: int) -> None: |
| if not 1 <= lookahead <= MAX_LOOKAHEAD: |
| raise ValueError(f"lookahead must be in [1, {MAX_LOOKAHEAD}]") |
| self.e_dim, self.lookahead, self.base = e_dim, lookahead, len(support) |
| self.masks: list[tuple[int, ...]] = [] |
| self.patterns: list[tuple[tuple[int, int], ...]] = [] |
| self.span: list[int] = [] |
| self.rdim: list[int] = [] |
| self.emask: list[int] = [] |
| self.first_from: list[tuple[int, ...]] = [] |
| for token, mask in enumerate(support): |
| if not mask: |
| raise ValueError("spatial token without support") |
| self._push(((0, token),), (int(mask),)) |
|
|
| def _push(self, pattern, masks) -> None: |
| span = max(o for o, m in enumerate(masks) if m) + 1 |
| masks = tuple(masks[:span]) |
| self.masks.append(masks) |
| self.patterns.append(tuple(pattern)) |
| self.span.append(span) |
| self.rdim.append(_ctz(masks[0])) |
| self.emask.append(sum(1 << o for o, m in enumerate(masks) if m >> self.e_dim & 1)) |
| first, following = [], -1 |
| for o in range(span - 1, -1, -1): |
| if masks[o]: |
| following = o |
| first.append(following) |
| self.first_from.append(tuple(reversed(first))) |
|
|
| def __len__(self) -> int: |
| return len(self.span) |
|
|
| def mask(self, token: int, offset: int) -> int: |
| masks = self.masks[token] |
| return masks[offset] if offset < len(masks) else 0 |
|
|
| def compose(self, left: int, right: int, offset: int) -> tuple[int, ...] | None: |
| """Union support of ``right`` rooted ``offset`` frames after ``left``, or ``None``.""" |
| if not 0 <= offset < self.lookahead: |
| return None |
| span = max(self.span[left], offset + self.span[right]) |
| if span > self.lookahead: |
| return None |
| out = [self.mask(left, o) for o in range(span)] |
| for o in range(self.span[right]): |
| if out[o + offset] & self.mask(right, o): |
| return None |
| out[o + offset] |= self.mask(right, o) |
| return tuple(out) |
|
|
| def add_merge(self, left: int, right: int, offset: int) -> int: |
| masks = self.compose(left, right, offset) |
| if masks is None: |
| raise ValueError("illegal temporal merge") |
| pattern = sorted( |
| self.patterns[left] + tuple((o + offset, t) for o, t in self.patterns[right]) |
| ) |
| self._push(pattern, masks) |
| return len(self) - 1 |
|
|
|
|
| def rule_key(left: int, right: int, offset: int) -> int: |
| return (left << 40) | (right << 16) | offset |
|
|
|
|
| class Instances: |
| """Mutable rooted token instances over one frame axis (reference implementation).""" |
|
|
| def __init__(self, tables: TokenTables, frames: int, dims: int) -> None: |
| self.tb, self.frames, self.dims = tables, frames, dims |
| self.tok: list[int] = [] |
| self.root: list[int] = [] |
| self.alive: list[bool] = [] |
| self.member = np.full((frames, dims), -1, dtype=np.int64) |
|
|
| def key(self, i: int) -> tuple[int, int, int, int]: |
| e = self.tb.emask[self.tok[i]] |
| root = self.root[i] |
| return (root, 1 if e else 0, root + _ctz(e) if e else 0, self.tb.rdim[self.tok[i]]) |
|
|
| def add(self, token: int, root: int) -> int: |
| self.tok.append(token) |
| self.root.append(root) |
| self.alive.append(True) |
| index = len(self.tok) - 1 |
| self._paint(index, index) |
| return index |
|
|
| def _paint(self, i: int, value: int) -> None: |
| token, root = self.tok[i], self.root[i] |
| for o, mask in enumerate(self.tb.masks[token]): |
| while mask: |
| d = _ctz(mask) |
| mask &= mask - 1 |
| if value >= 0 and self.member[root + o, d] >= 0: |
| raise RuntimeError("instance overlap") |
| self.member[root + o, d] = value |
|
|
| def kill(self, i: int) -> None: |
| self.alive[i] = False |
| self._paint(i, -1) |
|
|
| def _first(self, a: int, b: int) -> bool: |
| return (self.root[a], self.tb.rdim[self.tok[a]]) < (self.root[b], self.tb.rdim[self.tok[b]]) |
|
|
| def partners(self, i: int, allow_zero: bool): |
| """Ordered candidate pairs ``(a, b, offset)`` involving ``i`` in the lookahead window.""" |
| tb, L, root = self.tb, self.tb.lookahead, self.root[i] |
| lo, hi = max(0, root - L + 1), min(self.frames - 1, root + L - 1) |
| for f in range(lo, hi + 1): |
| for d in range(self.dims): |
| j = int(self.member[f, d]) |
| if j < 0 or j == i: |
| continue |
| tj, o = self.tok[j], f - self.root[j] |
| k = max(0, lo - self.root[j]) |
| if tb.first_from[tj][k] != o or _ctz(tb.mask(tj, o)) != d: |
| continue |
| a, b = (i, j) if self._first(i, j) else (j, i) |
| off = self.root[b] - self.root[a] |
| if off >= L or (off == 0 and not allow_zero): |
| continue |
| if max(tb.span[self.tok[a]], off + tb.span[self.tok[b]]) > L: |
| continue |
| yield a, b, off |
|
|
| def valid(self, a: int, b: int, off: int) -> bool: |
| """Whether merging ``a`` (rooted first) with ``b`` keeps the stream decodable.""" |
| tb, ta, tbb = self.tb, self.tok[a], self.tok[b] |
| span = max(tb.span[ta], off + tb.span[tbb]) |
| emask = tb.emask[ta] | (tb.emask[tbb] << off) |
| root = self.root[a] |
| merged = (root, 1 if emask else 0, root + _ctz(emask) if emask else 0, tb.rdim[ta]) |
| for o in range(span): |
| mask = tb.mask(ta, o) | (tb.mask(tbb, o - off) if o >= off else 0) |
| if not mask: |
| continue |
| u = root + o |
| h = int(self.member[u, tb.e_dim]) |
| if h >= 0 and h not in (a, b) and not self.key(h) > merged: |
| return False |
| if emask >> o & 1: |
| for d in range(self.dims): |
| j = int(self.member[u, d]) |
| if j >= 0 and j not in (a, b) and not self.key(j) < merged: |
| return False |
| return True |
|
|
| def serialize(self) -> list[int]: |
| live = sorted((i for i, alive in enumerate(self.alive) if alive), key=self.key) |
| return [self.tok[i] for i in live] |
|
|
|
|
| class SparseVocabulary: |
| """Fitted sparse token layer over one physical profile's dense cells.""" |
|
|
| def __init__(self, layout: SparseLayout, spatial_merges, temporal_merges, lookahead: int): |
| self.layout = layout |
| self.lookahead = _integer(lookahead, "lookahead", positive=True) |
| self.spatial = SpatialVocabulary(layout.sparse_bins, tuple(map(tuple, spatial_merges))) |
| self.temporal_merges = tuple( |
| tuple(_integer(v, "temporal merge value") for v in m) for m in temporal_merges |
| ) |
| self.tables = TokenTables([int(s) for s in self.spatial.support], layout.e_dim, lookahead) |
| self._rank = {} |
| for rank, (left, right, offset) in enumerate(self.temporal_merges): |
| if max(left, right) >= len(self.tables): |
| raise ValueError("a temporal merge must reference already-created tokens") |
| self.tables.add_merge(left, right, offset) |
| if self._rank.setdefault(rule_key(left, right, offset), rank) != rank: |
| raise ValueError("duplicate temporal rule") |
| self._native = None |
|
|
| @property |
| def vocab_size(self) -> int: |
| return len(self.tables) |
|
|
| def __getstate__(self): |
|
|
| return {**vars(self), "_native": None} |
|
|
| def _codec(self): |
| constructor = native_kernel("SparseCodec") |
| if constructor is not None and self._native is None: |
| self._native = constructor( |
| list(self.layout.bins), |
| self.layout.native_zero, |
| [tuple(m) for m in self.spatial.merges], |
| self.lookahead, |
| list(self.temporal_merges), |
| ) |
| return self._native if constructor is not None else None |
|
|
| def _cells(self, cells) -> np.ndarray: |
| values = np.asarray(cells) |
| if values.ndim != 3 or values.shape[0] < 1 or values.shape[1] < 1: |
| raise ValueError("cells must have shape (B, T, D) with B, T >= 1") |
| if values.shape[2] != self.layout.dims or values.dtype.kind not in "iu": |
| raise ValueError(f"cells must be integers with {self.layout.dims} slots") |
| if np.any(values < -1) or np.any(values >= np.asarray(self.layout.bins)): |
| raise ValueError("cell outside the primitive range") |
| return values |
|
|
| def encode_batch(self, cells, *, backend: str = "auto", threads: int = 0) -> list[list[int]]: |
| """Encode independent dense chunks ``(B, T, D)`` (``-1`` = inactive) into token rows.""" |
| values = self._cells(cells) |
| if resolve_backend(backend) == "native" and self._codec() is not None: |
| return self._codec().encode_batch(np.asarray(values, dtype=np.int32), threads) |
| return [self._encode_reference(chunk) for chunk in values] |
|
|
| def _encode_reference(self, cells: np.ndarray) -> list[int]: |
| sparse = self.layout.sparsify(cells) |
| state = Instances(self.tables, len(sparse) + self.lookahead, self.layout.dims + 1) |
| for t, row in enumerate(sparse): |
| if np.any(row >= 0): |
| for token in self.spatial.encode(row.tolist()): |
| state.add(int(token), t) |
| heap = [] |
|
|
| def offer(i): |
| for a, b, off in state.partners(i, True): |
| rank = self._rank.get(rule_key(state.tok[a], state.tok[b], off)) |
| if rank is not None: |
| item = (rank, state.root[a], self.tables.rdim[state.tok[a]], state.root[b]) |
| heapq.heappush(heap, (*item, a, b, off)) |
|
|
| for i in range(len(state.tok)): |
| offer(i) |
| while heap: |
| rank, *_, a, b, off = heapq.heappop(heap) |
| if not state.alive[a] or not state.alive[b] or not state.valid(a, b, off): |
| continue |
| root = state.root[a] |
| state.kill(a) |
| state.kill(b) |
| offer(state.add(self.tables.base + rank, root)) |
| return state.serialize() |
|
|
| def decode_batch( |
| self, rows: Sequence[Sequence[int]], horizon: int, active: Sequence[bool] |
| ) -> np.ndarray: |
| """Token rows to dense cells ``(B, horizon, D)``; inactive slots are ``-1``.""" |
| if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: |
| raise ValueError("horizon must be a positive integer") |
| flags = [bool(a) for a in active] |
| if len(flags) != self.layout.dims: |
| raise ValueError(f"active must have {self.layout.dims} flags") |
| checked = [] |
| for row in rows: |
| if any( |
| isinstance(t, (bool, np.bool_)) |
| or not isinstance(t, Integral) |
| or not 0 <= t < self.vocab_size |
| for t in row |
| ): |
| raise ValueError("token ids must be integers within the sparse vocabulary") |
| checked.append([int(t) for t in row]) |
| codec = self._codec() |
| if codec is not None: |
| return np.asarray(codec.decode_batch(checked, horizon, flags), dtype=np.int64) |
| out = np.empty((len(checked), horizon, self.layout.dims), dtype=np.int64) |
| for index, row in enumerate(checked): |
| out[index] = self.layout.densify(self._frames_reference(row, horizon, flags), flags) |
| return out |
|
|
| def _frames_reference(self, row: Sequence[int], horizon: int, active) -> np.ndarray: |
| tb, e_bit = self.tables, 1 << self.layout.e_dim |
| allowed = e_bit | sum(1 << d for d, flag in enumerate(active) if flag) |
| cover = [0] * (horizon + self.lookahead) |
| sparse = np.full((horizon, self.layout.dims + 1), -1, dtype=np.int64) |
| front = 0 |
| for token in row: |
| if not 0 <= token < len(tb): |
| raise ValueError("token outside vocab") |
| if front >= horizon: |
| raise ValueError("token after the last open frame") |
|
|
| own: dict[int, int] = {} |
| for offset, spatial in tb.patterns[token]: |
| u, mask = front + offset, int(self.spatial.support[spatial]) |
| if u >= horizon: |
| raise ValueError("token exceeds horizon") |
| if mask & ~allowed: |
| raise ValueError("token touches an inactive slot") |
| if own.get(u, 0) & mask: |
| raise ValueError("overlapping token supports") |
| if mask & e_bit and u == horizon - 1: |
| raise ValueError("end-of-frame atom on the last frame") |
| own[u] = own.get(u, 0) | mask |
| expansion = self.spatial.expansion[spatial] |
| sparse[u] = np.where(expansion >= 0, expansion, sparse[u]) |
| for u, mask in own.items(): |
| if cover[u] & mask: |
| raise ValueError("overlapping token supports") |
| if cover[u] & e_bit: |
| raise ValueError("token touches a closed frame") |
| cover[u] |= mask |
| while front < horizon and cover[front] & e_bit: |
| front += 1 |
| return sparse |
|
|
| def to_dict(self) -> dict: |
| return { |
| "format": FORMAT, |
| "format_version": FORMAT_VERSION, |
| "layout": {"bins": list(self.layout.bins), "zero": list(self.layout.zero)}, |
| "lookahead": self.lookahead, |
| "spatial_merges": [list(m) for m in self.spatial.merges], |
| "temporal_merges": [list(m) for m in self.temporal_merges], |
| } |
|
|
| @classmethod |
| def from_dict(cls, payload: dict) -> SparseVocabulary: |
| if payload.get("format") != FORMAT or payload.get("format_version") != FORMAT_VERSION: |
| raise ValueError(f"not an {FORMAT} v{FORMAT_VERSION} vocabulary") |
| layout = SparseLayout(tuple(payload["layout"]["bins"]), tuple(payload["layout"]["zero"])) |
| return cls( |
| layout, payload["spatial_merges"], payload["temporal_merges"], payload["lookahead"] |
| ) |
|
|
| def save_pretrained(self, directory: str | Path) -> Path: |
| path = Path(directory) / SPARSE_BPE_NAME |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(self.to_dict(), separators=(",", ":")) + "\n", encoding="utf-8") |
| return path |
|
|
| @classmethod |
| def from_pretrained(cls, directory: str | Path) -> SparseVocabulary: |
| path = Path(directory) / SPARSE_BPE_NAME |
| return cls.from_dict(json.loads(path.read_text(encoding="utf-8"))) |
|
|
