Update tokenizer.py

tokenizer.py

@@ -1,283 +1,175 @@
+"""
+tokenizer.py — RustBPETokenizer adapter.
+
+Drop-in replacement for the previous SentencePiece-based TokenizerWrapper.
+Exposes the SAME public API so train.py, _tok_worker.py, and finetune.py
+work unchanged. The only pipeline change required is that the on-disk
+tokenizer artifact is now a directory containing `tokenizer.pkl` (a
+pickled tiktoken.Encoding), rather than a single `tokenizer.model` file.
+
+Set TOKENIZER_MODEL_PATH in train.py to the directory (or to the pickle
+file directly — both are accepted).
+
+API compatibility
+-----------------
+Public attributes:
+    model_path      str   absolute path to the .pkl (hashed into fingerprint)
+    vocab_size      int   total vocab size, specials included
+    bos_id          int   id of <|bos|>
+    eos_id          int   same as bos_id (see Notes)
+    pad_id          int   same as bos_id (see Notes)
+    unk_id          int   same as bos_id (see Notes)
+
+Public methods:
+    encode(text, add_bos=True, add_eos=False)                -> List[int]
+    encode_large_text(text, add_bos, add_eos, chunk_chars)   -> List[int]
+    iter_encode_chunks(text, add_bos, add_eos, chunk_chars)  -> Iterator[np.ndarray]
+    encode_batch(texts, add_bos, add_eos, skip_errors)       -> List[List[int]]
+    decode(ids, skip_special_tokens=True)                    -> str
+    decode_batch(batch_ids, skip_special_tokens=True)        -> List[str]
+    save_config(path)
+    from_config(model_path, config_path=None)                -> TokenizerWrapper
+
+Notes
+-----
+pad_id / eos_id / unk_id all map to bos_id because:
+  - tiktoken is byte-level, so <unk> is never emitted
+    - rustbpe's SPECIAL_TOKENS list has no dedicated <|eos>; this adapter
+        treats each input to encode() as one document. For iter_encode_chunks(),
+        one BOS is placed at the start of the file and one EOS-equivalent BOS is
+        placed at the end, matching the flat SP-era stream contract rather than
+        nanochat's per-document stream.
+  - finetune.py needs SOME valid id for padding; <|bos|> is the standard
+    choice and is filtered by decode(skip_special_tokens=True).
+
+encode() uses tiktoken's encode_ordinary() fast path, which does NOT
+interpret "<|user_start|>" etc. as special tokens. SFT rendering must
+call encode_special_id() for control tokens and encode_ordinary() only
+for content. This is the same contract as nanochat's trainer.
+"""
+
 from __future__ import annotations
 
 import json
 import logging
+import os
+import pickle
 from pathlib import Path
-from typing import Sequence, List, Optional
+from typing import Iterator, List, Optional, Sequence, Union
 
 import numpy as np
-import sentencepiece as spm
-
-DEFAULT_CHUNK_CHARS = 2_000_000
 
 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.'
+DEFAULT_CHUNK_CHARS = 500_000
+_TIKTOKEN_THREADS = int(os.environ.get(
+    "TIKTOKEN_NUM_THREADS",
+    str(max(1, min(8, os.cpu_count() or 1))),
+))
+
+
+# ---------------------------------------------------------------------------
+# Helper: locate the pickle given a directory or a file path
+# ---------------------------------------------------------------------------
+
+def _resolve_pickle_path(model_path: Union[str, Path]) -> Path:
+    p = Path(model_path)
+    if p.is_dir():
+        candidate = p / "tokenizer.pkl"
+        if not candidate.exists():
+            raise FileNotFoundError(
+                f"{p} is a directory but contains no tokenizer.pkl "
+                f"(expected {candidate})"
             )
+        return candidate
+    if not p.exists():
+        raise FileNotFoundError(str(p))
+    return p
 
-    # 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")
 
+# ---------------------------------------------------------------------------
+# TokenizerWrapper
+# ---------------------------------------------------------------------------
 
 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.model_path = 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]:
