tokenizer.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Sequence, List, Optional

import sentencepiece as spm

logger = logging.getLogger(__name__)

# 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


def train_sentencepiece(
    data_files: Sequence[str],
    model_prefix: str = 'tokenizer',
    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 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:
        raise ValueError('data_files is empty')

    missing = [f for f in data_files if not Path(f).exists()]
    if missing:
        raise FileNotFoundError(f'Missing input files: {missing}')

    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,
    )

    logger.info(f"Training SentencePiece: vocab_size={vocab_size} model_type={model_type} "
                f"files={len(data_files)}")
    spm.SentencePieceTrainer.train(**kwargs)

    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,
    )
    return model_path


def _validate_trained_model(
    model_path: str,
    expected_vocab_size: int,
    expected_pad: int,
    expected_unk: int,
    expected_bos: int,
    expected_eos: int,
) -> None:
    """
    Self-critique validation pass — checks the things that actually broke
    in the previous (tiktoken) tokenizer, not just "does it load".
    """
    sp = spm.SentencePieceProcessor(model_file=model_path)

    # 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} "
                        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 = sp.decode(ids)
    if not decoded.strip():
        raise ValueError('Validation round-trip produced empty decode')

    # 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:
    def __init__(self, model_path: str):
        model_path = str(Path(model_path))
        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 < 0:
                raise ValueError(f'SentencePiece model missing <{name}>')
        self._special_ids = {self.pad_id, self.bos_id, self.eos_id}

    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:
            ids = list(self.sp.encode(text, out_type=int))
        if add_bos:
            ids = [self.bos_id] + ids
        if add_eos:
            ids = ids + [self.eos_id]
        return ids

    def encode_large_text(
        self,
        text: str,
        add_bos: bool = True,
        add_eos: bool = False,
        chunk_chars: int = 2_000_000,
    ) -> List[int]:
        """OOM-safe encode for large (multi-MB+) strings, e.g. a whole
        training corpus file read in one go. `encode()` above calls
        self.sp.encode() on the ENTIRE string in one call -- even though
        SentencePiece's C++ core is fast (not a slow Python BPE loop),
        the RETURN VALUE is still built as one giant Python list[int],
        and Python's own string handling means a large UTF-8 file can
        already be 2-4x its byte size once decoded into a `str` object.
        For a real, confirmed example: a 300MB single-file corpus OOM'd
        a ~12-13GB Colab/Kaggle box within ~30 seconds calling encode()
        on the whole file at once -- the failure was in Python-level
        memory (str + list[int] materialization), not inside
        SentencePiece's own encoding step.

        This method processes `text` in fixed-size character chunks,
        converting each chunk's ids to the accumulator and discarding
        the chunk's own Python list before moving to the next chunk, so
        peak memory is bounded by `chunk_chars` instead of len(text).
        Chunk boundaries are placed at whitespace (never mid-word/
        mid-token), so the output is IDENTICAL to what encode(text,
        add_bos, add_eos) would produce on the whole string at once --
        this is not an approximation, it's the same encoding, just
        computed incrementally. Verified by direct comparison across
        many chunk sizes and edge cases (empty string, no-whitespace
        text, exact chunk-boundary alignment) against the whole-string
        encode() path.

        Use this instead of encode() for anything that might be large
        (a whole file's contents) -- keep encode() for genuinely small
        strings (a single sentence/prompt at inference time) where the
        chunking overhead isn't worth paying.
        """
        if text is None:
            raise ValueError('encode_large_text() received None')

        if text == '':
            ids: List[int] = []
            if add_bos:
                ids = [self.bos_id] + ids
            if add_eos:
                ids = ids + [self.eos_id]
            return ids

        all_ids: List[int] = []
        pos = 0
        n = len(text)
        while pos < n:
            end = min(pos + chunk_chars, n)
            is_eof = (end >= n)
            window = text[pos:end]
            if is_eof:
                piece = window
                pos = end
            else:
                # Cut at the LAST whitespace boundary inside this
                # window, so no token is split across two encode()
                # calls (which would silently tokenize the seam
                # differently than encoding the whole string at once).
                cut = max(window.rfind(' '), window.rfind('\n'))
                if cut <= 0:
                    # cut == -1: no whitespace anywhere in this window
                    # (pathological -- one token/URL longer than
                    # chunk_chars). cut == 0: the window's FIRST
                    # character is whitespace, which would make
                    # `piece = window[:0]` empty and, worse, silently
                    # drop that whitespace character (it's before the
                    # cut, so it's never included in this piece OR the
                    # next one, since pos would advance past it without
                    # encoding it). Both cases: take the whole window
                    # as one piece and advance past it -- correctness
                    # at this one seam matters less than a silently
                    # dropped character or a stalled loop, and this
                    # only ever triggers when chunk_chars is smaller
                    # than a typical word (never true at the real
                    # ~2MB default; only reachable with a deliberately
                    # tiny chunk_chars in tests).
                    piece = window
                    pos = end
                else:
                    piece = window[:cut]
                    pos += cut
            if not piece:
                # Should be unreachable now that cut<=0 is folded into
                # the "take whole window" branch above, but keep this
                # as a hard backstop against any future edge case that
                # produces an empty piece with is_eof=False -- without
                # it, pos would never advance and the loop would spin
                # forever.
                if not is_eof:
                    pos += 1
                continue
            piece_ids = self.sp.encode(piece, out_type=int)
            all_ids.extend(piece_ids)
            del piece_ids

        if add_bos:
            all_ids = [self.bos_id] + all_ids
        if add_eos:
            all_ids = all_ids + [self.eos_id]
        return all_ids

    def encode_batch(
        self,
        texts: Sequence[str],
        add_bos: bool = True,
        add_eos: bool = False,
        skip_errors: bool = False,
    ) -> List[List[int]]:
        out: List[List[int]] = []
        for i, t in enumerate(texts):
            try:
                out.append(self.encode(t, add_bos=add_bos, add_eos=add_eos))
            except Exception as e:
                if skip_errors:
                    logger.warning(f"encode_batch: skipping item {i} ({e})")
                    continue
                raise
        return out

    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str:
        # Drop anything outside the valid piece-id range first. This is
        # required, not cosmetic: PyTorch's ignore_index=-100 convention for
        # masked label positions means `ids` is very commonly a raw labels
        # tensor, and sp.decode() raises IndexError on any id < 0 or
        # >= vocab_size instead of skipping it.
        ids = [int(i) for i in ids if 0 <= int(i) < self.vocab_size]
        if skip_special_tokens:
            filtered = [i for i in ids if i not in self._special_ids]
        else:
            filtered = [i for i in ids if i != self.pad_id]
        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]

    def save_config(self, path: str) -> None:
        Path(path).write_text(json.dumps({
            'vocab_size': self.vocab_size,
            'pad_id': self.pad_id,
            'unk_id': self.unk_id,
            'bos_id': self.bos_id,
            'eos_id': self.eos_id,
        }, indent=2), encoding='utf-8')

    @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'))
            mismatches = {
                k: (cfg[k], getattr(tok, k))
                for k in ('vocab_size', 'pad_id', 'unk_id', 'bos_id', 'eos_id')
                if k in cfg and cfg[k] != getattr(tok, k)
            }
            if mismatches:
                raise ValueError(f'Tokenizer/config mismatch: {mismatches}')
        return tok