byte_bpe.py

"""Constrained byte-fallback BPE vocabulary for TinyReceiptVQA.

The runtime in this module is deliberately implemented in pure Python.  The
optional ``tokenizers`` package is used only by :meth:`ByteFallbackBPE.build`
to make training fast; loading, encoding, decoding, and JSON serialization do
not depend on it.

This tokenizer differs from a general-purpose BPE in three useful ways for
receipt OCR:

* the structured-output tags and the ten ASCII digits are always atomic;
* whitespace is a hard BPE boundary;
* every UTF-8 byte has a reserved ``<0xNN>`` token, so unseen Unicode text
  never has to become ``<unk>``.

Encoding normalizes text to NFC.  Consequently, the round-trip guarantee is
``decode(encode(text)) == unicodedata.normalize("NFC", text)`` for valid
Unicode input.
"""
from __future__ import annotations

import hashlib
import json
import unicodedata
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import Any


DEFAULT_VOCAB_SIZE = 1024 + 512
DEFAULT_ALPHABET_SIZE = 768
BPE_TYPE = "byte_fallback_bpe"
BPE_VERSION = 1

SPECIAL_TOKENS = ("<pad>", "<bos>", "<eos>", "<unk>")
STRUCTURAL_TOKENS = (
    "<field>",
    "</field>",
    "<value>",
    "</value>",
    "<op>",
    "</op>",
    "<answer>",
    "</answer>",
)
DIGIT_TOKENS = tuple("0123456789")
ATOMIC_TOKENS = STRUCTURAL_TOKENS + DIGIT_TOKENS
BYTE_TOKENS = tuple(f"<0x{value:02X}>" for value in range(256))

# Unicode White_Space, kept explicit so Python and browser implementations
# use exactly the same boundary predicate.  In particular, this intentionally
# excludes U+FEFF and Python's legacy U+001C..U+001F isspace() characters.
_WHITESPACE_CODEPOINTS = frozenset(
    (
        0x0009,
        0x000A,
        0x000B,
        0x000C,
        0x000D,
        0x0020,
        0x0085,
        0x00A0,
        0x1680,
        0x2028,
        0x2029,
        0x202F,
        0x205F,
        0x3000,
    )
    + tuple(range(0x2000, 0x200B))
)


def _is_whitespace(character: str) -> bool:
    return ord(character) in _WHITESPACE_CODEPOINTS


def _normalized(text: object) -> str:
    return unicodedata.normalize("NFC", str(text))


def _iter_bpe_spans(text: str) -> Iterable[str]:
    """Yield only spans in which BPE merges are permitted."""
    text = _normalized(text)
    tags = sorted(STRUCTURAL_TOKENS, key=len, reverse=True)
    start = 0
    cursor = 0
    while cursor < len(text):
        tag = next((tag for tag in tags if text.startswith(tag, cursor)), None)
        boundary_length = len(tag) if tag is not None else 0
        if not boundary_length and (
            text[cursor] in DIGIT_TOKENS or _is_whitespace(text[cursor])
        ):
            boundary_length = 1
        if boundary_length:
            if start < cursor:
                yield text[start:cursor]
            cursor += boundary_length
            start = cursor
        else:
            cursor += 1
    if start < len(text):
        yield text[start:]