+    """Adapter exposing a pickled tiktoken.Encoding behind the SP-era API."""
+
+    def __init__(self, model_path: Union[str, Path]):
+        pickle_path = _resolve_pickle_path(model_path)
+        # Note: pickled tiktoken.Encoding. Only load pickles you created.
+        try:
+            with open(pickle_path, "rb") as f:
+                self.enc = pickle.load(f)
+        except Exception as exc:
+            raise ValueError(
+                f"Failed to unpickle {pickle_path}: {exc!r}. "
+                "This adapter expects a pickled tiktoken.Encoding "
+                "produced by RustBPETokenizer.save(). If you have an old "
+                "SentencePiece .model file, train a new tokenizer with rustbpe."
+            ) from exc
+
+        # Validate it looks like a tiktoken Encoding before trusting it.
+        for attr in ("n_vocab", "encode_ordinary", "encode_ordinary_batch",
+                     "decode", "encode_single_token", "special_tokens_set"):
+            if not hasattr(self.enc, attr):
+                raise TypeError(
+                    f"{pickle_path}: loaded object is not a tiktoken.Encoding "
+                    f"(missing attribute {attr!r}). Got {type(self.enc).__name__}."
+                )
+
+        self.model_path = str(pickle_path)
+        self.vocab_size = int(self.enc.n_vocab)
+
+        # BOS is required. If the vocab lacks it, that's a hard error.
+        self.bos_id = int(self.enc.encode_single_token("<|bos|>"))
+        # EOS/PAD/UNK reuse BOS — see module docstring.
+        self.eos_id = self.bos_id
+        self.pad_id = self.bos_id
+        self.unk_id = self.bos_id
+
+        # Everything tiktoken labels as special — used by decode() to drop
+        # control tokens when skip_special_tokens=True.
+        self._special_ids = set()
+        for name in self.enc.special_tokens_set:
+            try:
+                self._special_ids.add(int(self.enc.encode_single_token(name)))
+            except Exception:
+                # If a special name isn't encodable (shouldn't happen for a
+                # well-formed Encoding), skip it rather than crash.
+                pass
+        # Always include BOS even if special_tokens_set was empty for some
+        # reason.
+        self._special_ids.add(self.bos_id)
+
+    # -- single-sequence encode ---------------------------------------------
+
+    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 == '':
+            raise ValueError("encode() received None")
+        if text == "":
             ids: List[int] = []
         else:
-            ids = list(self.sp.encode(text, out_type=int))
+            ids = self.enc.encode_ordinary(text)
         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 = DEFAULT_CHUNK_CHARS,
-    ) -> List[int]:
-        """Backward-compatible wrapper that now delegates to the chunk iterator
-        so encode_large_text and iter_encode_chunks share exactly one loop."""
-        if text is None:
-            raise ValueError('encode_large_text() received None')
-        parts = [
-            arr for arr in self.iter_encode_chunks(
-                text,
-                add_bos=add_bos,
-                add_eos=add_eos,
-                chunk_chars=chunk_chars,
-            )
-        ]
-        if not parts:
-            return []
-        return np.concatenate(parts, axis=0).astype(np.int32).tolist()
-
-    def iter_encode_chunks(
-        self,
-        text: str,
-        add_bos: bool = True,
-        add_eos: bool = True,
-        chunk_chars: int = DEFAULT_CHUNK_CHARS,
-    ):
-        """Yield per-chunk int32 arrays while preserving the first/last chunk
-        BOS/EOS semantics of the existing whole-file encode path.
-
-        This iterator mirrors `encode_large_text()` chunk placement logic,
-        but streams one chunk's token ids out as a NumPy array instead of
-        materializing a Python list for the whole file.
-        """
-        if text is None:
-            raise ValueError('iter_encode_chunks() received None')
-        if text == '':
-            ids: List[int] = []
-            if add_bos:
-                ids = [self.bos_id] + ids
-            if add_eos:
-                ids = ids + [self.eos_id]
-            yield np.asarray(ids, dtype=np.int32)
-            return
+    # -- chunked streaming --------------------------------------------------
 
-        pos = 0
-        n = len(text)
-        first_chunk = True
+    @staticmethod
+    def _iter_chunks(text: str, chunk_chars: int) -> Iterator[str]:
+        """Yield whitespace-aligned chunks. Same boundaries as the SP-era
+        wrapper, so the resulting token stream is comparable."""
+        if chunk_chars <= 0:
+            raise ValueError(f"chunk_chars must be > 0, got {chunk_chars}")
+        pos, n = 0, len(text)
         while pos < n:
             end = min(pos + chunk_chars, n)
             is_eof = (end >= n)
@@ -286,26 +178,106 @@ class TokenizerWrapper:
                 piece = window
                 pos = end
             else:
-                cut = max(window.rfind(' '), window.rfind('\n'))
+                cut = max(window.rfind(" "), window.rfind("\n"))
                 if cut <= 0:
                     piece = window
                     pos = end
                 else:
                     piece = window[:cut]
                     pos += cut
