| |
| 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)) |
|
|
