'Strict ``bpe.json`` I/O for both tokenization schemes.' from __future__ import annotations import json from collections.abc import Mapping, Sequence from numbers import Integral from pathlib import Path from typing import TypeVar import numpy as np from ._validation import active_dimensions from .bpe_vocabulary import BPEVocabulary from .configuration import ActionCodec2Config, ConfigError, _integer from .native_serialization import NativeCacheMixin, decode_native, encode_native from .packing import EncodingSearch from .setbpe.spatial import SpatialVocabulary from .setbpe.temporal import TemporalVocabulary BPE_NAME = "bpe.json" SerializerT = TypeVar("SerializerT", bound="_Serializer") def _rules(payload: object, width: int, path: str) -> tuple[tuple[int, ...], ...]: if not isinstance(payload, Sequence) or isinstance(payload, (str, bytes, bytearray)): raise ConfigError(f"{path} must be a list") output = [] for rank, raw_rule in enumerate(payload): rule_path = f"{path}[{rank}]" if not isinstance(raw_rule, Sequence) or isinstance(raw_rule, (str, bytes, bytearray)): raise ConfigError(f"{rule_path} must be a list of {width} integers") if len(raw_rule) != width: raise ConfigError(f"{rule_path} must contain exactly {width} integers") output.append( tuple(_integer(value, f"{rule_path}[{index}]") for index, value in enumerate(raw_rule)) ) return tuple(output) def _unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: output: dict[str, object] = {} for key, value in pairs: if key in output: raise ConfigError(f"{BPE_NAME} contains duplicate field {key!r}") output[key] = value return output def _validate_rule_references( *, atom_count: int, lookahead: int, spatial_merges: tuple[tuple[int, int], ...], temporal_merges: tuple[tuple[int, int, int], ...], ) -> None: for rank, (left, right) in enumerate(spatial_merges): child = atom_count + rank if left >= child or right >= child: raise ConfigError(f"spatial_merges[{rank}] must reference already-created tokens") spatial_size = atom_count + len(spatial_merges) for rank, (left, right, offset) in enumerate(temporal_merges): child = spatial_size + rank if left >= child or right >= child: raise ConfigError(f"temporal_merges[{rank}] must reference already-created tokens") if not 1 <= offset < lookahead: raise ConfigError( f"temporal_merges[{rank}] offset must be positive and smaller than lookahead" ) def _read_payload(directory: str | Path) -> dict[str, object]: path = Path(directory) / BPE_NAME try: return json.loads(path.read_text(encoding="utf-8"), object_pairs_hook=_unique_json_object) except json.JSONDecodeError as error: raise ConfigError(f"cannot parse {path}: {error}") from error def _dense_cells( cells: np.ndarray, bins: Sequence[int], active_dims: Sequence[int] | None, *, ndim: int = 2, ) -> tuple[np.ndarray, tuple[int, ...] | None]: raw = np.asarray(cells) if raw.dtype.kind not in {"i", "u"}: raise ValueError("cells must contain integer indices") values = np.asarray(raw, dtype=np.int64) if values.ndim != ndim or min(values.shape) < 1: raise ValueError("cells must have non-empty shape [T, D_active]") if active_dims is None: if values.shape[-1] > len(bins): raise ValueError(f"cells width exceeds codebook dimension {len(bins)}") if values.shape[-1] < len(bins): dimensions = tuple(range(values.shape[-1])) dense = np.full((*values.shape[:-1], len(bins)), -1, dtype=np.int64) dense[..., : values.shape[-1]] = values values = dense else: dimensions = None else: dimensions = active_dimensions(active_dims, len(bins)) if len(dimensions) != values.shape[-1]: raise ValueError("active_dims must have one entry per compact cell coordinate") if np.any(values < 0): raise ValueError("compact cells must not contain -1") dense = np.full((*values.shape[:-1], len(bins)), -1, dtype=np.int64) dense[..., dimensions] = values values = dense counts = np.asarray(tuple(int(value) for value in bins), dtype=np.int64) if np.any(values < -1) or np.any(values >= counts[None, :]): raise ValueError("cells contain an invalid active or absent index") return values, dimensions class _Serializer(NativeCacheMixin): _merge_fields: tuple[str, ...] def _init_dimensions(self, bins: Sequence[int], lookahead: int) -> None: self.encoding_search = EncodingSearch() self.lookahead = _integer(lookahead, "lookahead", positive=True) self.bins = tuple( _integer(value, f"bins[{index}]", positive=True) for index, value in enumerate(bins) ) if not self.bins: raise ConfigError("bins must contain at least one cell slot") def configure_encoding( self, *, spatial_search: bool = False, temporal_search: bool = False, search_window: int = 32, search_beam: int = 4, ): 'Select independent runtime packing stages; both default to disabled.' options = EncodingSearch(spatial_search, temporal_search, search_window, search_beam) options.validate_codec(self.scheme, self.bins) self.encoding_search = options return self def encode_batch(self, cells: np.ndarray, *, active_dims=None) -> list[list[int]]: values = np.asarray(cells) if values.ndim != 3 or min(values.shape) < 1: raise ValueError("cells must have non-empty shape [B,T,D]") dense, _ = _dense_cells(values, self.bins, active_dims, ndim=3) result = encode_native(self, dense) return ( result if result is not None else [self.encode(row, active_dims=active_dims) for row in values] ) def decode_batch(self, streams, *, active_dims=None) -> np.ndarray: dims = active_dimensions(active_dims, len(self.bins)) result = decode_native(self, streams, dims, validate=True) if result is not None: return result rows = [] for stream in streams: row = [] for token in stream: if isinstance(token, bool) or not isinstance(token, Integral): raise ValueError("tokens must be integers") if not 0 <= int(token) < self.vocab_size: raise ValueError("token is outside the vocabulary") row.append(int(token)) if not row: raise ValueError("token stream must not be empty") rows.append(row) if not rows: raise ValueError("streams must not be empty") result = decode_native(self, rows, dims) if result is not None: return result cells = [self.decode(row, active_dims=active_dims) for row in rows] if len({len(row) for row in cells}) != 1: raise ValueError("all token streams must decode to the same length") return np.stack(cells) def encode( self, cells: np.ndarray, *, active_dims: Sequence[int] | None = None, ) -> list[int]: 'Serialize one non-empty trajectory without emitting absent slots.' values, _ = _dense_cells(cells, self.bins, active_dims) result = encode_native(self, values[None]) if result is not None: return result[0] return self.vocabulary.encode(values) def decode( self, tokens: Sequence[int], *, active_dims: Sequence[int] | None = None, ) -> np.ndarray: """Recover dense cells, writing ``-1`` in semantically absent slots.""" if len(tokens) < 1: raise ValueError("token stream must not be empty") checked = [] for index, token in enumerate(tokens): if isinstance(token, bool) or not isinstance(token, Integral): raise ValueError(f"tokens[{index}] must be an integer") value = int(token) if not 0 <= value < self.vocab_size: raise ValueError(f"tokens[{index}] is outside the vocabulary") checked.append(value) result = decode_native(self, [checked], active_dimensions(active_dims, len(self.bins))) return ( result[0] if result is not None else self.vocabulary.decode(checked, active_dims=active_dims) ) def to_dict(self) -> dict[str, object]: """Return only the merge lists; physical metadata lives in config.yaml.""" return { field: [list(rule) for rule in getattr(self, field)] for field in self._merge_fields } @classmethod def from_dict( cls: type[SerializerT], payload: object, *, bins: Sequence[int], lookahead: int ) -> SerializerT: """Validate the scheme's exact artifact fields before constructing it.""" if not isinstance(payload, Mapping) or not all(isinstance(key, str) for key in payload): raise ConfigError(f"{BPE_NAME} must contain one JSON object") allowed = set(cls._merge_fields) for fields, message in ( (set(payload) - allowed, "contains unknown field(s)"), (allowed - set(payload), "is missing field(s)"), ): if fields: raise ConfigError(f"{BPE_NAME} {message}: {', '.join(sorted(fields))}") return cls(bins=bins, lookahead=lookahead, **payload) @classmethod def from_pretrained( cls: type[SerializerT], directory: str | Path, *, bins: Sequence[int], lookahead: int, ) -> SerializerT: return cls.from_dict(_read_payload(directory), bins=bins, lookahead=lookahead) def save_pretrained(self, directory: str | Path) -> Path: """Write ``bpe.json`` and return its path.""" path = Path(directory) / BPE_NAME path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8") return path class SetBPESerializer(_Serializer): 'Lossless mapping between ``[T,D]`` physical cells and token IDs.' scheme = "set_bpe" _merge_fields = ("spatial_merges", "temporal_merges") def __init__( self, *, bins: Sequence[int], lookahead: int, spatial_merges: Sequence[Sequence[int]], temporal_merges: Sequence[Sequence[int]], ) -> None: self._init_dimensions(bins, lookahead) self.spatial_merges = _rules(spatial_merges, 2, "spatial_merges") self.temporal_merges = _rules(temporal_merges, 3, "temporal_merges") _validate_rule_references( atom_count=sum(self.bins), lookahead=self.lookahead, spatial_merges=self.spatial_merges, temporal_merges=self.temporal_merges, ) try: self.spatial = SpatialVocabulary(self.bins, self.spatial_merges) self.vocabulary = TemporalVocabulary(self.spatial, self.lookahead, self.temporal_merges) except ValueError as error: raise ConfigError(str(error)) from error self.patterns = self.vocabulary.patterns self.root = self.vocabulary.root self.vocab_size = self.vocabulary.vocab_size self.token_cell_counts = np.asarray( [ sum( int(self.spatial.support[spatial_token]).bit_count() for _, spatial_token in self.patterns[token] ) for token in range(self.vocab_size) ], dtype=np.int64, ) class BPESerializer(_Serializer): 'Lossless mapping between ``[T,D]`` physical cells and plain-BPE ids.' scheme = "bpe" _merge_fields = ("merges",) def __init__( self, *, bins: Sequence[int], lookahead: int, merges: Sequence[Sequence[int]], ) -> None: self._init_dimensions(bins, lookahead) self.merges = _rules(merges, 2, "merges") atom_count = sum(self.bins) for rank, (left, right) in enumerate(self.merges): child = atom_count + rank if left >= child or right >= child: raise ConfigError(f"merges[{rank}] must reference already-created tokens") try: self.vocabulary = BPEVocabulary(self.bins, self.merges) except ValueError as error: raise ConfigError(str(error)) from error self.max_token_cells = self.lookahead * len(self.bins) self.token_cell_counts = self.vocabulary.token_cells oversized = int(np.count_nonzero(self.token_cell_counts > self.max_token_cells)) if oversized: raise ConfigError( f"{oversized} merge rule(s) cover more than the {self.max_token_cells} " "primitive cells allowed by lookahead" ) self.vocab_size = self.vocabulary.vocab_size SERIALIZERS: dict[str, type[SetBPESerializer] | type[BPESerializer]] = { SetBPESerializer.scheme: SetBPESerializer, BPESerializer.scheme: BPESerializer, } def serializer_class(config: ActionCodec2Config) -> type[SetBPESerializer] | type[BPESerializer]: """Return the serializer the configured tokenization scheme requires.""" scheme = config.set_bpe.scheme if scheme == "unified_set_bpe": from .unified_set_bpe import UnifiedSetBPESerializer return UnifiedSetBPESerializer if scheme not in SERIALIZERS: raise ConfigError( f"set_bpe.scheme={scheme!r} is not an ActionCodec2 serializer; " "load sparse_set_bpe artifacts with ActionCodec2Sparse" ) return SERIALIZERS[scheme]