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