Delete tokenizer.py

tokenizer.py

@@ -1,282 +0,0 @@
-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_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