| """Strict ``bpe.json`` I/O for both tokenization schemes. |
| |
| The serialization layer owns only the artifact schema. Set-BPE spatial and temporal |
| merge semantics stay in :mod:`actioncodec2.setbpe` and plain-BPE merge |
| semantics in :mod:`actioncodec2.tokenization.bpe_vocabulary`, so training and deployment |
| cannot drift through two independent implementations of the same algorithm. |
| |
| ``config.set_bpe.scheme`` selects which serializer an artifact carries: |
| ``"set_bpe"`` (the production default) or ``"bpe"``, the FAST-style baseline |
| that flattens the quantized cell matrix row-major before merging. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from collections.abc import Mapping, Sequence |
| from numbers import Integral |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from ..common._validation import active_dimensions |
| from .bpe_vocabulary import BPEVocabulary |
| from ..physical.configuration import ActionCodec2Config, ConfigError |
| from .setbpe_spatial import SpatialVocabulary |
| from .setbpe_temporal import TemporalVocabulary |
|
|
| BPE_NAME = "bpe.json" |
|
|
|
|
| def _integer(value: object, path: str, *, positive: bool = False) -> int: |
| if isinstance(value, bool) or not isinstance(value, int): |
| raise ConfigError(f"{path} must be an integer") |
| minimum = 1 if positive else 0 |
| if value < minimum: |
| qualifier = "positive" if positive else "non-negative" |
| raise ConfigError(f"{path} must be {qualifier}") |
| return int(value) |
|
|
|
|
| def _rules(payload: object, width: int, path: str) -> tuple[tuple[int, ...], ...]: |
| """Parse an ordered merge list without accepting lossy coercions.""" |
| 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: |
| """Give artifact-level errors before constructing vocabulary metadata.""" |
| 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]: |
| """Parse ``bpe.json`` from an artifact directory, rejecting duplicate keys.""" |
| path = Path(directory) / BPE_NAME |
| try: |
| return json.loads( |
| path.read_text(encoding="utf-8"), object_pairs_hook=_unique_json_object |
| ) |
| except OSError: |
| raise |
| 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, |
| ) -> tuple[np.ndarray, tuple[int, ...] | None]: |
| """Validate one trajectory and expand it to dense ``[T,D_max]`` cells. |
| |
| Compact ``[T,D_active]`` input may provide explicit global ``active_dims``. |
| Without it, compact input uses prefix dimensions. Dense ``[T,D_max]`` input |
| may mark absent slots with ``-1``. |
| """ |
| 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 != 2 or len(values) < 1 or values.shape[1] < 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((len(values), 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((len(values), 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 SetBPESerializer: |
| """Lossless mapping between ``[T,D]`` physical cells and token IDs. |
| |
| Merge lists are interpreted in creation order. Token IDs are therefore |
| stable across physical save/load cycles, and decoding can infer any positive |
| trajectory length directly from stream coverage. |
| """ |
|
|
| scheme = "set_bpe" |
|
|
| def __init__( |
| self, |
| *, |
| bins: Sequence[int], |
| lookahead: int, |
| spatial_merges: Sequence[Sequence[int]], |
| temporal_merges: Sequence[Sequence[int]], |
| ) -> None: |
| 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") |
| 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, |
| ) |
|
|
| def encode( |
| self, |
| cells: np.ndarray, |
| *, |
| active_dims: Sequence[int] | None = None, |
| ) -> list[int]: |
| """Serialize one non-empty trajectory without emitting absent slots. |
| |
| Compact ``[T,D_active]`` input may provide explicit global |
| ``active_dims``. Without it, compact input uses prefix dimensions. |
| Dense ``[T,D_max]`` input may mark absent slots with ``-1``. |
| """ |
| values, _ = _dense_cells(cells, self.bins, active_dims) |
| 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) |
| return self.vocabulary.decode(checked, active_dims=active_dims) |
|
|
| def to_dict(self) -> dict[str, object]: |
| """Return the intentionally minimal, path-independent BPE artifact.""" |
| return { |
| "spatial_merges": [list(rule) for rule in self.spatial_merges], |
| "temporal_merges": [list(rule) for rule in self.temporal_merges], |
| } |
|
|
| @classmethod |
| def from_dict( |
| cls, |
| payload: object, |
| *, |
| bins: Sequence[int], |
| lookahead: int, |
| ) -> SetBPESerializer: |
| 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 = {"spatial_merges", "temporal_merges"} |
| unknown = sorted(set(payload) - allowed) |
| if unknown: |
| raise ConfigError( |
| f"{BPE_NAME} contains unknown field(s): {', '.join(unknown)}" |
| ) |
| missing = sorted(allowed - set(payload)) |
| if missing: |
| raise ConfigError(f"{BPE_NAME} is missing field(s): {', '.join(missing)}") |
| return cls( |
| bins=bins, |
| lookahead=lookahead, |
| spatial_merges=payload["spatial_merges"], |
| temporal_merges=payload["temporal_merges"], |
| ) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| directory: str | Path, |
| *, |
| bins: Sequence[int], |
| lookahead: int, |
| ) -> SetBPESerializer: |
| 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 BPESerializer: |
| """Lossless mapping between ``[T,D]`` physical cells and plain-BPE ids. |
| |
| The FAST-style baseline: cells are serialized row-major over the active |
| support and merged as an ordinary symbol sequence. It shares the primitive |
| atoms, the artifact layout, and the total vocabulary budget with |
| :class:`SetBPESerializer`, so the two can be fitted on one corpus and |
| compared directly. |
| |
| ``lookahead`` bounds a token at ``lookahead * len(bins)`` primitive cells, |
| the plain-BPE analogue of the Set-BPE rule that a temporal token spans at |
| most ``lookahead`` control steps. |
| """ |
|
|
| scheme = "bpe" |
|
|
| def __init__( |
| self, |
| *, |
| bins: Sequence[int], |
| lookahead: int, |
| merges: Sequence[Sequence[int]], |
| ) -> None: |
| 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") |
| 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 |
|
|
| def encode( |
| self, |
| cells: np.ndarray, |
| *, |
| active_dims: Sequence[int] | None = None, |
| ) -> list[int]: |
| """Serialize one non-empty trajectory without emitting absent slots. |
| |
| Input conventions match :meth:`SetBPESerializer.encode`: compact |
| ``[T,D_active]`` may carry explicit global ``active_dims`` and |
| otherwise uses prefix dimensions, while dense ``[T,D_max]`` marks |
| absent slots with ``-1``. |
| """ |
| values, _ = _dense_cells(cells, self.bins, active_dims) |
| 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) |
| return self.vocabulary.decode(checked, active_dims=active_dims) |
|
|
| def to_dict(self) -> dict[str, object]: |
| """Return the intentionally minimal, path-independent BPE artifact.""" |
| return {"merges": [list(rule) for rule in self.merges]} |
|
|
| @classmethod |
| def from_dict( |
| cls, |
| payload: object, |
| *, |
| bins: Sequence[int], |
| lookahead: int, |
| ) -> BPESerializer: |
| 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") |
| unknown = sorted(set(payload) - {"merges"}) |
| if unknown: |
| raise ConfigError( |
| f"{BPE_NAME} contains unknown field(s): {', '.join(unknown)}" |
| ) |
| if "merges" not in payload: |
| raise ConfigError(f"{BPE_NAME} is missing field(s): merges") |
| return cls(bins=bins, lookahead=lookahead, merges=payload["merges"]) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| directory: str | Path, |
| *, |
| bins: Sequence[int], |
| lookahead: int, |
| ) -> BPESerializer: |
| 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 |
|
|
|
|
| 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.""" |
| return SERIALIZERS[config.set_bpe.scheme] |
|
|
