# Têtu Code 2 — inference code (tetu-AI). License: Apache 2.0, see README.md.
from tokenizers import Regex, Tokenizer, decoders, models, pre_tokenizers, trainers
SPECIALS = ['<|pad|>', '<|eos|>', '<|system|>', '<|user|>', '<|assistant|>', '<|tool|>', '<|call|>', '<|/call|>', '<|fim_prefix|>', '<|fim_suffix|>', '<|fim_middle|>', '<|file|>', '<|repo|>', '<|think|>', '<|/think|>', '<|trace|>', '<|/trace|>', '<|tools|>', '<|commit|>', '<|arg|>', '<|end|>'] + [f'<|reserved_{i}|>' for i in range(2, 16)]
PATTERN = '\\r?\\n[ \\t]*|\\p{N}| ?[\\p{L}\\p{M}_]+| ?[^\\s\\p{L}\\p{M}\\p{N}_]+|[ \\t]+|\\s'
VOCAB = 32768

def new_tokenizer():
    tok = Tokenizer(models.BPE())
    tok.pre_tokenizer = pre_tokenizers.Sequence([pre_tokenizers.Split(Regex(PATTERN), behavior='isolated'), pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)])
    tok.decoder = decoders.ByteLevel()
    return tok

def train(texts, vocab=VOCAB, length=None):
    tok = new_tokenizer()
    trainer = trainers.BpeTrainer(vocab_size=vocab, min_frequency=2, special_tokens=SPECIALS, show_progress=False, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), max_token_length=32)
    tok.train_from_iterator(texts, trainer, length=length)
    return tok

def load(path):
    return Tokenizer.from_file(str(path))