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)