RAM_Altman_GPT / src /tokenizer_utils.py
brunuu's picture
1.0
7de15fc verified
Raw History Blame Contribute Delete
955 Bytes
from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
from transformers import GPT2TokenizerFast
from config import CHAT_TEMPLATE, DATA_FILE, END_TOKEN, VOCAB_SIZE
def build_tokenizer() -> tuple[GPT2TokenizerFast, int]:
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder = decoders.ByteLevel()
# Keep the vocabulary small
tokenizer.train([DATA_FILE], trainers.BpeTrainer(
vocab_size=VOCAB_SIZE,
special_tokens=[END_TOKEN], # makes the model learn to stop
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
show_progress=False,
))
hf_tokenizer = GPT2TokenizerFast(
tokenizer_object=tokenizer,
bos_token=END_TOKEN,
eos_token=END_TOKEN,
unk_token=END_TOKEN,
)
hf_tokenizer.chat_template = CHAT_TEMPLATE
return hf_tokenizer, len(hf_tokenizer)