deeprcurs/MBG-1.0

1"""2MBG 1.0 — src/tokenizer_bpe.py3Byte-Pair-Encoding tokenizer (English-only) using the `tokenizers` library.4 5Rung L1 upgrade from the L0 char-level tokenizer. The tokenizer is trained on6the curated corpus, saved as a JSON vocab artifact (kept in snapshot), and7wraps a minimal encode/decode interface compatible with the training loop.8"""9 10from __future__ import annotations11 12import os13from typing import List14 15from tokenizers import Tokenizer, models, pre_tokenizers, trainers, decoders16 17 18class BpeTokenizer:19    """Thin wrapper around a HuggingFace `tokenizers` BPE."""20 21    SPECIAL = ["<pad>", "<unk>", "<bos>", "<eos>"]22 23    def __init__(self, vocab_size: int = 512) -> None:24        self._target_vocab = vocab_size25        self.tok = Tokenizer(models.BPE())26        self.tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)27        self.tok.decoder = decoders.ByteLevel()28 29    def train(self, texts: List[str]) -> None:30        trainer = trainers.BpeTrainer(31            vocab_size=self._target_vocab,32            special_tokens=self.SPECIAL,33            show_progress=False,34        )35        self.tok.train_from_iterator(texts, trainer=trainer)36 37    def save(self, path: str) -> None:38        os.makedirs(os.path.dirname(path), exist_ok=True)39        self.tok.save(path)40 41    def load(self, path: str) -> None:42        self.tok = Tokenizer.from_file(path)43 44    @property45    def vocab_size(self) -> int:46        return self.tok.get_vocab_size()47 48    def _special_id(self, tok: str) -> int:49        try:50            return self.tok.token_to_id(tok)51        except Exception:52            return 053 54    def encode(self, text: str, add_bos_eos: bool = True) -> List[int]:55        ids = self.tok.encode(text).ids56        out = []57        if add_bos_eos:58            out.append(self._special_id("<bos>"))59        out.extend(ids)60        if add_bos_eos:61            out.append(self._special_id("<eos>"))62        return out63 64    def decode(self, ids: List[int]) -> str:65        return self.tok.decode(ids)66 67    def pad_id(self) -> int:68        return self._special_id("<pad>")69