Implement unk_token, fuse_unk, byte_fallback and ignore_merges

docs/components.md

@@ -37,8 +37,38 @@ byte-mapped space.
 Turns one byte-mapped piece into subword tokens.
 
 - `BPE(vocab, merges, *, unk_token=None, ...)` — byte-pair encoding. `tokenize`
-  greedily applies the lowest-rank applicable merge until none remain (identical
-  to the UbiTok reference runtime).
+  resolves each character onto vocabulary tokens, then greedily applies the
+  lowest-rank applicable merge until none remain (identical to the UbiTok
+  reference runtime).
+
+Resolution happens **before** merging, because the tokens it produces take part
+in merges themselves — a vocabulary containing `<0xE4><0xB8>` plus a merge over
+its two halves tokenizes `中` as `['<0xE4><0xB8>', '<0xAD>']`.
+
+| Option | Effect |
+| --- | --- |
+| `unk_token` | out-of-vocabulary characters become this token. Without it they are **dropped**, matching HuggingFace |
+| `fuse_unk` | consecutive unresolvable characters collapse into a single unk |
+| `byte_fallback` | an out-of-vocabulary character becomes its UTF-8 bytes as `<0xXX>` tokens (uppercase, zero-padded). All-or-nothing per character: one missing byte token sends the whole character to unk |
+| `ignore_merges` | a piece that is already a vocabulary entry is emitted whole, skipping merges. Required by GPT-4o / Llama-3 tokenizers |
+
+Setting `unk_token` to a string that is absent from the vocabulary is allowed at
+construction and raises `ModelError` only if unresolvable input actually arrives,
+which is when HuggingFace raises too.
+
+### Two behaviours worth knowing
+
+**A pending unk survives a byte-fallback hit.** HuggingFace defers at most one
+unresolved character, and a successful byte expansion does not flush it. With
+`<0xE4>` missing from the vocabulary, `'中é'` gives
+`['<0xC3>', '<0xA9>', '<unk>']` — the unk for `中` lands *after* the bytes of `é`.
+This is likely an upstream quirk, but byte-parity with HuggingFace is the point
+of this package, so it is reproduced rather than corrected.
+
+**`byte_fallback` is encode-only.** The `ByteLevel` decoder maps each character
+through the byte table, so it renders `<0xE4>` as the six ASCII characters of its
+own name rather than as part of `中`. Round-tripping byte-fallback output needs a
+`ByteFallback` decoder, which does not exist yet.
 
 ## PostProcessor (`ubi_tokenizer.processors`)
 
@@ -61,3 +91,24 @@ Reconstructs text from tokens.
 - `BpeTrainer(vocab_size, min_frequency, special_tokens=None, ...)` — in-memory
   greedy trainer. For large/distributed training use UbiTok; this is for small
   corpora and tests.
+
+## Unimplemented options
+
+`BPE` and `BpeTrainer` accept the full HuggingFace option set so that a
+`tokenizer.json` round-trips unchanged. Options that are not implemented raise
+`UnsupportedFeatureError` when set to a value that would change tokenization,
+rather than being accepted and ignored:
+
+| Option | Accepted values | Rejected values |
+| --- | --- | --- |
+| `dropout` | `None`, `0.0` | anything in `(0, 1]` — outside `[0, 1]` raises `ModelError` |
+| `continuing_subword_prefix` | `None`, `""` | any non-empty string |
+| `end_of_word_suffix` | `None`, `""` | any non-empty string |
+
+The accepted values are genuine no-ops in HuggingFace as well, so files that
+tokenize correctly today keep loading — GPT-2-lineage exports ship `""` for both
+affix options.
+
+`cache_capacity` is accepted and retained but not yet used. It is not rejected
+because a cache is a pure optimization: its absence cannot change output.
+HuggingFace does not serialize it into `tokenizer.json`, so `to_dict()` omits it.

pyproject.toml

@@ -21,6 +21,10 @@ dependencies = []
 
 [project.optional-dependencies]
 transcription = []
+# Only needed to regenerate tests/fixtures/hf_bpe_flags.json; the test suite
+# itself runs on the standard library alone. The tokenizers pin is deliberate:
+# the recorded ground truth includes version-sensitive behaviour.
+dev = ["pytest", "tokenizers==0.23.1"]
 
 [project.scripts]
 ubi-tokenizer = "ubi_tokenizer.__main__:main"

