Update tokenizer.py

tokenizer.py

@@ -5,52 +5,62 @@ import logging
 from pathlib import Path
 from typing import Sequence, List, Optional
 
-from tokenizers import Tokenizer, pre_tokenizers, decoders, trainers
-from tokenizers.models import BPE
-from tokenizers.processors import TemplateProcessing
-import tiktoken
+import sentencepiece as spm
 
 logger = logging.getLogger(__name__)
 
-# GPT-2 style byte-level regex pretokenizer pattern. This is what actually
-# prevents "train." and "train!" from ever becoming distinct merged tokens
-# in the first place: punctuation, whitespace, and word characters are
-# split into separate pretoken chunks BEFORE BPE merges run, so BPE only
-# ever sees clean "train" as a candidate word-piece, with "." / "!" as
-# their own separate single-character pretokens. This is a pretokenization
-# property, not something achieved by pruning merges after training.
-GPT2_SPLIT_PATTERN = (
-    r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
-)
+# 32K is the well-established baseline vocab size for BPE/SentencePiece
+# LLM tokenizers (Llama-1/2, T5, Gopher, Chinchilla all use exactly this).
+# 128K+ only pays off for heavy multilingual/code coverage; for a small,
+# largely-English, narrow-domain model, 32K is the standard, safe default.
+DEFAULT_VOCAB_SIZE = 32000
 
-SPECIAL_TOKENS = ["<unk>", "<pad>", "<bos>", "<eos>"]
-UNK_ID, PAD_ID, BOS_ID, EOS_ID = 0, 1, 2, 3
 
