| """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)) |
|
|
| |
| |
| |
| _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: |
| 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 |
| |
| |
| |
| 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, |
| ) |
|
|
| |
| 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"))) |
|
|
|
|
| |
| 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", "")), |
| } |
|
|
