File size: 1,928 Bytes
a962f7c | 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 | from tokenizers import Tokenizer as HFTokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
import os
class BPETokenizerWrapper():
def __init__(self, hf_tokenizer):
self.tok = hf_tokenizer
self.word_to_idx = hf_tokenizer.get_vocab()
self.idx_to_word = {v: k for k, v in self.word_to_idx.items()}
def __len__(self):
return self.tok.get_vocab_size()
def encode(self, text):
ids = self.tok.encode(text, add_special_tokens=False).ids
bos = self.word_to_idx["<bos>"]
eos = self.word_to_idx["<eos>"]
return [bos] + ids + [eos]
def decode(self, ids):
special_ids = {self.word_to_idx[t] for t in ["<pad>", "<bos>", "<eos>", "<unk>"]}
filtered = [i for i in ids if i not in special_ids]
return self.tok.decode(filtered)
def build_bpe_tokenizer(combined_texts, vocab_size=32000, save_path="fineweb_bpe.json"):
if os.path.exists(save_path):
tokenizer_bpe = HFTokenizer.from_file(save_path)
else:
tokenizer_bpe = HFTokenizer(BPE(unk_token="<unk>"))
tokenizer_bpe.pre_tokenizer = ByteLevel()
tokenizer_bpe.decoder = ByteLevelDecoder()
trainer = BpeTrainer(
vocab_size=vocab_size,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"],
initial_alphabet=ByteLevel.alphabet(),
)
tokenizer_bpe.train_from_iterator(combined_texts, trainer=trainer)
tokenizer_bpe.save(save_path)
return BPETokenizerWrapper(tokenizer_bpe)
def load_bpe_tokenizer(path):
if not os.path.exists(path):
raise FileNotFoundError(f"{path} missing. Run prepare_fineweb.py first.")
return BPETokenizerWrapper(HFTokenizer.from_file(str(path))) |