File size: 10,248 Bytes
895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 | """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)
|