-
-def token_train(
+def train_sentencepiece(
     data_files: Sequence[str],
     model_prefix: str = 'tokenizer',
-    vocab_size: int = 32000,
-    min_frequency: int = 2,
+    vocab_size: int = DEFAULT_VOCAB_SIZE,
+    model_type: str = 'bpe',
+    character_coverage: float = 0.9995,
+    byte_fallback: bool = True,
+    pad_id: int = 1,
+    unk_id: int = 0,
+    bos_id: int = 2,
+    eos_id: int = 3,
+    add_dummy_prefix: bool = True,
+    num_threads: int = 8,
+    input_sentence_size: int = 5_000_000,
+    shuffle_input_sentence: bool = True,
+    max_sentence_length: int = 16384,
+    split_digits: bool = True,
+    allow_whitespace_only_pieces: bool = True,
+    train_extremely_large_corpus: bool = False,
 ) -> str:
     """
-    Train a byte-level BPE tokenizer (tiktoken/GPT-2 style) using HF
-    `tokenizers` as the trainer backend, then export a tiktoken-loadable
-    merges file for fast inference.
-
-    Why HF `tokenizers` as the trainer and not tiktoken directly: tiktoken
-    itself is inference-only — it has no production-grade vocab trainer
-    (its reference trainer in `_educational.py` is explicitly slow/
-    reference-quality, not meant for real corpora). `tokenizers` is the
-    standard fast Rust-backed BPE trainer; we train there, then hand off
-    the resulting merges to tiktoken's `Encoding` for fast BPE inference
-    at train/eval time.
-
-    The GPT2_SPLIT_PATTERN pretokenizer is what keeps "train.", "train!",
-    "train," etc. from ever being learned as single fused tokens — the
-    regex splits words from adjacent punctuation before BPE ever runs, so
-    BPE always sees "train" as a standalone candidate, with punctuation as
-    its own separate byte-level pretoken.
+    Train a SentencePiece BPE tokenizer with byte-fallback — the same
+    scheme used by Llama-2, Mistral, and EuroLLM (BPE + byte_fallback via
+    SentencePiece specifically, not a hand-rolled BPE implementation).
+
+    Why SentencePiece and not a hand-written tiktoken export: SentencePiece's
+    C++ core does encode/decode and merge-rank bookkeeping internally and
+    natively — there is no manual ID-renumbering or rank-export step for
+    calling code to get wrong. (A prior tiktoken-based rewrite of this
+    tokenizer had exactly that class of bug: hand-exported merge ranks were
+    non-contiguous because special tokens occupied ids 0-3 in the source
+    vocab, silently corrupting merge-priority order and decode() mappings —
+    manifesting as repetitive garbage output like "to to to" despite a
+    healthy training loss. Delegating to SentencePiece's own encode/decode
+    removes that entire class of bug by construction.)
+
+    Notes on defaults:
+      - character_coverage < 1.0 with byte_fallback=True: rare glyphs fall
+        back to byte pieces instead of bloating the vocab with singletons.
+      - input_sentence_size + shuffle_input_sentence: without shuffling,
+        SentencePiece samples from the START of the concatenated corpus,
+        which silently biases vocab toward whichever domain file comes
+        first if you hand it multiple files back to back.
+      - split_digits: keeps numbers as individual digit tokens, which
+        generally helps arithmetic/math task tokenization consistency.
     """
     data_files = [str(Path(p)) for p in data_files]
     if not data_files:
@@ -60,278 +70,153 @@ def token_train(
     if missing:
         raise FileNotFoundError(f'Missing input files: {missing}')
 
-    tokenizer = Tokenizer(BPE(unk_token="<unk>", byte_fallback=True))
-    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(
-        add_prefix_space=False, use_regex=True,
-    )
-    tokenizer.decoder = decoders.ByteLevel()
-
-    trainer = trainers.BpeTrainer(
-        vocab_size=vocab_size,
-        min_frequency=min_frequency,
-        special_tokens=SPECIAL_TOKENS,
-        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
-        show_progress=True,
-    )
-
-    logger.info(f"Training BPE tokenizer: vocab_size={vocab_size} files={len(data_files)}")
-    tokenizer.train(files=data_files, trainer=trainer)
-
-    tokenizer.post_processor = TemplateProcessing(
-        single="<bos> $A <eos>",
-        special_tokens=[
-            ("<bos>", tokenizer.token_to_id("<bos>")),
-            ("<eos>", tokenizer.token_to_id("<eos>")),
-        ],
+    kwargs = dict(
+        input=','.join(data_files),
+        model_prefix=model_prefix,
+        vocab_size=int(vocab_size),
+        model_type=model_type,
+        character_coverage=character_coverage,
+        pad_id=pad_id,
+        unk_id=unk_id,
+        bos_id=bos_id,
+        eos_id=eos_id,
+        byte_fallback=byte_fallback,
+        hard_vocab_limit=False,
+        normalization_rule_name='nmt_nfkc',
+        add_dummy_prefix=add_dummy_prefix,
+        num_threads=num_threads,
+        input_sentence_size=input_sentence_size,
+        shuffle_input_sentence=shuffle_input_sentence,
+        max_sentence_length=max_sentence_length,
+        split_digits=split_digits,
+        allow_whitespace_only_pieces=allow_whitespace_only_pieces,
+        train_extremely_large_corpus=train_extremely_large_corpus,
     )
 
-    hf_json_path = f'{model_prefix}.tokenizer.json'
-    tokenizer.save(hf_json_path)
-
-    tiktoken_path = _export_tiktoken_format(tokenizer, model_prefix)
-    _validate_trained_model(hf_json_path, tiktoken_path, vocab_size)
-
-    return tiktoken_path
-
-
-def _bytes_to_unicode() -> dict:
-    """
-    The canonical GPT-2 byte<->unicode bijection (same table used inside
-    HF's ByteLevel pre_tokenizer/decoder, and in tiktoken's own reference
-    implementations). Maps every raw byte value 0-255 to a printable
-    unicode code point, so byte-level BPE can be trained/represented as
-    ordinary text.
+    logger.info(f"Training SentencePiece: vocab_size={vocab_size} model_type={model_type} "
+                f"files={len(data_files)}")
+    spm.SentencePieceTrainer.train(**kwargs)
 
-    Returns: dict {byte_value (0-255): unicode_char}
-    """
-    bs = (
-        list(range(ord("!"), ord("~") + 1))
-        + list(range(ord("¡"), ord("¬") + 1))
-        + list(range(ord("®"), ord("ÿ") + 1))
+    model_path = f'{model_prefix}.model'
+    _validate_trained_model(
+        model_path, vocab_size,
+        expected_pad=pad_id, expected_unk=unk_id, expected_bos=bos_id, expected_eos=eos_id,
     )
-    cs = bs[:]
-    n = 0
-    for b in range(256):
-        if b not in bs:
-            bs.append(b)
-            cs.append(256 + n)
-            n += 1
-    cs = [chr(c) for c in cs]
-    return dict(zip(bs, cs))
-
-
-# Built once at import time: char -> original raw byte value. This is the
-# ONLY correct way to recover a token's raw bytes from its GPT-2
-# byte-level-alphabet string form.
-_UNICODE_TO_BYTE = {v: k for k, v in _bytes_to_unicode().items()}
-
-
-def _token_str_to_bytes(token_str: str) -> bytes:
-    """
-    Converts a GPT-2 byte-level-alphabet token string back to its raw
-    bytes by inverting the bijection CHARACTER BY CHARACTER.
-
-    CRITICAL: this must NOT go through `tokenizer.decoder.decode(...)`
-    followed by `.encode('utf-8')`. That round-trip treats the decoded
-    result as TEXT, and a lone raw byte in the 128-255 range is not valid
-    UTF-8 on its own — decoding it in isolation silently produces the
-    Unicode replacement character (U+FFFD) instead of the original byte.
-    Every byte value 128-255 collapses into that same wrong 3-byte
-    sequence this way, which is exactly what caused the "no entry found
-    for key" Rust panic: 128 real single-byte entries were silently
-    replaced by 127 colliding, wrong entries, deleting the entire
-    non-ASCII byte range from the exported vocabulary. Any text
-    containing so much as one accented character, curly quote, or emoji
-    (i.e. almost any real-world corpus) then has no byte-fallback entry
-    to encode with, and tiktoken's Rust core panics.
-    """
-    return bytes(_UNICODE_TO_BYTE[ch] for ch in token_str)
+    return model_path
 
 
-def _export_tiktoken_format(tokenizer: "Tokenizer", model_prefix: str) -> str:
+def _validate_trained_model(
+    model_path: str,
+    expected_vocab_size: int,
+    expected_pad: int,
+    expected_unk: int,
+    expected_bos: int,
+    expected_eos: int,
+) -> None:
     """
-    Export the trained HF tokenizer's merges into a tiktoken-loadable
-    `.tiktoken` file (rank<TAB>base64-token-bytes per line) plus a sidecar
-    JSON with special-token ids, so TokenizerWrapper can load via
-    tiktoken.Encoding for fast inference.
-
-    CRITICAL: tiktoken requires mergeable_ranks to be DENSE, starting at 0
-    (rank IS the merge-priority order — see any real tiktoken.Encoding
-    setup, e.g. Llama's tokenizer: special tokens are always assigned
-    `len(mergeable_ranks) + i`, i.e. right after the merge ranks end).
-    HF's BpeTrainer, because special_tokens is passed first, assigns them
-    ids 0-3 and pushes actual merges to start at id 4 — a gapped,
-    non-contiguous rank space if reused directly. Writing those original
-    ids into the .tiktoken file silently corrupts BPE merge-priority
-    ordering (tokenization no longer matches what was trained), and
-    decode() maps ids to the wrong byte sequences — this was the actual
-    cause of a real repetition/garbage-decode bug in production ("to to
-    to", ",,,"), not a decoding-strategy issue. Fix: renumber merge ranks
-    densely from 0 in sorted-original-id order (preserves relative merge
-    priority, which is all that matters), then assign special token ids
-    as len(mergeable_ranks) + i, matching tiktoken's actual convention.
+    Self-critique validation pass — checks the things that actually broke
+    in the previous (tiktoken) tokenizer, not just "does it load".
     """
-    import base64
-
-    vocab = tokenizer.get_vocab()  # token string -> original HF id
-
-    # Sort NON-special tokens by their original HF id (preserves the
-    # relative merge-priority order learned during training), then
-    # renumber them densely 0..N-1 — this is the rank space tiktoken
-    # actually requires.
-    non_special = sorted(
-        ((s, i) for s, i in vocab.items() if s not in SPECIAL_TOKENS),
-        key=lambda kv: kv[1],
-    )
-
-    tiktoken_path = f'{model_prefix}.tiktoken'
-    with open(tiktoken_path, 'w', encoding='utf-8') as f:
-        for new_rank, (token_str, _old_id) in enumerate(non_special):
-            # Direct char-by-char inversion of the GPT-2 byte<->unicode
-            # bijection — NOT tokenizer.decoder.decode() + UTF-8 encode,
-            # which is lossy for any byte value 128-255 (see
-            # _token_str_to_bytes docstring for exactly why).
-            token_bytes = _token_str_to_bytes(token_str)
-            f.write(f"{base64.b64encode(token_bytes).decode('ascii')} {new_rank}\n")
-
-    # Special tokens go right after the dense merge-rank space ends —
-    # matching tiktoken's real convention (len(mergeable_ranks) + i), NOT
-    # HF's original 0-3 ids.
-    num_merges = len(non_special)
-    special_ids = {
-        name.strip('<>'): num_merges + i for i, name in enumerate(SPECIAL_TOKENS)
-    }
-    Path(f'{model_prefix}.special_tokens.json').write_text(
-        json.dumps(special_ids, indent=2), encoding='utf-8'
-    )
-
-    # ── Hard completeness check ────────────────────────────────────────
-    # tiktoken's BPE core requires every possible single byte (0-255) to
-    # have a mergeable_ranks entry — it's the mandatory fallback base case
-    # for encoding any byte the trained merges don't otherwise cover. If
-    # even one byte value is missing (e.g. from a bug like the one this
-    # function just fixed), encoding ANY text containing that byte panics
-    # deep inside tiktoken's Rust core with a cryptic "no entry found for
-    # key" — with no indication of which byte or why. Catch that here,
-    # immediately after export, with a clear Python-level error instead.
-    exported_single_bytes = set()
-    with open(tiktoken_path, 'r', encoding='utf-8') as f:
-        for line in f:
-            line = line.strip()
-            if not line:
-                continue
-            b64_token, _ = line.split()
-            raw = base64.b64decode(b64_token)
-            if len(raw) == 1:
-                exported_single_bytes.add(raw[0])
-    missing = sorted(set(range(256)) - exported_single_bytes)
-    if missing:
-        raise RuntimeError(
-            f"Tokenizer export is missing {len(missing)}/256 single-byte "
-            f"entries (byte values: {missing[:20]}{'...' if len(missing) > 20 else ''}). "
-            f"Any training text containing these byte values will crash "
-            f"tiktoken's Rust core with 'no entry found for key'. This "
-            f"means the byte<->unicode inversion during export is broken — "
-            f"do not proceed with this .tiktoken file."
-        )
-
-    return tiktoken_path
-
+    sp = spm.SentencePieceProcessor(model_file=model_path)
 
-def _validate_trained_model(hf_json_path: str, tiktoken_path: str, expected_vocab_size: int) -> None:
-    """Round-trip check on the HF tokenizer (source of truth) before we
-    trust the exported tiktoken artifact."""
-    tok = Tokenizer.from_file(hf_json_path)
-    actual_vocab = tok.get_vocab_size()
+    # 1. Vocab size sanity
+    actual_vocab = sp.vocab_size()
     if actual_vocab != expected_vocab_size:
-        logger.warning(f"Trained vocab_size={actual_vocab} differs from requested={expected_vocab_size}")
-
-    for name in SPECIAL_TOKENS:
-        if tok.token_to_id(name) is None:
-            raise ValueError(f'Trained model missing special token {name}')
-
-    probe = "The quick brown fox jumps over 42 lazy dogs. Did it train? train! train."
-    encoding = tok.encode(probe)
-    if not encoding.ids:
+        logger.warning(f"Trained vocab_size={actual_vocab} differs from requested={expected_vocab_size} "
+                        f"(hard_vocab_limit=False allows this if the corpus is small)")
+
+    # 2. Special token IDs must be EXACTLY what was requested — not just
+    # ">= 0". A previous bug class involved special-token ids silently
+    # drifting from what calling code assumed. Check explicitly, not
+    # loosely.
+    checks = [
+        ('pad', sp.pad_id(), expected_pad),
+        ('unk', sp.unk_id(), expected_unk),
+        ('bos', sp.bos_id(), expected_bos),
+        ('eos', sp.eos_id(), expected_eos),
+    ]
+    for name, actual, expected in checks:
+        if actual < 0:
+            raise ValueError(f'Trained model missing <{name}> special token')
+        if actual != expected:
+            raise ValueError(
+                f'<{name}> id drift: requested {expected}, SentencePiece '
+                f'assigned {actual}. This mismatch is exactly the class of '
+                f'bug that broke a previous tokenizer version — refusing '
+                f'to silently proceed.'
+            )
+
+    # 3. Basic round-trip: encode -> decode must reproduce recognizable text
+    probe = "The quick brown fox jumps over 42 lazy dogs. def foo(): return None"
+    ids = sp.encode(probe, out_type=int)
+    if not ids:
         raise ValueError('Validation encode produced empty output')
-    decoded = tok.decode(encoding.ids)
+    decoded = sp.decode(ids)
     if not decoded.strip():
         raise ValueError('Validation round-trip produced empty decode')
 
-    logger.info(f"Validation OK: vocab={actual_vocab} probe_tokens={len(encoding.ids)}")
+    # 4. SPECIFIC regression check for the actual reported failure mode:
+    # repetitive-token degenerate decode ("to to to", ",,,"). This won't
+    # catch a MODEL that's actually stuck in a repetition loop (that's a
+    # decoding-strategy issue, separate from the tokenizer), but it DOES
+    # catch a tokenizer that maps distinct ids to the same or corrupted
+    # text, which was the real bug here: encode the same repeated-word
+    # probe multiple times and confirm token ids are stable and decode
+    # is exact, not degenerating into duplicated/garbled pieces.
+    repeat_probe = "to to to , , , the the the"
+    repeat_ids = sp.encode(repeat_probe, out_type=int)
+    repeat_decoded = sp.decode(repeat_ids)
+    # Re-encoding the decoded output should reproduce the same ids
+    # (idempotency) — this is the real symptom check: a corrupted rank/id
+    # mapping breaks exactly this property even when a single encode/decode
+    # pass looks fine.
+    reencoded_ids = sp.encode(repeat_decoded, out_type=int)
+    if reencoded_ids != repeat_ids:
+        raise ValueError(
+            f'Round-trip idempotency FAILED on repeated-token probe: '
+            f'encode->decode->encode did not reproduce the same ids. '
+            f'original={repeat_ids} reencoded={reencoded_ids}. This is '
+            f'the specific failure signature of an id/rank mapping bug.'
+        )
+
+    # 5. Byte-fallback sanity: an unusual/rare unicode character must not
+    # crash and must not silently become <unk> if byte_fallback is on —
+    # it should decompose into byte pieces instead.
+    exotic_probe = "emoji test \U0001F600 and rare char \u0800"
+    exotic_ids = sp.encode(exotic_probe, out_type=int)
+    if not exotic_ids:
+        raise ValueError('Byte-fallback validation: exotic-character probe produced empty encode')
+    exotic_decoded = sp.decode(exotic_ids)
+    if not exotic_decoded.strip():
+        raise ValueError('Byte-fallback validation: exotic-character round-trip produced empty decode')
 
+    logger.info(f"✓ Validation OK: vocab={actual_vocab} probe_tokens={len(ids)} "
+                f"round-trip idempotency verified, byte-fallback verified")
 
-class TokenizerWrapper:
-    """
-    Fast-path wrapper using tiktoken for encode/decode at train/eval time,
-    backed by a vocab trained via HF `tokenizers` (see train_bpe_tokenizer).
-    """
 
+class TokenizerWrapper:
     def __init__(self, model_path: str):
         model_path = str(Path(model_path))
-        tiktoken_path = model_path if model_path.endswith('.tiktoken') else f'{model_path}.tiktoken'
-        special_path = tiktoken_path.replace('.tiktoken', '.special_tokens.json')
-
-        if not Path(tiktoken_path).exists():
-            raise FileNotFoundError(tiktoken_path)
-        if not Path(special_path).exists():
-            raise FileNotFoundError(special_path)
-
-        special_ids = json.loads(Path(special_path).read_text(encoding='utf-8'))
-        self.unk_id = special_ids['unk']
-        self.pad_id = special_ids['pad']
-        self.bos_id = special_ids['bos']
-        self.eos_id = special_ids['eos']
+        if not Path(model_path).exists():
+            raise FileNotFoundError(model_path)
+        self.sp = spm.SentencePieceProcessor(model_file=model_path)
+        self.vocab_size = int(self.sp.vocab_size())
+        self.pad_id = self.sp.pad_id()
+        self.unk_id = self.sp.unk_id()
+        self.bos_id = self.sp.bos_id()
+        self.eos_id = self.sp.eos_id()
         for name, val in [('pad', self.pad_id), ('unk', self.unk_id), ('bos', self.bos_id), ('eos', self.eos_id)]:
-            if val is None or val < 0:
-                raise ValueError(f'Tokenizer missing <{name}>')
-
-        mergeable_ranks = self._load_tiktoken_ranks(tiktoken_path)
-        self.vocab_size = len(mergeable_ranks) + len(special_ids)
-
-        self.enc = tiktoken.Encoding(
-            name=Path(tiktoken_path).stem,
-            pat_str=GPT2_SPLIT_PATTERN,
-            mergeable_ranks=mergeable_ranks,
-            special_tokens={
-                "<unk>": self.unk_id, "<pad>": self.pad_id,
-                "<bos>": self.bos_id, "<eos>": self.eos_id,
-            },
-        )
+            if val < 0:
+                raise ValueError(f'SentencePiece model missing <{name}>')
         self._special_ids = {self.pad_id, self.bos_id, self.eos_id}
 
-    @staticmethod
-    def _load_tiktoken_ranks(tiktoken_path: str) -> dict:
-        ranks = {}
-        with open(tiktoken_path, 'r', encoding='utf-8') as f:
-            for line in f:
-                line = line.strip()
-                if not line:
-                    continue
-                b64_token, rank = line.split()
-                import base64
-                ranks[base64.b64decode(b64_token)] = int(rank)
-        return ranks
-
     def encode(self, text: str, add_bos: bool = True, add_eos: bool = False) -> List[int]:
         if text is None:
             raise ValueError('encode() received None')
         if text == '':
             ids: List[int] = []
         else:
-            # disallowed_special=() (not allowed_special=set()): training
-            # corpora can contain LITERAL text that happens to match a
-            # special token string — e.g. WikiText-103 ships with literal
-            # "<unk>" markers baked into its raw text from its original
-            # preprocessing. tiktoken's default behavior treats any string
-            # matching a registered special token as a forbidden injection
-            # attempt and raises. We want the opposite: literal "<unk>" in
-            # training data should be encoded as ordinary text (broken into
-            # bytes/subwords), never treated as an actual special token
-            # unless WE insert it programmatically (e.g. via add_bos/add_eos
-            # below, which append the token ID directly, bypassing string
-            # matching entirely — so this doesn't weaken bos/eos handling).
-            ids = self.enc.encode(text, disallowed_special=())
+            ids = list(self.sp.encode(text, out_type=int))
         if add_bos:
             ids = [self.bos_id] + ids
         if add_eos:
@@ -361,7 +246,7 @@ class TokenizerWrapper:
             filtered = [int(i) for i in ids if int(i) not in self._special_ids]
         else:
             filtered = [int(i) for i in ids if int(i) != self.pad_id]
-        return self.enc.decode(filtered)
+        return self.sp.decode(filtered)
 
     def decode_batch(self, batch_ids: Sequence[Sequence[int]], skip_special_tokens: bool = True) -> List[str]:
         return [self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in batch_ids]
@@ -377,6 +262,7 @@ class TokenizerWrapper:
 
     @classmethod
     def from_config(cls, model_path: str, config_path: Optional[str] = None) -> 'TokenizerWrapper':
+        """Load and, if a config is given, verify special-id consistency against it."""
         tok = cls(model_path)
         if config_path and Path(config_path).exists():
             cfg = json.loads(Path(config_path).read_text(encoding='utf-8'))