-            if not piece:
-                if not is_eof:
-                    pos += 1
+            if piece:
+                yield piece
+            elif not is_eof:
+                pos += 1
+
+    def iter_encode_chunks(
+        self,
+        text: str,
+        add_bos: bool = True,
+        add_eos: bool = True,
+        chunk_chars: int = DEFAULT_CHUNK_CHARS,
+        batch_size: int | None = None,
+    ) -> Iterator[np.ndarray]:
+        """Yield per-chunk int32 arrays with BOS on first, EOS on last."""
+        if text is None:
+            raise ValueError("iter_encode_chunks() received None")
+
+        if text == "":
+            ids: List[int] = []
+            if add_bos:
+                ids.append(self.bos_id)
+            if add_eos:
+                ids.append(self.eos_id)
+            yield np.asarray(ids, dtype=np.int32)
+            return
+
+        if batch_size is None:
+            memory_budget_bytes = 32_000_000
+            chunks_per_batch = memory_budget_bytes // max(1, chunk_chars)
+            batch_size = max(1, min(16, chunks_per_batch))
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be > 0, got {batch_size}")
+        first_emitted = False
+        pending: List[str] = []
+        current: List[str] | None = None
+
+        def encode_batch(
+            chunks: List[str], is_last: bool, add_bos_here: bool,
+        ) -> Iterator[np.ndarray]:
+            encodings = self.enc.encode_ordinary_batch(
+                chunks, num_threads=_TIKTOKEN_THREADS,
+            )
+            for i, ids in enumerate(encodings):
+                if add_bos_here and i == 0:
+                    ids = [self.bos_id] + ids
+                if add_eos and is_last and i == len(encodings) - 1:
+                    ids = ids + [self.eos_id]
+                yield np.asarray(ids, dtype=np.int32)
+
+        for piece in self._iter_chunks(text, chunk_chars):
+            pending.append(piece)
+            if len(pending) < batch_size:
                 continue
+            if current is not None:
+                yield from encode_batch(current, False, add_bos and not first_emitted)
+                first_emitted = True
+            current, pending = pending, []
+
+        if current is None:
+            current = pending
+        elif pending:
+            yield from encode_batch(current, False, add_bos and not first_emitted)
+            first_emitted = True
+            current = pending
+
+        if current:
+            yield from encode_batch(current, True, add_bos and not first_emitted)
+        else:
+            # Input was all whitespace.
+            ids = []
+            if add_bos:
+                ids.append(self.bos_id)
+            if add_eos:
+                ids.append(self.eos_id)
+            yield np.asarray(ids, dtype=np.int32)
 
-            piece_ids = self.sp.encode(piece, out_type=int)
-            if add_bos and first_chunk:
-                piece_ids = [self.bos_id] + piece_ids
-            if add_eos and is_eof:
-                piece_ids = piece_ids + [self.eos_id]
+    def encode_large_text(
+        self,
+        text: str,
+        add_bos: bool = True,
+        add_eos: bool = False,
+        chunk_chars: int = DEFAULT_CHUNK_CHARS,
+    ) -> List[int]:
+        if text is None:
+            raise ValueError("encode_large_text() received None")
+        parts = list(self.iter_encode_chunks(
+            text, add_bos=add_bos, add_eos=add_eos, chunk_chars=chunk_chars,
+        ))
+        if not parts:
+            return []
+        return np.concatenate(parts, axis=0).astype(np.int32).tolist()
 
-            yield np.asarray(piece_ids, dtype=np.int32)
-            first_chunk = False
+    # -- batch ------------------------------------------------------------
 
     def encode_batch(
         self,
@@ -314,53 +286,94 @@ class TokenizerWrapper:
         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:
+        texts = list(texts)
+        if skip_errors:
+            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:
                     logger.warning(f"encode_batch: skipping item {i} ({e})")
-                    continue
-                raise
+            return out
+
+        for t in texts:
+            if t is None:
+                raise ValueError("encode_batch() received None")
+        if not texts:
+            return []
+
+        encodings = self.enc.encode_ordinary_batch(
+            texts, num_threads=_TIKTOKEN_THREADS,
+        )
+        out = []
+        for ids in encodings:
+            if add_bos:
+                ids = [self.bos_id] + ids
+            if add_eos:
+                ids = ids + [self.eos_id]
+            out.append(ids)
         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]
+    # -- decode -----------------------------------------------------------
+
+    def decode(
+        self,
+        ids: Sequence[int],
+        skip_special_tokens: bool = True,
+    ) -> str:
+        # Filter out-of-range ids first: PyTorch's ignore_index=-1 convention
+        # for masked positions means callers frequently hand us raw label
+        # tensors. tiktoken.decode() raises on any id < 0 or >= vocab_size.
+        clean: List[int] = []
+        for i in ids:
+            iv = int(i)
+            if 0 <= iv < self.vocab_size:
+                clean.append(iv)
         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)
+            clean = [i for i in clean if i not in self._special_ids]
+        if not clean:
+            return ""
+        return self.enc.decode(clean)
+
+    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]
+
+    # -- special-token helpers (extra, not in the SP wrapper) -------------
+
+    def encode_special_id(self, name: str) -> int:
+        """Look up the id of a named special token (e.g. '<|user_start|>')."""
+        return int(self.enc.encode_single_token(name))
 
-    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]
+    # -- config round-trip ------------------------------------------------
 
     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')
+            "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."""
+    def from_config(
+        cls,
+        model_path: str,
+        config_path: Optional[str] = None,
+    ) -> "TokenizerWrapper":
         tok = cls(model_path)
         if config_path and Path(config_path).exists():
-            cfg = json.loads(Path(config_path).read_text(encoding='utf-8'))
+            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')
+                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}')
+                raise ValueError(f"Tokenizer/config mismatch: {mismatches}")
         return tok