def _fingerprint_payload(payload: Mapping[str, Any]) -> str:
    unhashed = dict(payload)
    unhashed.pop("tokenizer_hash", None)
    encoded = json.dumps(
        unhashed,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


class ByteFallbackBPE:
    """NFC, structured-output-aware byte-fallback BPE vocabulary."""

    def __init__(
        self,
        itos: Sequence[str],
        merges: Sequence[Sequence[str]],
        *,
        atomic_tokens: Sequence[str] = ATOMIC_TOKENS,
        byte_tokens: Sequence[str] = BYTE_TOKENS,
        unused_tokens: Sequence[str] = (),
    ) -> None:
        self.itos = list(itos)
        self.stoi = {token: index for index, token in enumerate(self.itos)}
        self.merges = [(str(pair[0]), str(pair[1])) for pair in merges]
        self.atomic_tokens = tuple(atomic_tokens)
        self.byte_tokens = tuple(byte_tokens)
        self.unused_tokens = tuple(unused_tokens)

        self._atomic_set = frozenset(self.atomic_tokens)
        self._byte_set = frozenset(self.byte_tokens)
        self._unused_set = frozenset(self.unused_tokens)
        self._tags_longest_first = tuple(
            sorted(
                (token for token in self.atomic_tokens if len(token) > 1),
                key=len,
                reverse=True,
            )
        )
        self._merge_ranks = {
            (left, right): rank for rank, (left, right) in enumerate(self.merges)
        }
        self._validate()

    @classmethod
    def build(
        cls,
        records: Iterable[Mapping[str, Any]],
        *,
        vocab_size: int = DEFAULT_VOCAB_SIZE,
        min_frequency: int = 2,
        alphabet_size: int = DEFAULT_ALPHABET_SIZE,
        text_fields: Sequence[str] = ("question", "target"),
    ) -> "ByteFallbackBPE":
        """Train from the records explicitly supplied by the caller.

        Callers should pass training records only.  This method never discovers
        or reads validation/heldout data on its own.
        """
        reserved = list(SPECIAL_TOKENS + ATOMIC_TOKENS + BYTE_TOKENS)
        if vocab_size < len(reserved):
            raise ValueError(
                f"vocab_size={vocab_size} is smaller than the "
                f"{len(reserved)} required special/atomic/byte tokens"
            )
        if min_frequency < 1:
            raise ValueError("min_frequency must be at least 1")
        maximum_alphabet = vocab_size - len(reserved)
        if not 1 <= alphabet_size <= maximum_alphabet:
            raise ValueError(
                f"alphabet_size must be between 1 and {maximum_alphabet}, "
                f"got {alphabet_size}"
            )
        if not text_fields:
            raise ValueError("text_fields must not be empty")

        try:
            from tokenizers import Tokenizer, models, normalizers, trainers
        except ImportError as exc:  # pragma: no cover - environment dependent
            raise RuntimeError(
                "ByteFallbackBPE.build() requires the optional 'tokenizers' "
                "package; install it for training. Runtime loading and "
                "encode/decode do not require it."
            ) from exc

        def training_spans() -> Iterable[str]:
            for record in records:
                if not isinstance(record, Mapping):
                    raise TypeError(
                        "each training record must be a mapping containing "
                        f"{tuple(text_fields)!r}"
                    )
                for field in text_fields:
                    value = record.get(field)
                    if value is None:
                        continue
                    yield from _iter_bpe_spans(str(value))

        tokenizer = Tokenizer(
            models.BPE(unk_token="<unk>", byte_fallback=True)
        )
        tokenizer.normalizer = normalizers.NFC()
        trainer = trainers.BpeTrainer(
            vocab_size=vocab_size,
            min_frequency=min_frequency,
            show_progress=False,
            special_tokens=reserved,
            limit_alphabet=alphabet_size,
        )
        tokenizer.train_from_iterator(training_spans(), trainer=trainer)

        trained = json.loads(tokenizer.to_str())
        model_data = trained["model"]
        vocab_by_token = {
            str(token): int(index)
            for token, index in model_data["vocab"].items()
        }
        ordered = sorted(vocab_by_token.items(), key=lambda item: item[1])
        if any(index != expected for expected, (_, index) in enumerate(ordered)):
            raise ValueError("trainer returned a non-contiguous vocabulary")
        itos = [token for token, _ in ordered]

        merges: list[tuple[str, str]] = []
        for raw_pair in model_data.get("merges") or []:
            if isinstance(raw_pair, str):
                pair = raw_pair.split(" ", 1)
            else:
                pair = list(raw_pair)
            if len(pair) != 2:
                raise ValueError(f"invalid BPE merge from trainer: {raw_pair!r}")
            left, right = str(pair[0]), str(pair[1])
            merged = left + right
            # Training spans already exclude these boundaries.  Keep this
            # defensive filter so a future trainer/pre-tokenizer change cannot
            # silently weaken the OCR constraints.
            if (
                left in BYTE_TOKENS
                or right in BYTE_TOKENS
                or left in ATOMIC_TOKENS
                or right in ATOMIC_TOKENS
                or any(char in DIGIT_TOKENS for char in merged)
                or any(_is_whitespace(char) for char in merged)
            ):
                continue
            if merged not in vocab_by_token:
                raise ValueError(
                    f"merge output {merged!r} is absent from trained vocabulary"
                )
            merges.append((left, right))

        unused_tokens: list[str] = []
        occupied = set(itos)
        counter = 0
        while len(itos) < vocab_size:
            candidate = f"<unused_{counter:04d}>"
            counter += 1
            if candidate in occupied:
                continue
            occupied.add(candidate)
            unused_tokens.append(candidate)
            itos.append(candidate)
        if len(itos) != vocab_size:
            raise ValueError(
                f"trainer returned {len(itos)} tokens for vocab_size={vocab_size}"
            )

        return cls(
            itos,
            merges,
            atomic_tokens=ATOMIC_TOKENS,
            byte_tokens=BYTE_TOKENS,
            unused_tokens=unused_tokens,
        )

    # A name that reads naturally at integration call sites.
    train_from_records = build

    def _validate(self) -> None:
        if len(self.itos) != len(self.stoi):
            raise ValueError("itos contains duplicate token strings")
        if tuple(self.itos[: len(SPECIAL_TOKENS)]) != SPECIAL_TOKENS:
            raise ValueError(
                f"vocabulary must begin with {SPECIAL_TOKENS!r}; pad must be id 0"
            )
        if self.byte_tokens != BYTE_TOKENS:
            raise ValueError("byte_tokens must contain <0x00> through <0xFF>")
        missing = [
            token
            for token in SPECIAL_TOKENS + self.atomic_tokens + self.byte_tokens
            if token not in self.stoi
        ]
        if missing:
            raise ValueError(f"vocabulary is missing required tokens: {missing!r}")
        if len(self._merge_ranks) != len(self.merges):
            raise ValueError("duplicate merge pairs are not allowed")
        if not self._unused_set.issubset(self.stoi):
            raise ValueError("unused_tokens contains a token absent from itos")

        for left, right in self.merges:
            merged = left + right
            if left not in self.stoi or right not in self.stoi:
                raise ValueError(f"merge input is absent from vocabulary: {(left, right)!r}")
            if merged not in self.stoi:
                raise ValueError(f"merge output is absent from vocabulary: {merged!r}")
            if left in self._byte_set or right in self._byte_set:
                raise ValueError("byte fallback tokens must never participate in merges")
            if left in self._atomic_set or right in self._atomic_set:
                raise ValueError("atomic tokens must never participate in merges")
            if any(char in DIGIT_TOKENS for char in merged):
                raise ValueError("digits must never participate in merges")
            if any(_is_whitespace(char) for char in merged):
                raise ValueError("whitespace must never participate in merges")

    @property
    def pad(self) -> int:
        return self.stoi["<pad>"]

    @property
    def bos(self) -> int:
        return self.stoi["<bos>"]

    @property
    def eos(self) -> int:
        return self.stoi["<eos>"]

    @property
    def unk(self) -> int:
        return self.stoi["<unk>"]

    @property
    def vocab_size(self) -> int:
        return len(self.itos)

    def __len__(self) -> int:
        return len(self.itos)

    def _fallback_ids(self, character: str) -> list[int]:
        return [
            self.stoi[f"<0x{byte:02X}>"]
            for byte in character.encode("utf-8", errors="strict")
        ]

    def _apply_bpe(self, initial: list[str]) -> list[str]:
        tokens = initial
        while len(tokens) > 1:
            best_pair: tuple[str, str] | None = None
            best_rank: int | None = None
            for index in range(len(tokens) - 1):
                pair = (tokens[index], tokens[index + 1])
                rank = self._merge_ranks.get(pair)
                if rank is not None and (best_rank is None or rank < best_rank):
                    best_pair = pair
                    best_rank = rank
            if best_pair is None:
                break

            merged: list[str] = []
            index = 0
            while index < len(tokens):
                if (
                    index + 1 < len(tokens)
                    and (tokens[index], tokens[index + 1]) == best_pair
                ):
                    merged.append(tokens[index] + tokens[index + 1])
                    index += 2
                else:
                    merged.append(tokens[index])
                    index += 1
            tokens = merged
        return tokens

    def _encode_mergeable_span(self, span: str) -> list[int]:
        initial: list[str] = []
        for character in span:
            if (
                character in self.stoi
                and character not in self._byte_set
                and character not in self._unused_set
            ):
                initial.append(character)
            else:
                initial.extend(
                    f"<0x{byte:02X}>"
                    for byte in character.encode("utf-8", errors="strict")
                )
        return [self.stoi[token] for token in self._apply_bpe(initial)]

    def encode(
        self,
        text: str,
        add_bos: bool = False,
        add_eos: bool = False,
        max_len: int = 0,
    ) -> list[int]:
        normalized = _normalized(text)
        ids: list[int] = []
        span_start = 0
        cursor = 0
        while cursor < len(normalized):
            tag = next(
                (
                    token
                    for token in self._tags_longest_first
                    if normalized.startswith(token, cursor)
                ),
                None,
            )
            is_digit = normalized[cursor] in DIGIT_TOKENS
            is_space = _is_whitespace(normalized[cursor])
            if tag is None and not is_digit and not is_space:
                cursor += 1
                continue

            if span_start < cursor:
                ids.extend(self._encode_mergeable_span(normalized[span_start:cursor]))
            if tag is not None:
                ids.append(self.stoi[tag])
                cursor += len(tag)
            elif is_digit:
                ids.append(self.stoi[normalized[cursor]])
                cursor += 1
            else:
                character = normalized[cursor]
                if character in self.stoi and character not in self._byte_set:
                    ids.append(self.stoi[character])
                else:
                    ids.extend(self._fallback_ids(character))
                cursor += 1
            span_start = cursor

        if span_start < len(normalized):
            ids.extend(self._encode_mergeable_span(normalized[span_start:]))
        if add_bos:
            ids.insert(0, self.bos)
        if add_eos:
            ids.append(self.eos)
        if max_len:
            ids = ids[:max_len]
            if add_eos and ids[-1] != self.eos:
                ids[-1] = self.eos
        return ids

    def decode(self, ids: Iterable[int], *, errors: str = "replace") -> str:
        output = bytearray()
        for raw_id in ids:
            token_id = int(raw_id)
            if token_id == self.eos:
                break
            if token_id in (self.pad, self.bos):
                continue
            if not 0 <= token_id < len(self.itos):
                continue
            token = self.itos[token_id]
            if token in self._unused_set:
                continue
            if token in self._byte_set:
                output.append(int(token[3:5], 16))
            else:
                output.extend(token.encode("utf-8", errors="strict"))
        return output.decode("utf-8", errors=errors)

    def to_json(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "type": BPE_TYPE,
            "version": BPE_VERSION,
            "vocab_size": len(self.itos),
            "itos": list(self.itos),
            "merges": [[left, right] for left, right in self.merges],
            "normalization": "NFC",
            "atomic_tokens": list(self.atomic_tokens),
            "byte_tokens": list(self.byte_tokens),
            "unused_tokens": list(self.unused_tokens),
            "special_tokens": {
                "pad": "<pad>",
                "bos": "<bos>",
                "eos": "<eos>",
                "unk": "<unk>",
            },
        }
        payload["tokenizer_hash"] = _fingerprint_payload(payload)
        return payload

    @classmethod
    def from_json(cls, obj: Mapping[str, Any]) -> "ByteFallbackBPE":
        if not is_byte_fallback_bpe_json(obj):
            raise ValueError("expected a byte_fallback_bpe version 1 vocabulary")
        if obj.get("normalization") != "NFC":
            raise ValueError("byte_fallback_bpe version 1 requires NFC normalization")

        expected_hash = obj.get("tokenizer_hash")
        if expected_hash is not None and expected_hash != _fingerprint_payload(obj):
            raise ValueError("tokenizer_hash does not match vocabulary contents")

        tokenizer = cls(
            obj["itos"],
            obj.get("merges") or (),
            atomic_tokens=obj.get("atomic_tokens") or ATOMIC_TOKENS,
            byte_tokens=obj.get("byte_tokens") or BYTE_TOKENS,
            unused_tokens=obj.get("unused_tokens") or (),
        )
        declared_size = int(obj.get("vocab_size", -1))
        if declared_size != len(tokenizer):
            raise ValueError(
                f"declared vocab_size={declared_size} but itos has "
                f"{len(tokenizer)} entries"
            )
        return tokenizer

    def save(self, path: str | Path) -> None:
        Path(path).write_text(
            json.dumps(self.to_json(), ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

    @classmethod
    def load(cls, path: str | Path) -> "ByteFallbackBPE":
        return cls.from_json(json.loads(Path(path).read_text(encoding="utf-8")))


# Slightly more explicit alias for call sites that use "Vocab" terminology.
ByteFallbackBPEVocab = ByteFallbackBPE


def is_byte_fallback_bpe_json(obj: object) -> bool:
    return (
        isinstance(obj, Mapping)
        and obj.get("type") == BPE_TYPE
        and int(obj.get("version", 0)) == BPE_VERSION
    )


def load_bpe1536_vocab(obj: Mapping[str, Any]) -> ByteFallbackBPEVocab:
    """Load the single tokenizer contract supported by release tooling."""
    tokenizer = ByteFallbackBPEVocab.from_json(obj)
    tokenizer_hash = obj.get("tokenizer_hash")
    if (
        not isinstance(tokenizer_hash, str)
        or len(tokenizer_hash) != 64
        or tokenizer_hash != _fingerprint_payload(obj)
    ):
        raise ValueError("BPE1536 release vocabulary requires a valid tokenizer_hash")
    if len(tokenizer) != DEFAULT_VOCAB_SIZE:
        raise ValueError(
            f"release tokenizer must contain {DEFAULT_VOCAB_SIZE} tokens, "
            f"got {len(tokenizer)}"
        )
    return tokenizer


def bpe1536_contract(obj: Mapping[str, Any]) -> dict[str, Any]:
    """Return the manifest fields after fully validating a release vocabulary."""
    tokenizer = load_bpe1536_vocab(obj)
    return {
        "type": BPE_TYPE,
        "version": BPE_VERSION,
        "vocab_size": len(tokenizer),
        "normalization": "NFC",
        "tokenizer_hash": str(obj.get("tokenizer_hash", "")),
    }