File size: 1,933 Bytes
8f61432 | 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 | import json
from pathlib import Path
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
VOCAB_SIZE = 32000
SPECIAL_TOKENS = ["<pad>", "<unk>", "<bos>", "<eos>"]
class BPETokenizer:
def __init__(self, tokenizer: Tokenizer):
self._tok = tokenizer
@property
def vocab_size(self) -> int:
return self._tok.get_vocab_size()
def encode(self, text: str) -> list[int]:
return self._tok.encode(text).ids
def decode(self, ids) -> str:
return self._tok.decode(list(ids))
def save(self, path: Path):
self._tok.save(str(path))
@classmethod
def load(cls, path: Path) -> "BPETokenizer":
return cls(Tokenizer.from_file(str(path)))
@classmethod
def build_from_text(cls, text: str, vocab_size: int = VOCAB_SIZE) -> "BPETokenizer":
tok = Tokenizer(models.BPE(unk_token="<unk>"))
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
tok.decoder = decoders.ByteLevel()
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=SPECIAL_TOKENS,
min_frequency=2,
)
tok.train_from_iterator(_chunk(text), trainer=trainer)
return cls(tok)
@classmethod
def build_from_files(cls, paths: list[Path], vocab_size: int = VOCAB_SIZE) -> "BPETokenizer":
tok = Tokenizer(models.BPE(unk_token="<unk>"))
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
tok.decoder = decoders.ByteLevel()
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=SPECIAL_TOKENS,
min_frequency=2,
)
tok.train([str(p) for p in paths], trainer=trainer)
return cls(tok)
def _chunk(text: str, size: int = 1_000_000):
for i in range(0, len(text), size):
yield text[i:i + size]
CharTokenizer = BPETokenizer
|