"""ActionCodec2Sparse vocabulary training: optional spatial stage, then certified temporal merges.""" from __future__ import annotations import heapq import time from collections import defaultdict from collections.abc import Sequence import numpy as np from ..native_backend import native_kernel, resolve_backend from ..setbpe.spatial import aggregate_records from ..setbpe.spatial_training import train_spatial from .layout import SparseLayout, sparsify_chunks from .vocabulary import Instances, SparseVocabulary, TokenTables, rule_key MIN_PAIR_COUNT = 2 CERTIFY_BATCH = 32 def _unpack(key: int) -> tuple[int, int, int]: return key >> 40, (key >> 16) & 0xFFFFFF, key & 0xFFFF class _Stat: __slots__ = ( "retired", "raw", "version", "exact", "exact_version", "added", "added_at_exact", "dirty_epoch", "live", "queued_version", "queued_count", ) def __init__(self) -> None: self.retired = self.live = False self.raw = self.version = self.exact_version = self.added = self.added_at_exact = 0 self.dirty_epoch = self.queued_version = self.queued_count = 0 self.exact = -1 @property def certified(self) -> bool: return self.exact >= 0 and self.exact_version == self.version @property def bound(self) -> int: """Upper bound on the exact count: certified counts grow only by inserted pairs.""" if self.certified: return self.exact if self.exact >= 0: return min(self.raw, self.exact + self.added - self.added_at_exact) return self.raw class ReferenceTrainer: """Python reference of the native temporal trainer (same interface and merges).""" def __init__( self, spatial_support, dims, lookahead, vocab_size, allow_zero_offset, min_count=2 ): self.tables = TokenTables([int(s) for s in spatial_support], dims - 1, lookahead) self.dims, self.vocab_size = dims, vocab_size self.allow_zero, self.min_count = allow_zero_offset, min_count self.stats: dict[int, _Stat] = defaultdict(_Stat) self.heap: list[tuple[int, int, int]] = [] self._merges: list[tuple[int, int, int]] = [] self.epoch = 0 self.counters = dict(dry_runs=0, rejected_by_validity=0, applied=0) def set_instances(self, anchors, tokens, frames) -> None: self.state = Instances(self.tables, int(frames), self.dims) self.occurrence: dict[int, list[int]] = defaultdict(list) for anchor, token in zip( np.asarray(anchors).tolist(), np.asarray(tokens).tolist(), strict=True ): if not 0 <= token < self.tables.base: raise ValueError("token outside spatial vocab") self.occurrence[token].append(self.state.add(token, anchor)) for values in self.occurrence.values(): values.sort(key=lambda i: self.state.root[i]) self.initial_instances = len(self.state.tok) def merges(self) -> list[tuple[int, int, int]]: return list(self._merges) def _queue(self, key: int, stat: _Stat, count: int) -> None: heapq.heappush(self.heap, (-count, key, stat.version)) stat.live, stat.queued_version, stat.queued_count = True, stat.version, count def _pairs(self, left: int, right: int, offset: int): st = self.state for a in self.occurrence[left]: if not st.alive[a]: continue f = st.root[a] + offset if f >= st.frames: continue b = int(st.member[f, self.tables.rdim[right]]) if b >= 0 and b != a and st.alive[b] and st.tok[b] == right and st.root[b] == f: yield a, b def _count_exact(self, key: int) -> int: left, right, offset = _unpack(key) consumed, count = set(), 0 for a, b in self._pairs(left, right, offset): if left == right and a in consumed: continue if not self.state.valid(a, b, offset): continue if left == right: consumed.add(b) count += 1 return count def _take(self): count, key, version = heapq.heappop(self.heap) stat = self.stats.get(key) if stat is None or stat.retired or not stat.live or version != stat.queued_version: return None stat.live = False if version == stat.version: return -count, key, stat if stat.bound > 0: self._queue(key, stat, stat.bound) return None def _current_certified(self, entry) -> bool: _, key, version = entry stat = self.stats.get(key) return ( stat is not None and not stat.retired and stat.live and version == stat.queued_version == stat.version and stat.certified ) def _certify(self, keys: list[int]) -> None: self.counters["dry_runs"] += len(keys) exact = [self._count_exact(key) for key in keys] for key, count in zip(keys, exact, strict=True): stat = self.stats[key] stat.exact, stat.exact_version, stat.added_at_exact = count, stat.version, stat.added if count >= self.min_count: self._queue(key, stat, count) def _apply(self, left: int, right: int, offset: int, child: int) -> list[int]: st, delta = self.state, defaultdict(int) pairs = list(self._pairs(left, right, offset)) tok = st.tok for a, b in pairs: if not st.alive[a] or not st.alive[b]: continue if not st.valid(a, b, offset): self.counters["rejected_by_validity"] += 1 continue for x, y, o in st.partners(a, self.allow_zero): delta[rule_key(tok[x], tok[y], o)] -= 1 for x, y, o in st.partners(b, self.allow_zero): if a not in (x, y): delta[rule_key(tok[x], tok[y], o)] -= 1 root = st.root[a] st.kill(a) st.kill(b) merged = st.add(child, root) self.occurrence[child].append(merged) for x, y, o in st.partners(merged, self.allow_zero): delta[rule_key(tok[x], tok[y], o)] += 1 self.counters["applied"] += 1 dirty = [] for key, change in delta.items(): stat = self.stats[key] stat.raw += change stat.added += max(change, 0) stat.version += 1 if stat.dirty_epoch != self.epoch: stat.dirty_epoch = self.epoch dirty.append(key) return dirty def train(self, progress: bool = False) -> dict: del progress st = self.state for i in range(len(st.tok)): for a, b, o in st.partners(i, self.allow_zero): if a == i: self.stats[rule_key(st.tok[a], st.tok[b], o)].raw += 1 initial_types = len(self.stats) for key, stat in list(self.stats.items()): self._queue(key, stat, stat.raw) while len(self.tables) < self.vocab_size and self.heap: taken = self._take() if taken is None: continue count, key, stat = taken if not stat.certified: batch = [key] while self.heap and len(batch) < CERTIFY_BATCH: if self._current_certified(self.heap[0]): break following = self._take() if following is not None: batch.append(following[1]) self._certify(batch) continue if count != stat.exact or stat.exact < self.min_count: continue left, right, offset = _unpack(key) if self.tables.compose(left, right, offset) is None: stat.exact = 0 continue child = self.tables.add_merge(left, right, offset) self._merges.append((left, right, offset)) stat.retired = True self.epoch += 1 for dirty in self._apply(left, right, offset, child): other = self.stats[dirty] if not other.retired and other.bound > 0: if not other.live or other.bound > other.queued_count: self._queue(dirty, other, other.bound) return dict( initial_instances=self.initial_instances, final_instances=sum(st.alive), initial_rule_types=initial_types, merges=len(self._merges), **self.counters, ) def _spatial_instances(vocabulary: SparseVocabulary, frames: np.ndarray): bins = np.asarray(vocabulary.layout.sparse_bins, dtype=np.int64) + 1 if float(np.prod(bins.astype(float))) < 2.0**62: keys = np.ravel_multi_index((frames.astype(np.int64) + 1).T, bins) unique, inverse = np.unique(keys, return_inverse=True) rows = np.stack(np.unravel_index(unique, bins), axis=1) - 1 else: rows, inverse = np.unique(frames.astype(np.int64), axis=0, return_inverse=True) inverse = np.asarray(inverse).reshape(-1) nonempty = np.flatnonzero(np.any(rows >= 0, axis=1)) counts = np.zeros(len(rows), dtype=np.int64) codec = vocabulary._codec() if codec is not None: encoded_counts, flat = codec.spatial_encode(np.asarray(rows[nonempty], dtype=np.int32)) counts[nonempty] = encoded_counts flat = np.asarray(flat, dtype=np.int32) else: parts = [vocabulary.spatial.encode(rows[i].tolist()) for i in nonempty] counts[nonempty] = [len(p) for p in parts] flat = np.asarray([t for p in parts for t in sorted(p)], dtype=np.int32) offsets = np.concatenate(([0], np.cumsum(counts))) per_frame = counts[inverse] total = int(per_frame.sum()) position = np.arange(total) - np.repeat(np.cumsum(per_frame) - per_frame, per_frame) return per_frame, flat[np.repeat(offsets[:-1][inverse], per_frame) + position] def fit_sparse_vocabulary( layout: SparseLayout, chunks: Sequence[np.ndarray], *, spatial_vocab_size: int, vocab_size: int, lookahead: int, backend: str = "auto", threads: int = 0, progress: bool = False, ) -> tuple[SparseVocabulary, dict]: """Fit spatial then temporal sparse Set-BPE on independent dense chunks ``[(T_i, D)]``.""" backend = resolve_backend(backend) started = time.perf_counter() frames, starts = sparsify_chunks(layout, chunks) atoms = sum(layout.sparse_bins) if spatial_vocab_size < atoms: raise ValueError(f"spatial_vocab_size must be at least the {atoms} sparse atoms") spatial_merges, spatial_report = (), {} if spatial_vocab_size > atoms: records = frames[np.any(frames >= 0, axis=1)] codes, weights = aggregate_records(records, layout.sparse_bins) spatial, spatial_report = train_spatial( codes, weights, layout.sparse_bins, vocab_size=spatial_vocab_size, backend=backend, threads=threads, progress=progress, ) spatial_merges = spatial.merges spatial_done = time.perf_counter() vocabulary = SparseVocabulary(layout, spatial_merges, (), lookahead) per_frame, tokens = _spatial_instances(vocabulary, frames) chunk = np.repeat(np.arange(len(starts) - 1), np.diff(starts)) anchors = np.repeat(np.arange(len(frames), dtype=np.int64) + chunk * lookahead, per_frame) support = [int(s) for s in vocabulary.spatial.support] arguments = (support, layout.dims + 1, lookahead, vocab_size, True, MIN_PAIR_COUNT) if backend == "native": trainer = native_kernel("SparseTrainer")( *arguments, threads, bool(getattr(layout, "dense", False)) ) else: trainer = ReferenceTrainer(*arguments) trainer.set_instances(anchors, tokens.astype(np.int32), len(frames) + len(chunks) * lookahead) del anchors temporal_report = dict(trainer.train(progress)) vocabulary = SparseVocabulary(layout, spatial_merges, trainer.merges(), lookahead) report = dict( chunks=len(chunks), frames=len(frames), sparse_atoms=atoms, atoms_per_frame=float(np.count_nonzero(frames >= 0) / len(frames)), spatial_vocab_size=vocabulary.spatial.vocab_size, vocab_size=vocabulary.vocab_size, spatial_instances=len(tokens), spatial=dict(spatial_report), temporal=temporal_report, seconds=dict(spatial=spatial_done - started, total=time.perf_counter() - started), ) return vocabulary, report