| 'Temporal Set-BPE fitting implementations.' | |
| from __future__ import annotations | |
| import heapq | |
| from array import array | |
| from collections import Counter | |
| from collections.abc import Iterable, Iterator, Sequence | |
| import numpy as np | |
| from ..native_backend import require_native, resolve_backend | |
| from .spatial import SpatialVocabulary | |
| from .temporal import Pattern, Rule, TemporalVocabulary, _merge_pattern, _validate_episode | |
| __all__ = ["train_temporal"] | |
| class _SparseTemporalTrainer: | |
| def __init__( | |
| self, | |
| episodes: Sequence[np.ndarray], | |
| spatial: SpatialVocabulary, | |
| lookahead: int, | |
| max_vocab_size: int, | |
| ) -> None: | |
| if lookahead < 1: | |
| raise ValueError("lookahead must be positive") | |
| self.spatial = spatial | |
| self.lookahead = lookahead | |
| self.max_vocab_size = max_vocab_size | |
| self.patterns: list[Pattern] = [((0, token),) for token in range(spatial.vocab_size)] | |
| self.spans: list[int] = [1] * spatial.vocab_size | |
| self.masks: list[tuple[int, ...]] = [ | |
| (int(spatial.support[token]),) for token in range(spatial.vocab_size) | |
| ] | |
| self._mask_table = np.zeros((max_vocab_size, lookahead), dtype=np.uint64) | |
| self._mask_table[: spatial.vocab_size, 0] = spatial.support.astype(np.uint64, copy=False) | |
| self.merges: list[Rule] = [] | |
| self.tokens = array("I") | |
| self.anchors = array("I") | |
| self.slots = array("I") | |
| total_anchors = sum(len(episode) + lookahead for episode in episodes) | |
| self.member = np.full((total_anchors, spatial.dimension), -1, dtype=np.int32) | |
| self.occupied = np.zeros(total_anchors, dtype=np.uint64) | |
| self._slot_mask = (1 << spatial.dimension) - 1 | |
| self.occurrences: list[dict[int, int]] = [dict() for _ in range(spatial.vocab_size)] | |
| self.episode_count = len(episodes) | |
| self.action_count = sum(len(episode) for episode in episodes) | |
| self.initial_count = 0 | |
| self.counts: dict[int, int] = {} | |
| self.self_versions: dict[int, int] = {} | |
| self.heap: list[tuple[int, int, int, int, int, bool]] = [] | |
| self.self_cache: dict[int, tuple[int, list[tuple[int, int]]]] = {} | |
| self._legal_table = np.zeros((lookahead, max_vocab_size, max_vocab_size), dtype=np.uint8) | |
| for offset in range(1, lookahead): | |
| self._legal_table[offset, : spatial.vocab_size, : spatial.vocab_size] = 2 | |
| self.statistics = Counter() | |
| self._active_count = 0 | |
| self._load_episodes(episodes) | |
| self.counts = self._initial_pair_counts() | |
| self.heap = self._initial_heap() | |
| self.statistics["initial_edge_types"] = len(self.counts) | |
| self.statistics["initial_edges"] = sum(self.counts.values()) | |
| def _load_episodes(self, episodes: Sequence[np.ndarray]) -> None: | |
| base = 0 | |
| encoding_cache: dict[bytes, tuple[int, ...]] = {} | |
| for episode in episodes: | |
| values = np.asarray(episode, dtype=np.int64) | |
| _validate_episode(values, self.spatial) | |
| needed = base + len(values) + self.lookahead | |
| for time, row in enumerate(values): | |
| cache_key = row.tobytes() | |
| encoded = encoding_cache.get(cache_key) | |
| if encoded is None: | |
| encoded = tuple(self.spatial.encode(row)) | |
| encoding_cache[cache_key] = encoded | |
| anchor = base + time | |
| for token in encoded: | |
| self._new_instance(token, anchor) | |
| self.initial_count += 1 | |
| base = needed | |
| def _initial_pair_counts(self) -> dict[int, int]: | |
| base_vocab = self.spatial.vocab_size | |
| dense_size = base_vocab * base_vocab * self.lookahead | |
| if dense_size <= 1 << 25: | |
| dense_counts = np.zeros(dense_size, dtype=np.int64) | |
| for offset in range(1, self.lookahead): | |
| left_rows = self.member[:-offset] | |
| right_rows = self.member[offset:] | |
| for left_slot in range(self.spatial.dimension): | |
| left = left_rows[:, left_slot] | |
| for right_slot in range(self.spatial.dimension): | |
| right = right_rows[:, right_slot] | |
| valid = (left >= 0) & (right >= 0) | |
| if not valid.any(): | |
| continue | |
| packed = ( | |
| left[valid].astype(np.int64) * base_vocab + right[valid] | |
| ) * self.lookahead + offset | |
| np.add.at(dense_counts, packed, 1) | |
| initial_counts: dict[int, int] = {} | |
| for raw_key in np.flatnonzero(dense_counts): | |
| pair, offset = divmod(int(raw_key), self.lookahead) | |
| left, right = divmod(pair, base_vocab) | |
| packed_key = (left * self.max_vocab_size + right) * self.lookahead + offset | |
| initial_counts[packed_key] = int(dense_counts[raw_key]) | |
| return initial_counts | |
| return self._initial_pair_counts_sparse() | |
| def _initial_pair_counts_sparse(self) -> dict[int, int]: | |
| fallback: Counter[int] = Counter() | |
| for anchor, left_row in enumerate(self.member): | |
| left_tokens = left_row[left_row >= 0] | |
| for offset in range(1, self.lookahead): | |
| target = anchor + offset | |
| if target >= len(self.member): | |
| break | |
| right_tokens = self.member[target] | |
| right_tokens = right_tokens[right_tokens >= 0] | |
| for raw_left in left_tokens: | |
| for raw_right in right_tokens: | |
| packed_key = ( | |
| int(raw_left) * self.max_vocab_size + int(raw_right) | |
| ) * self.lookahead + offset | |
| fallback[packed_key] += 1 | |
| return dict(fallback) | |
| def _initial_heap(self) -> list[tuple[int, int, int, int, int, bool]]: | |
| heap: list[tuple[int, int, int, int, int, bool]] = [] | |
| for packed_key, count in self.counts.items(): | |
| pair, offset = divmod(packed_key, self.lookahead) | |
| left, right = divmod(pair, self.max_vocab_size) | |
| heap.append((-count, left, right, offset, 0, False)) | |
| heapq.heapify(heap) | |
| return heap | |
| def _new_instance(self, token: int, anchor: int) -> int: | |
| while token >= len(self.occurrences): | |
| self.occurrences.append({}) | |
| if anchor in self.occurrences[token]: | |
| raise RuntimeError("a temporal type cannot occur twice at one anchor") | |
| occupied = int(self.occupied[anchor]) | |
| available = self._slot_mask ^ occupied | |
| if available == 0: | |
| raise RuntimeError("temporal anchor exceeded the spatial token-width bound") | |
| lowest = available & -available | |
| slot = lowest.bit_length() - 1 | |
| self.occupied[anchor] = occupied | lowest | |
| instance = len(self.tokens) | |
| self.tokens.append(token) | |
| self.anchors.append(anchor) | |
| self.slots.append(slot) | |
| self.member[anchor, slot] = token | |
| self._active_count += 1 | |
| self.occurrences[token][anchor] = instance | |
| return instance | |
| def _initialize_legality(self, child: int) -> None: | |
| count = child + 1 | |
| spans = np.asarray(self.spans[:count], dtype=np.int64) | |
| masks = self._mask_table[:count] | |
| child_masks = self._mask_table[child] | |
| child_span = self.spans[child] | |
| for offset in range(1, self.lookahead): | |
| overlap = np.zeros(count, dtype=bool) | |
| for relative in range(self.lookahead - offset): | |
| overlap |= (child_masks[offset + relative] & masks[:, relative]) != 0 | |
| legal = (offset + spans <= self.lookahead) & ~overlap | |
| self._legal_table[offset, child, :count] = np.where(legal, 2, 1) | |
| overlap.fill(False) | |
| for relative in range(min(child_span, self.lookahead - offset)): | |
| overlap |= (masks[:, offset + relative] & child_masks[relative]) != 0 | |
| legal = (offset + child_span <= self.lookahead) & ~overlap | |
| self._legal_table[offset, :count, child] = np.where(legal, 2, 1) | |
| def _entry( | |
| self, packed_key: int, priority: int, certified: bool | |
| ) -> tuple[int, int, int, int, int, bool]: | |
| pair, offset = divmod(packed_key, self.lookahead) | |
| left, right = divmod(pair, self.max_vocab_size) | |
| return ( | |
| -priority, | |
| left, | |
| right, | |
| offset, | |
| self.self_versions.get(packed_key, 0) if left == right else 0, | |
| certified, | |
| ) | |
| def _select(self, key: Rule) -> list[tuple[int, int]]: | |
| left, right, offset = key | |
| left_occurrences = self.occurrences[left] | |
| right_occurrences = self.occurrences[right] | |
| if left != right: | |
| return [ | |
| (left_instance, right_occurrences[anchor + offset]) | |
| for anchor, left_instance in left_occurrences.items() | |
| if anchor + offset in right_occurrences | |
| ] | |
| selected: list[tuple[int, int]] = [] | |
| consumed: set[int] = set() | |
| for anchor in sorted(left_occurrences): | |
| target = anchor + offset | |
| if anchor in consumed or target in consumed or target not in right_occurrences: | |
| continue | |
| selected.append((left_occurrences[anchor], right_occurrences[target])) | |
| consumed.add(anchor) | |
| consumed.add(target) | |
| return selected | |
| def best(self) -> tuple[Rule, list[tuple[int, int]]] | None: | |
| while self.heap: | |
| negative, left, right, offset, version, certified = heapq.heappop(self.heap) | |
| self.statistics["heap_pops"] += 1 | |
| key = (left, right, offset) | |
| packed_key = (left * self.max_vocab_size + right) * self.lookahead + offset | |
| current_count = self.counts.get(packed_key, 0) | |
| if current_count <= 0: | |
| continue | |
| if left == right: | |
| current_version = self.self_versions.get(packed_key, 0) | |
| stale = current_version != version | |
| else: | |
| stale = current_count != -negative | |
| if stale: | |
| heapq.heappush(self.heap, self._entry(packed_key, current_count, False)) | |
| continue | |
| if left == right and not certified: | |
| matches = self._select(key) | |
| self.self_cache[packed_key] = (version, matches) | |
| heapq.heappush(self.heap, self._entry(packed_key, len(matches), True)) | |
| self.statistics["self_certifications"] += 1 | |
| continue | |
| if certified: | |
| cached = self.self_cache.get(packed_key) | |
| if cached is None or cached[0] != version: | |
| continue | |
| matches = cached[1] | |
| else: | |
| matches = self._select(key) | |
| gain = -negative | |
| if len(matches) != gain: | |
| raise RuntimeError("sparse temporal counter disagrees with exact matching") | |
| return key, matches | |
| return None | |
| def _accumulate_incident( | |
| self, | |
| instances: Sequence[int] | set[int], | |
| multiplier: int, | |
| deltas: Counter[int], | |
| ) -> None: | |
| grouped: dict[int, list[int]] = {} | |
| for instance in instances: | |
| grouped.setdefault(self.tokens[instance], []).append(self.anchors[instance]) | |
| width = len(self.patterns) | |
| total_anchors = self.member.shape[0] | |
| vocab_width = self.max_vocab_size | |
| lookahead = self.lookahead | |
| for token, raw_anchors in grouped.items(): | |
| anchors = np.asarray(raw_anchors, dtype=np.int64) | |
| for offset in range(1, self.lookahead): | |
| future_anchors = anchors[anchors + offset < total_anchors] + offset | |
| if future_anchors.size: | |
| values = self.member[future_anchors].reshape(-1) | |
| values = values[values >= 0] | |
| if values.size: | |
| counts = np.bincount(values, minlength=width) | |
| neighbors = np.flatnonzero(counts) | |
| neighbors = neighbors[self._legal_table[offset, token, neighbors] == 2] | |
| for raw_neighbor in neighbors: | |
| neighbor = int(raw_neighbor) | |
| packed_key = (token * vocab_width + neighbor) * lookahead + offset | |
| deltas[packed_key] += multiplier * int(counts[neighbor]) | |
| past_anchors = anchors[anchors >= offset] - offset | |
| if past_anchors.size: | |
| values = self.member[past_anchors].reshape(-1) | |
| values = values[values >= 0] | |
| if values.size: | |
| counts = np.bincount(values, minlength=width) | |
| neighbors = np.flatnonzero(counts) | |
| neighbors = neighbors[self._legal_table[offset, neighbors, token] == 2] | |
| for raw_neighbor in neighbors: | |
| neighbor = int(raw_neighbor) | |
| packed_key = (neighbor * vocab_width + token) * lookahead + offset | |
| deltas[packed_key] += multiplier * int(counts[neighbor]) | |
| def _correct_internal_edges( | |
| self, | |
| instances: Sequence[int] | set[int], | |
| correction: int, | |
| deltas: Counter[int], | |
| ) -> None: | |
| by_anchor: dict[int, list[int]] = {} | |
| for instance in instances: | |
| by_anchor.setdefault(self.anchors[instance], []).append(self.tokens[instance]) | |
| vocab_width = self.max_vocab_size | |
| lookahead = self.lookahead | |
| for anchor, left_tokens in by_anchor.items(): | |
| for offset in range(1, self.lookahead): | |
| right_tokens = by_anchor.get(anchor + offset) | |
| if right_tokens is None: | |
| continue | |
| for left in left_tokens: | |
| for right in right_tokens: | |
| if self._legal_table[offset, left, right] == 2: | |
| packed_key = (left * vocab_width + right) * lookahead + offset | |
| deltas[packed_key] += correction | |
| def _apply_deltas(self, deltas: Counter[int], selected: Rule) -> None: | |
| counts = self.counts | |
| self_versions = self.self_versions | |
| heap = self.heap | |
| updates = 0 | |
| vocab_width = self.max_vocab_size | |
| lookahead = self.lookahead | |
| selected_key = (selected[0] * vocab_width + selected[1]) * lookahead + selected[2] | |
| for packed_key, delta in deltas.items(): | |
| previous = counts.get(packed_key, 0) | |
| count = previous + delta | |
| if count < 0: | |
| raise RuntimeError("temporal edge counter became negative") | |
| pair = packed_key // lookahead | |
| is_self = pair % (vocab_width + 1) == 0 | |
| if is_self: | |
| version = self_versions.get(packed_key, 0) + 1 | |
| self_versions[packed_key] = version | |
| self.self_cache.pop(packed_key, None) | |
| else: | |
| version = 0 | |
| if count: | |
| counts[packed_key] = count | |
| if previous == 0 or packed_key == selected_key: | |
| pair, offset = divmod(packed_key, lookahead) | |
| left, right = divmod(pair, vocab_width) | |
| heapq.heappush(heap, (-count, left, right, offset, version, False)) | |
| else: | |
| counts.pop(packed_key, None) | |
| updates += abs(delta) | |
| self.statistics["counter_updates"] += updates | |
| def commit(self, key: Rule, matches: Sequence[tuple[int, int]]) -> None: | |
| left, right, offset = key | |
| pattern, masks = _merge_pattern( | |
| self.patterns[left], | |
| self.patterns[right], | |
| offset, | |
| self.lookahead, | |
| self.spatial, | |
| ) | |
| child = len(self.patterns) | |
| self.patterns.append(pattern) | |
| self.spans.append(len(masks)) | |
| self.masks.append(masks) | |
| self.occurrences.append({}) | |
| self.merges.append(key) | |
| self._mask_table[child, : len(masks)] = masks | |
| self._initialize_legality(child) | |
| removed = {instance for pair in matches for instance in pair} | |
| if len(removed) != 2 * len(matches): | |
| raise RuntimeError("selected temporal matches overlap") | |
| deltas: Counter[int] = Counter() | |
| self._accumulate_incident(removed, -1, deltas) | |
| self._correct_internal_edges(removed, 1, deltas) | |
| for instance in removed: | |
| token = self.tokens[instance] | |
| anchor = self.anchors[instance] | |
| self.member[anchor, self.slots[instance]] = -1 | |
| self._active_count -= 1 | |
| self.occupied[anchor] = int(self.occupied[anchor]) & ~(1 << self.slots[instance]) | |
| del self.occurrences[token][anchor] | |
| children = [ | |
| self._new_instance(child, self.anchors[left_instance]) for left_instance, _ in matches | |
| ] | |
| self._accumulate_incident(children, 1, deltas) | |
| self._correct_internal_edges(children, -1, deltas) | |
| self._apply_deltas(deltas, key) | |
| def active_count(self) -> int: | |
| return self._active_count | |
| def _prepare_episodes( | |
| episodes: Iterable[np.ndarray], spatial: SpatialVocabulary | |
| ) -> Sequence[np.ndarray]: | |
| if getattr(episodes, "_actioncodec2_validated", False): | |
| if not isinstance(episodes, Sequence) or len(episodes) == 0: | |
| raise ValueError("validated temporal corpus must be a non-empty sequence") | |
| return episodes | |
| prepared: list[np.ndarray] = [] | |
| for episode in episodes: | |
| values = np.asarray(episode) | |
| if values.dtype.kind not in {"i", "u"}: | |
| raise ValueError("episodes must contain integer cell indices") | |
| _validate_episode(values, spatial) | |
| prepared.append(np.ascontiguousarray(values)) | |
| if not prepared or sum(len(episode) for episode in prepared) == 0: | |
| raise ValueError("temporal training requires at least one non-empty episode") | |
| return prepared | |
| def _train_temporal_python( | |
| episodes: Iterable[np.ndarray], | |
| spatial: SpatialVocabulary, | |
| *, | |
| vocab_size: int, | |
| lookahead: int, | |
| progress: bool = False, | |
| ) -> tuple[TemporalVocabulary, dict[str, object]]: | |
| prepared = _prepare_episodes(episodes, spatial) | |
| if vocab_size < spatial.vocab_size: | |
| raise ValueError( | |
| f"vocab_size {vocab_size} is below the {spatial.vocab_size} spatial tokens" | |
| ) | |
| trainer = _SparseTemporalTrainer(prepared, spatial, lookahead, vocab_size) | |
| length = trainer.initial_count | |
| for step in range(vocab_size - spatial.vocab_size): | |
| winner = trainer.best() | |
| if winner is None: | |
| break | |
| key, matches = winner | |
| if not matches: | |
| break | |
| trainer.commit(key, matches) | |
| length -= len(matches) | |
| if progress and step % 64 == 0: | |
| print( | |
| f" temporal merge {step:>5} tokens/action {length / trainer.action_count:.4f}", | |
| flush=True, | |
| ) | |
| if length != trainer.active_count: | |
| raise RuntimeError("temporal token accounting disagrees with active tiling") | |
| vocabulary = TemporalVocabulary(spatial, lookahead, tuple(trainer.merges)) | |
| report: dict[str, object] = { | |
| "episodes": trainer.episode_count, | |
| "actions": trainer.action_count, | |
| "lookahead": lookahead, | |
| "spatial_vocab_size": spatial.vocab_size, | |
| "merges": len(trainer.merges), | |
| "vocab_size": vocabulary.vocab_size, | |
| "tokens_per_action_initial": trainer.initial_count / trainer.action_count, | |
| "tokens_per_action": length / trainer.action_count, | |
| "compression": trainer.initial_count / length, | |
| "trainer": "sparse_exact", | |
| **{key: int(value) for key, value in trainer.statistics.items()}, | |
| } | |
| return vocabulary, report | |
| def _flatten_instances_native( | |
| episodes: Sequence[np.ndarray], | |
| spatial: SpatialVocabulary, | |
| lookahead: int, | |
| threads: int, | |
| ) -> tuple[Iterator[tuple[np.ndarray, np.ndarray]], int, int]: | |
| if getattr(episodes, "_actioncodec2_validated", False): | |
| return _flatten_cell_corpus(episodes, spatial, lookahead, threads) | |
| actions = sum(len(episode) for episode in episodes) | |
| total_anchors = actions + len(episodes) * lookahead | |
| max_batch_actions = 1_000_000 | |
| merges = np.asarray(spatial.merges, dtype=np.int32).reshape(-1, 2) | |
| def generate() -> Iterator[tuple[np.ndarray, np.ndarray]]: | |
| pending: list[np.ndarray] = [] | |
| pending_anchors: list[np.ndarray] = [] | |
| pending_actions = 0 | |
| base = 0 | |
| def encode() -> tuple[np.ndarray, np.ndarray]: | |
| codes = np.concatenate(pending, axis=0).astype(np.int32, copy=False) | |
| row_anchors = np.concatenate(pending_anchors) | |
| tokens, anchors = require_native().spatial_encode_instances( | |
| np.ascontiguousarray(codes), | |
| np.ascontiguousarray(row_anchors), | |
| [int(count) for count in spatial.bins], | |
| np.ascontiguousarray(merges), | |
| int(threads), | |
| ) | |
| return np.asarray(tokens, dtype=np.int32), np.asarray(anchors, dtype=np.int64) | |
| for episode in episodes: | |
| length = len(episode) | |
| pending.append(episode) | |
| pending_anchors.append(np.arange(base, base + length, dtype=np.int64)) | |
| pending_actions += length | |
| base += length + lookahead | |
| if pending_actions >= max_batch_actions: | |
| yield encode() | |
| pending.clear() | |
| pending_anchors.clear() | |
| pending_actions = 0 | |
| if pending: | |
| yield encode() | |
| return generate(), total_anchors, actions | |
| def _flatten_cell_corpus( | |
| episodes: Sequence[np.ndarray], | |
| spatial: SpatialVocabulary, | |
| lookahead: int, | |
| threads: int, | |
| ) -> tuple[Iterator[tuple[np.ndarray, np.ndarray]], int, int]: | |
| actions = sum(len(codes) for codes in episodes.chunks if codes is not None) | |
| total_anchors = actions + len(episodes) * lookahead | |
| merges = np.asarray(spatial.merges, dtype=np.int32).reshape(-1, 2) | |
| chunks = episodes.chunks | |
| offsets_by_chunk = episodes.offsets | |
| def generate() -> Iterator[tuple[np.ndarray, np.ndarray]]: | |
| base = 0 | |
| for chunk_index, (codes, offsets) in enumerate(zip(chunks, offsets_by_chunk, strict=True)): | |
| if codes is None: | |
| raise RuntimeError("cell corpus chunk has already been consumed") | |
| lengths = np.diff(offsets) | |
| row_anchors = ( | |
| base | |
| + np.arange(len(codes), dtype=np.int64) | |
| + np.repeat(np.arange(len(lengths), dtype=np.int64) * lookahead, lengths) | |
| ) | |
| tokens, anchors = require_native().spatial_encode_instances( | |
| np.ascontiguousarray(codes), | |
| row_anchors, | |
| [int(count) for count in spatial.bins], | |
| np.ascontiguousarray(merges), | |
| int(threads), | |
| ) | |
| episodes.release_chunk(chunk_index) | |
| yield np.asarray(tokens, dtype=np.int32), np.asarray(anchors, dtype=np.int64) | |
| base += len(codes) + len(lengths) * lookahead | |
| return generate(), total_anchors, actions | |
| def train_temporal( | |
| episodes: Iterable[np.ndarray], | |
| spatial: SpatialVocabulary, | |
| *, | |
| vocab_size: int, | |
| lookahead: int, | |
| backend: str = "auto", | |
| progress: bool = False, | |
| threads: int = 0, | |
| ) -> tuple[TemporalVocabulary, dict[str, object]]: | |
| 'Fit rooted temporal Set-BPE on top of a frozen spatial vocabulary.' | |
| prepared = _prepare_episodes(episodes, spatial) | |
| if resolve_backend(backend) == "python": | |
| return _train_temporal_python( | |
| prepared, | |
| spatial, | |
| vocab_size=vocab_size, | |
| lookahead=lookahead, | |
| progress=progress, | |
| ) | |
| episode_count = len(prepared) | |
| instance_chunks, total_anchors, actions = _flatten_instances_native( | |
| prepared, spatial, lookahead, threads | |
| ) | |
| merges, report = require_native().train_temporal( | |
| instance_chunks, | |
| int(total_anchors), | |
| np.ascontiguousarray(spatial.support, dtype=np.uint64), | |
| np.ascontiguousarray(spatial.anchor, dtype=np.int32), | |
| int(spatial.dimension), | |
| int(lookahead), | |
| int(vocab_size), | |
| int(actions), | |
| int(threads), | |
| ) | |
| vocabulary = TemporalVocabulary( | |
| spatial, int(lookahead), tuple(tuple(int(value) for value in rule) for rule in merges) | |
| ) | |
| payload = dict(report) | |
| payload["episodes"] = episode_count | |
| return vocabulary, payload | |