| """Runtime rooted-temporal Set-BPE vocabulary and lossless serialization.""" |
|
|
| from __future__ import annotations |
|
|
| import heapq |
| from collections.abc import Sequence |
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| from .._validation import active_dimensions |
| from ..native_serialization import NativeCacheMixin, decode_native, encode_native |
| from .spatial import SpatialVocabulary |
|
|
| __all__ = ["TemporalVocabulary"] |
|
|
| Rule = tuple[int, int, int] |
| Pattern = tuple[tuple[int, int], ...] |
|
|
|
|
| def _active_support(active_dims: Sequence[int] | None, dimension: int) -> int: |
| dims = active_dimensions(active_dims, dimension) |
| return sum(1 << value for value in dims) |
|
|
|
|
| def _validate_episode(values: np.ndarray, spatial: SpatialVocabulary) -> int: |
| if values.ndim != 2 or values.shape[1] != spatial.dimension or len(values) < 1: |
| raise ValueError( |
| f"every episode must have non-empty shape [T, {spatial.dimension}], got {values.shape}" |
| ) |
| counts = np.asarray(spatial.bins, dtype=np.int64) |
| if np.any(values < -1) or np.any(values >= counts[None, :]): |
| raise ValueError("episodes must contain -1 or an in-range cell index") |
| support = values >= 0 |
| if np.any(np.sum(support, axis=1) == 0): |
| raise ValueError("every timestep needs at least one active coordinate") |
| if np.any(support != support[:1]): |
| raise ValueError("active dimensions must remain constant within an episode") |
| return sum(1 << int(value) for value in np.flatnonzero(support[0])) |
|
|
|
|
| @dataclass |
| class TemporalVocabulary(NativeCacheMixin): |
| """A spatial vocabulary plus rooted, forward-only temporal merge rules.""" |
|
|
| spatial: SpatialVocabulary |
| lookahead: int |
| merges: tuple[Rule, ...] |
|
|
| def __post_init__(self) -> None: |
| if self.lookahead < 1: |
| raise ValueError("lookahead must be positive") |
| self.patterns: list[Pattern] = [((0, token),) for token in range(self.spatial.vocab_size)] |
| self.spans: list[int] = [1] * self.spatial.vocab_size |
| self.root: list[int] = list(range(self.spatial.vocab_size)) |
| self.support_masks: list[tuple[int, ...]] = [ |
| (int(self.spatial.support[token]),) for token in range(self.spatial.vocab_size) |
| ] |
| self._rules: dict[Rule, tuple[int, int]] = {} |
| for rank, rule in enumerate(self.merges): |
| left, right, offset = (int(value) for value in rule) |
| child = self.spatial.vocab_size + rank |
| if left >= child or right >= child or left < 0 or right < 0: |
| raise ValueError("a temporal merge must reference already-created tokens") |
| pattern, masks = _merge_pattern( |
| self.patterns[left], |
| self.patterns[right], |
| offset, |
| self.lookahead, |
| self.spatial, |
| ) |
| self.patterns.append(pattern) |
| self.spans.append(len(masks)) |
| self.root.append(self.root[left]) |
| self.support_masks.append(masks) |
| key = (left, right, offset) |
| if key in self._rules: |
| raise ValueError("duplicate temporal merge rule") |
| self._rules[key] = (child, rank) |
|
|
| @property |
| def vocab_size(self) -> int: |
| return self.spatial.vocab_size + len(self.merges) |
|
|
| def encode(self, codes: np.ndarray) -> list[int]: |
| """Encode a dense episode; ``-1`` slots are semantically absent.""" |
| values = np.asarray(codes, dtype=np.int64) |
| _validate_episode(values, self.spatial) |
| result = encode_native(self, values[None]) |
| if result is not None: |
| return result[0] |
|
|
| at: list[dict[int, int]] = [dict() for _ in range(values.shape[0])] |
| tokens: list[int] = [] |
| anchors: list[int] = [] |
| active: list[bool] = [] |
| for anchor, row in enumerate(values): |
| for token in self.spatial.encode(row): |
| instance = len(tokens) |
| tokens.append(token) |
| anchors.append(anchor) |
| active.append(True) |
| at[anchor][token] = instance |
|
|
| frontier: list[tuple[int, int, int, int, int, int]] = [] |
|
|
| def offer(left_instance: int, right_instance: int) -> None: |
| left = tokens[left_instance] |
| right = tokens[right_instance] |
| offset = anchors[right_instance] - anchors[left_instance] |
| rule = self._rules.get((left, right, offset)) |
| if rule is not None: |
| child, rank = rule |
| heapq.heappush( |
| frontier, |
| ( |
| rank, |
| anchors[left_instance], |
| int(self.spatial.anchor[self.root[child]]), |
| left_instance, |
| right_instance, |
| child, |
| ), |
| ) |
|
|
| for anchor, left_items in enumerate(at): |
| for offset in range(1, min(self.lookahead, len(at) - anchor)): |
| for left_instance in left_items.values(): |
| for right_instance in at[anchor + offset].values(): |
| offer(left_instance, right_instance) |
|
|
| while frontier: |
| _, anchor, _, left_instance, right_instance, child = heapq.heappop(frontier) |
| if not active[left_instance] or not active[right_instance]: |
| continue |
| if anchors[left_instance] != anchor: |
| continue |
| rule = ( |
| tokens[left_instance], |
| tokens[right_instance], |
| anchors[right_instance] - anchor, |
| ) |
| resolved = self._rules.get(rule) |
| if resolved is None or resolved[0] != child: |
| continue |
|
|
| active[left_instance] = False |
| active[right_instance] = False |
| del at[anchor][tokens[left_instance]] |
| del at[anchors[right_instance]][tokens[right_instance]] |
| new_instance = len(tokens) |
| tokens.append(child) |
| anchors.append(anchor) |
| active.append(True) |
| at[anchor][child] = new_instance |
|
|
| for offset in range(1, self.lookahead): |
| future = anchor + offset |
| if future < len(at): |
| for other in at[future].values(): |
| offer(new_instance, other) |
| past = anchor - offset |
| if past >= 0: |
| for other in at[past].values(): |
| offer(other, new_instance) |
|
|
| ordered: list[tuple[int, int, int]] = [] |
| for anchor, instances in enumerate(at): |
| for token, instance in instances.items(): |
| if active[instance]: |
| ordered.append((anchor, int(self.spatial.anchor[self.root[token]]), token)) |
| ordered.sort() |
| return [token for _, _, token in ordered] |
|
|
| def decode( |
| self, |
| stream: Sequence[int], |
| *, |
| active_dims: Sequence[int] | None = None, |
| ) -> np.ndarray: |
| """Decode ids using the route's requested active-dimension support.""" |
| if not stream: |
| return np.empty((0, self.spatial.dimension), dtype=np.int64) |
| dims = active_dimensions(active_dims, self.spatial.dimension) |
| result = decode_native(self, [[int(token) for token in stream]], dims) |
| if result is not None: |
| return result[0] |
| full = _active_support(active_dims, self.spatial.dimension) |
| coverage: list[int] = [0] |
| spatial_tokens: list[list[int]] = [[]] |
| front = 0 |
| for raw_token in stream: |
| token = int(raw_token) |
| if not 0 <= token < self.vocab_size: |
| raise ValueError(f"temporal token {token} is outside the vocabulary") |
| while front < len(coverage) and coverage[front] == full: |
| front += 1 |
| if front == len(coverage): |
| coverage.append(0) |
| spatial_tokens.append([]) |
|
|
| uncovered = full & ~coverage[front] |
| expected_anchor = int(uncovered & -uncovered).bit_length() - 1 |
| actual_anchor = int(self.spatial.anchor[self.root[token]]) |
| if actual_anchor != expected_anchor: |
| raise ValueError( |
| "temporal stream violates root order at the earliest incomplete step" |
| ) |
|
|
| for offset, spatial_token in self.patterns[token]: |
| step = front + offset |
| while step >= len(coverage): |
| coverage.append(0) |
| spatial_tokens.append([]) |
| support = int(self.spatial.support[spatial_token]) |
| if coverage[step] & support: |
| raise ValueError("temporal stream contains overlapping spatial support") |
| coverage[step] |= support |
| spatial_tokens[step].append(spatial_token) |
|
|
| if any(mask != full for mask in coverage): |
| raise ValueError("temporal stream ended with an incomplete control step") |
| return np.stack( |
| [self.spatial.decode(tokens, active_dims=active_dims) for tokens in spatial_tokens] |
| ) |
|
|
|
|
| def _merge_pattern( |
| left: Pattern, |
| right: Pattern, |
| offset: int, |
| lookahead: int, |
| spatial: SpatialVocabulary, |
| ) -> tuple[Pattern, tuple[int, ...]]: |
| if not 1 <= offset < lookahead: |
| raise ValueError("a temporal merge offset must be in [1, lookahead)") |
| shifted = tuple((time + offset, token) for time, token in right) |
| pattern = tuple(sorted((*left, *shifted))) |
| span = 1 + max(time for time, _ in pattern) |
| if span > lookahead: |
| raise ValueError("temporal merge exceeds the lookahead horizon") |
| masks = [0] * span |
| roots = 0 |
| for time, token in pattern: |
| support = int(spatial.support[token]) |
| if masks[time] & support: |
| raise ValueError("temporal merge joins overlapping spatial supports") |
| masks[time] |= support |
| roots += int(time == 0) |
| if roots != 1: |
| raise ValueError("a rooted temporal token must have exactly one time-zero spatial token") |
| return pattern, tuple(masks) |
|
|
