Blame for runtime/backends/actioncodec2/setbpe/temporal_training.py

'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)

    @property
    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