'Plain-BPE fitting implementations.' from __future__ import annotations import heapq from collections import defaultdict from collections.abc import Iterable, Iterator, Sequence import numpy as np from .bpe_vocabulary import BPEVocabulary from .native_backend import require_native, resolve_backend __all__ = ["train_bpe", "cell_chunks"] def cell_chunks( episodes: Iterable[np.ndarray], *, max_chunk_rows: int = 1_000_000, ) -> Iterator[tuple[np.ndarray, np.ndarray]]: 'Yield ``(cells, episode_offsets)`` blocks for the native trainer.' chunks = getattr(episodes, "chunks", None) offsets = getattr(episodes, "offsets", None) if getattr(episodes, "_actioncodec2_validated", False) and chunks is not None: for index, (codes, bounds) in enumerate(zip(chunks, offsets, strict=True)): if codes is None: raise RuntimeError("cell corpus chunk has already been consumed") yield np.ascontiguousarray(codes), np.ascontiguousarray(bounds, dtype=np.int64) episodes.release_chunk(index) return pending: list[np.ndarray] = [] lengths: list[int] = [] rows = 0 def flush() -> tuple[np.ndarray, np.ndarray]: block = np.concatenate(pending, axis=0) bounds = np.concatenate(([0], np.cumsum(lengths, dtype=np.int64))) pending.clear() lengths.clear() return np.ascontiguousarray(block), bounds for episode in episodes: values = np.asarray(episode) if values.dtype.kind not in {"i", "u"}: raise ValueError("episodes must contain integer cell indices") if values.ndim != 2 or len(values) < 1: raise ValueError("every episode must have non-empty shape [T, D]") pending.append(values) lengths.append(len(values)) rows += len(values) if rows >= max_chunk_rows: yield flush() rows = 0 if pending: yield flush() def _flatten( chunks: Iterable[tuple[np.ndarray, np.ndarray]], bins: Sequence[int] ) -> tuple[list[list[int]], int]: counts = np.asarray(tuple(int(value) for value in bins), dtype=np.int64) offsets = np.concatenate(([0], np.cumsum(counts))) streams: list[list[int]] = [] actions = 0 for codes, bounds in chunks: values = np.asarray(codes, dtype=np.int64) if values.ndim != 2 or values.shape[1] != len(counts): raise ValueError("cells must have one column per bin count") if np.any(values < -1) or np.any(values >= counts[None, :]): raise ValueError("cells must contain -1 or an in-range cell index") for start, stop in zip(bounds[:-1], bounds[1:], strict=True): episode = values[int(start) : int(stop)] if not len(episode): raise ValueError("every episode must be non-empty") support = episode >= 0 if np.any(support != support[:1]): raise ValueError("active dimensions must remain constant within an episode") present = np.flatnonzero(support[0]) if not len(present): raise ValueError("every control step needs at least one active coordinate") shifted = episode[:, present] + offsets[present][None, :] streams.append(shifted.reshape(-1).tolist()) actions += len(episode) if not streams: raise ValueError("BPE training requires at least one complete trajectory") return streams, actions def _train_bpe_python( chunks: Iterable[tuple[np.ndarray, np.ndarray]], bins: Sequence[int], *, vocab_size: int, max_token_cells: int, progress: bool = False, ) -> tuple[BPEVocabulary, dict[str, object]]: atoms = int(sum(int(value) for value in bins)) if vocab_size < atoms: raise ValueError(f"vocab_size {vocab_size} is below the {atoms} primitive atoms") if max_token_cells < 1: raise ValueError("max_token_cells must be positive") streams, actions = _flatten(chunks, bins) initial_length = sum(len(stream) for stream in streams) length = initial_length counts: dict[tuple[int, int], int] = defaultdict(int) carriers: dict[tuple[int, int], set[int]] = defaultdict(set) for index, stream in enumerate(streams): for pair in zip(stream[:-1], stream[1:], strict=True): counts[pair] += 1 carriers[pair].add(index) heap = [(-value, pair) for pair, value in counts.items()] heapq.heapify(heap) token_cells = [1] * atoms merges: list[tuple[int, int]] = [] oversized: set[tuple[int, int]] = set() observed_cells = 1 while len(merges) + atoms < vocab_size and heap: negative, pair = heapq.heappop(heap) frequency = counts.get(pair, 0) if frequency <= 0 or frequency != -negative or pair in oversized: continue span = token_cells[pair[0]] + token_cells[pair[1]] if span > max_token_cells: oversized.add(pair) continue child = atoms + len(merges) token_cells.append(span) observed_cells = max(observed_cells, span) touched: set[tuple[int, int]] = set() for index in sorted(carriers.pop(pair, ())): stream = streams[index] output: list[int] = [] position = 0 while position < len(stream): if ( position + 1 < len(stream) and stream[position] == pair[0] and stream[position + 1] == pair[1] ): if output: counts[(output[-1], pair[0])] -= 1 touched.add((output[-1], pair[0])) counts[(output[-1], child)] += 1 carriers[(output[-1], child)].add(index) touched.add((output[-1], child)) if position + 2 < len(stream): follower = stream[position + 2] counts[(pair[1], follower)] -= 1 touched.add((pair[1], follower)) counts[(child, follower)] += 1 carriers[(child, follower)].add(index) touched.add((child, follower)) counts[pair] -= 1 length -= 1 output.append(child) position += 2 else: output.append(stream[position]) position += 1 streams[index] = output merges.append(pair) for member in touched: value = counts.get(member, 0) if value > 0: heapq.heappush(heap, (-value, member)) if progress and len(merges) % 256 == 0: print(f" bpe merge {len(merges):>5} tokens/action {length / actions:.4f}", flush=True) vocabulary = BPEVocabulary(tuple(int(value) for value in bins), tuple(merges)) report: dict[str, object] = { "episodes": len(streams), "actions": actions, "atoms": atoms, "merges": len(merges), "vocab_size": vocabulary.vocab_size, "cells": initial_length, "tokens_per_action_initial": initial_length / actions, "tokens_per_action": length / actions, "compression": initial_length / length, "max_token_cells": int(max_token_cells), "max_token_cells_observed": observed_cells, "oversized_pairs_skipped": len(oversized), "trainer": "python_sparse_exact", } return vocabulary, report def train_bpe( episodes: Iterable[np.ndarray], bins: Sequence[int], *, vocab_size: int, max_token_cells: int, backend: str = "auto", progress: bool = False, threads: int = 0, ) -> tuple[BPEVocabulary, dict[str, object]]: 'Fit plain BPE on the row-major flattened primitive-cell stream.' if resolve_backend(backend) == "python": return _train_bpe_python( cell_chunks(episodes), bins, vocab_size=vocab_size, max_token_cells=max_token_cells, progress=progress, ) merges, report = require_native().train_bpe( cell_chunks(episodes), [int(count) for count in bins], int(vocab_size), int(max_token_cells), int(threads), ) vocabulary = BPEVocabulary( tuple(int(count) for count in bins), tuple((int(left), int(right)) for left, right in merges), ) return vocabulary, dict(report)