scripts/gen_hf_flag_fixtures.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""Regenerate the HuggingFace ground-truth fixtures for BPE option behaviour.
+
+``tests/test_bpe_flags.py`` checks ``ubi_tokenizer``'s BPE against recorded
+HuggingFace output. This script produces that recording, so the expectations are
+measured rather than derived from reading the Rust source.
+
+It is the only part of the test setup that needs a third-party package, and it is
+run by hand -- the test suite itself stays standard-library only, which is the
+point of the whole project. Install the dev extra first:
+
+    pip install -e ".[dev]"
+    python scripts/gen_hf_flag_fixtures.py
+
+The ``tokenizers`` version is pinned in the dev extra and recorded in the output.
+Several behaviours captured here are quirky enough (see ``bf_partial_bytes``)
+that regenerating against a different version could silently move ground truth.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+try:
+    import tokenizers
+    from tokenizers import models
+except ImportError:  # pragma: no cover - developer-facing script
+    raise SystemExit(
+        'huggingface tokenizers is required to regenerate fixtures: pip install -e ".[dev]"'
+    )
+
+OUTPUT = Path(__file__).resolve().parent.parent / "tests" / "fixtures" / "hf_bpe_flags.json"
+
+BYTE_TOKENS = [f"<0x{value:02X}>" for value in range(256)]
+
+
+def vocab_of(*groups: list[str]) -> dict[str, int]:
+    """Assign ids in listed order, skipping duplicates."""
+    vocab: dict[str, int] = {}
+    for group in groups:
+        for token in group:
+            if token not in vocab:
+                vocab[token] = len(vocab)
+    return vocab
+
+
+def bytes_except(*excluded: str) -> list[str]:
+    return [token for token in BYTE_TOKENS if token not in excluded]
+
+
+# Inputs reused across the byte-fallback cases. "\t" (0x09), "\n" (0x0A) and
+# "\x01" exercise bytes below 0x10, where an implementation that produced
+# uppercase-but-not-zero-padded hex would otherwise pass unnoticed -- every
+# CJK/accented character has bytes >= 0x80 and so is always two hex digits.
+LOW_BYTE_INPUTS = ["\t", "\n", "\x01", "\x0f"]
+
+CASES: list[dict[str, Any]] = [
+    {
+        "name": "baseline_no_flags",
+        "why": "merge loop must not regress",
+        "vocab": vocab_of(["a", "b", "c", "ab", "abc"]),
+        "merges": [("a", "b"), ("ab", "c")],
+        "flags": {},
+        "inputs": ["abc", "ab", "a", "abz", "zz", "az", ""],
+    },
+    {
+        "name": "ignore_merges_off_control",
+        "why": "control for the case below; proves the flag is what changes output",
+        "vocab": vocab_of(["x", "y", "xy", "z", "xz"]),
+        "merges": [("x", "z")],
+        "flags": {},
+        "inputs": ["xy", "xz", "xyz", "xyxy", ""],
+    },
+    {
+        "name": "ignore_merges_discriminating",
+        "why": "'xy' is in vocab but no merge reaches it, so the flag alone decides",
+        "vocab": vocab_of(["x", "y", "xy", "z", "xz"]),
+        "merges": [("x", "z")],
+        "flags": {"ignore_merges": True},
+        "inputs": ["xy", "xz", "xyz", "xyxy", ""],
+    },
+    {
+        "name": "drop_when_no_unk",
+        "why": "out-of-vocab characters are dropped, not turned into an invalid id",
+        "vocab": vocab_of(["a", "b"]),
+        "merges": [],
+        "flags": {},
+        "inputs": ["zz", "azb", "az", "za", "中"],
+    },
+    {
+        "name": "unk_only",
+        "why": "one unk per out-of-vocab character",
+        "vocab": vocab_of(["a", "b", "ab", "<unk>"]),
+        "merges": [("a", "b")],
+        "flags": {"unk_token": "<unk>"},
+        "inputs": ["ab", "az", "zz", "az z", "中", "<unk>"],
+    },
+    {
+        "name": "unk_fused",
+        "why": "consecutive pending unks coalesce",
+        "vocab": vocab_of(["a", "b", "ab", "<unk>"]),
+        "merges": [("a", "b")],
+        "flags": {"unk_token": "<unk>", "fuse_unk": True},
+        "inputs": ["ab", "az", "zz", "az z", "中", "azza"],
+    },
+    {
+        "name": "unk_token_present_in_vocab",
+        "why": "an in-vocab unk string must not swallow a following genuine unk",
+        "vocab": vocab_of(["a", "U"]),
+        "merges": [],
+        "flags": {"unk_token": "U", "fuse_unk": True},
+        "inputs": ["Uz", "zU", "aUza", "UU"],
+    },
+    {
+        "name": "byte_fallback_with_unk",
+        "why": "byte expansion, including bytes below 0x10",
+        "vocab": vocab_of(["a", "<unk>"], BYTE_TOKENS),
+        "merges": [],
+        "flags": {"unk_token": "<unk>", "byte_fallback": True},
+        "inputs": ["a", "az", "中", "é", "abc", "€", *LOW_BYTE_INPUTS],
+    },
+    {
+        "name": "byte_fallback_without_unk",
+        "why": "no unk to fall back on; misses are dropped",
+        "vocab": vocab_of(["a"], bytes_except("<0xE4>", "<0xB8>", "<0xAD>")),
+        "merges": [],
+        "flags": {"byte_fallback": True},
+        "inputs": ["a", "az", "中", "é", "a中a", *LOW_BYTE_INPUTS],
+    },
+    {
+        "name": "bare_ascii_absent_but_byte_token_present",
+        "why": "byte fallback also fires for plain ASCII when the bare char is missing",
+        "vocab": vocab_of(["<unk>"], ["<0x61>", "<0x62>", "<0x63>"]),
+        "merges": [],
+        "flags": {"unk_token": "<unk>", "byte_fallback": True},
+        "inputs": ["abc", "a", "abd"],
+    },
+    {
+        "name": "bf_partial_bytes",
+        "why": "all-or-nothing per character, and the pending-unk slot is NOT flushed by a byte-fallback hit",
+        "vocab": vocab_of(["a", "<unk>"], bytes_except("<0xE4>")),
+        "merges": [],
+        "flags": {"unk_token": "<unk>", "byte_fallback": True},
+        "inputs": ["中", "中é", "é中", "中中é", "中é中", "中中", "aé中", "中a"],
+    },
+    {
+        "name": "bf_partial_bytes_fused",
+        "why": "same slot behaviour with coalescing on top",
+        "vocab": vocab_of(["a", "<unk>"], bytes_except("<0xE4>")),
+        "merges": [],
+        "flags": {"unk_token": "<unk>", "byte_fallback": True, "fuse_unk": True},
+        "inputs": ["中", "中é", "é中", "中中é", "中é中", "中中", "aé中", "中a"],
+    },
+    {
+        "name": "merge_over_byte_tokens",
+        "why": "byte-fallback tokens participate in merges, so resolution precedes merging",
+        "vocab": vocab_of(
+            ["a", "<unk>", "<0xE4>", "<0xB8>", "<0xAD>", "<0xE4><0xB8>", "a<0xE4>"]
+        ),
+        "merges": [("<0xE4>", "<0xB8>")],
+        "flags": {"unk_token": "<unk>", "byte_fallback": True},
+        "inputs": ["中", "a中"],
+    },
+    {
+        "name": "merge_over_byte_tokens_rank_order",
+        "why": "merge rank over byte tokens changes the result",
+        "vocab": vocab_of(
+            ["a", "<unk>", "<0xE4>", "<0xB8>", "<0xAD>", "<0xE4><0xB8>", "a<0xE4>"]
+        ),
+        "merges": [("a", "<0xE4>"), ("<0xE4>", "<0xB8>")],
+        "flags": {"unk_token": "<unk>", "byte_fallback": True},
+        "inputs": ["中", "a中"],
+    },
+    {
+        "name": "merge_over_unk_tokens",
+        "why": "generated unk tokens participate in merges too",
+        "vocab": vocab_of(["a", "<unk>", "<unk><unk>"]),
+        "merges": [("<unk>", "<unk>")],
+        "flags": {"unk_token": "<unk>"},
+        "inputs": ["zz", "azza", "z", "zzz"],
+    },
+    {
+        "name": "fuse_unk_without_unk_token",
+        "why": "fuse_unk is inert with no unk token to fuse",
+        "vocab": vocab_of(["a", "b"]),
+        "merges": [],
+        "flags": {"fuse_unk": True},
+        "inputs": ["azza", "zz", "ab"],
+    },
+    {
+        "name": "ignore_merges_falls_through_to_unk",
+        "why": "a piece absent from vocab still takes the normal resolve path",
+        "vocab": vocab_of(["a", "b", "ab", "<unk>"]),
+        "merges": [("a", "b")],
+        "flags": {"ignore_merges": True, "unk_token": "<unk>"},
+        "inputs": ["ab", "abz", "zz", "a"],
+    },
+]
+
+
+def build(case: dict[str, Any]) -> dict[str, Any]:
+    model = models.BPE(
+        vocab=dict(case["vocab"]),
+        merges=[tuple(pair) for pair in case["merges"]],
+        **case["flags"],
+    )
+    expected = {}
+    for text in case["inputs"]:
+        expected[text] = [token.value for token in model.tokenize(text)]
+    return {
+        "name": case["name"],
+        "why": case["why"],
+        "vocab": case["vocab"],
+        "merges": [list(pair) for pair in case["merges"]],
+        "flags": case["flags"],
+        "expected": expected,
+    }
+
+
+def main() -> None:
+    payload = {
+        "_generated_by": "scripts/gen_hf_flag_fixtures.py",
+        "_ground_truth": f"huggingface tokenizers {tokenizers.__version__}",
+        "_warning": (
+            "Do not hand-edit. Regenerate with the pinned tokenizers version; some "
+            "recorded behaviour is version-sensitive."
+        ),
+        "cases": [build(case) for case in CASES],
+    }
+    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
+    OUTPUT.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    total = sum(len(case["expected"]) for case in payload["cases"])
+    print(f"wrote {OUTPUT.relative_to(Path.cwd())}: {len(payload['cases'])} cases, {total} inputs")
+
+
+if __name__ == "__main__":
+    main()

