from tokenizers import Tokenizer, models, trainers, pre_tokenizers, processors, decoders
from tokenizers.normalizers import NFC
from transformers import PreTrainedTokenizerFast
from pathlib import Path

SPECIAL_TOKENS = [
    "<|pad|>", "<|bos|>", "<|eos|>", "<|unk|>", "<|sep|>",
    "<|fim_prefix|>", "<|fim_middle|>", "<|fim_suffix|>",
    "<|python|>", "<|javascript|>", "<|typescript|>",
    "<|cpp|>", "<|rust|>", "<|go|>", "<|java|>", "<|bash|>",
]

def get_gpt2_tokenizer_for_code():
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("gpt2")
    tok.pad_token = tok.eos_token
    tok.add_special_tokens({"additional_special_tokens": [
        "<|fim_prefix|>", "<|fim_middle|>", "<|fim_suffix|>",
        "<|python|>", "<|javascript|>", "<|rust|>", "<|go|>",
    ]})
    return tok

def load_tokenizer(save_dir="./tokenizer"):
    return PreTrainedTokenizerFast.from_pretrained(save_dir)