src/ubi_tokenizer/_validation.py

@@ -0,0 +1,54 @@
+"""Shared rejection rules for recognized-but-unimplemented options.
+
+The BPE model and the BPE trainer both accept options that they store, and in
+the model's case serialize, without ever acting on them. Silently accepting
+such an option hands the caller wrong tokenization with no signal, so a value
+that would change behaviour is rejected instead.
+
+The rule has to stay identical in both places, so it lives here rather than
+being duplicated.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from ubi_tokenizer.errors import ModelError, UnsupportedFeatureError
+
+_HINT = (
+    "recognized but not implemented by ubi_tokenizer. Drop the option, or use "
+    "huggingface/tokenizers if you need it."
+)
+
+
+def reject_unimplemented_affixes(
+    continuing_subword_prefix: Optional[str],
+    end_of_word_suffix: Optional[str],
+) -> None:
+    """Reject affix options that would change tokenization.
+
+    ``None`` and ``""`` are genuine no-ops in HuggingFace and stay accepted --
+    GPT-2-lineage ``tokenizer.json`` files ship ``""`` for both, and rejecting
+    those would break loading files that tokenize correctly today.
+    """
+    for name, value in (
+        ("continuing_subword_prefix", continuing_subword_prefix),
+        ("end_of_word_suffix", end_of_word_suffix),
+    ):
+        if value:
+            raise UnsupportedFeatureError(f"{name}={value!r} is {_HINT}")
+
+
+def reject_unimplemented_dropout(dropout: Optional[float]) -> None:
+    """Reject BPE dropout.
+
+    ``None`` and ``0.0`` both mean "no dropout" and stay accepted. Values
+    outside ``[0, 1]`` are invalid rather than unimplemented, which is why they
+    raise :class:`ModelError`; HuggingFace rejects them at construction too.
+    """
+    if dropout is None:
+        return
+    if not 0.0 <= dropout <= 1.0:
+        raise ModelError(f"dropout={dropout!r} must be between 0 and 1, inclusive")
+    if dropout > 0.0:
+        raise UnsupportedFeatureError(f"dropout={dropout!r} is {_HINT}")

src/ubi_tokenizer/errors.py

@@ -13,3 +13,12 @@ class DeserializationError(UbiTokenizerError):
 
 class ModelError(UbiTokenizerError):
     """Raised for invalid model state (missing tokens, bad merges, etc.)."""
+
+
+class UnsupportedFeatureError(UbiTokenizerError):
+    """Raised when a requested option is recognized but not implemented.
+
+    Distinct from :class:`ModelError`: the request is well-formed, it simply
+    cannot be honoured. Raised eagerly so the caller finds out when the
+    tokenizer is built rather than getting silently wrong tokenization.
+    """