File size: 2,997 Bytes
c9721e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """Native ByteLevel BPE tokenizer and explicit role-token chat formatting."""
from __future__ import annotations
from pathlib import Path
SPECIAL=['<|unk|>','<|pad|>','<|bos|>','<|eos|>','<|im_start|>','<|im_end|>','<|eot_id|>']
TOKENIZER_FILE='tokenizer.json'
def load_tokenizer():
from tokenizers import Tokenizer
from .settings import path
p=path('tokenizer.json')
if not p.exists():raise FileNotFoundError('Tokenizer missing: run python -m velqora tokenizer')
return Tokenizer.from_file(str(p))
def train_tokenizer(corpus_file,vocab_size=32000):
from tokenizers import Tokenizer,models,pre_tokenizers,decoders,trainers
from .settings import path
tokenizer=Tokenizer(models.BPE(unk_token='<|unk|>'))
tokenizer.pre_tokenizer=pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder=decoders.ByteLevel()
trainer=trainers.BpeTrainer(vocab_size=vocab_size,min_frequency=2,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
special_tokens=SPECIAL,show_progress=True)
tokenizer.train([str(corpus_file)],trainer)
tokenizer.save(str(path(TOKENIZER_FILE)))
print(f'Saved tokenizer with {tokenizer.get_vocab_size()} tokens')
return tokenizer
def require_ids(tokenizer):
return {key: token_id(tokenizer,key) for key in SPECIAL}
def token_id(tokenizer,text):
x=tokenizer.token_to_id(text)
if x is None:raise ValueError(f'Missing special token: {text}')
return x
def text_ids(tokenizer,text):
return tokenizer.encode(text,add_special_tokens=False).ids
def chat_ids(tokenizer,messages,add_assistant_prefix=False,system=None):
"""Return (ids, loss_mask). Targets are assistant text/eot, not prompts or role markers."""
spec=require_ids(tokenizer)
start,end,eot=spec['<|im_start|>'],spec['<|im_end|>'],spec['<|eot_id|>']
result=[];mask=[]
if system is not None:
messages=[{'role':'system','content':system}]+list(messages)
for m in messages:
role=m.get('role')
content=m.get('content')
if role not in ('user','assistant','system') or not isinstance(content,str):continue
prefix=[start]+text_ids(tokenizer,role+'\n')
content_tokens=text_ids(tokenizer,content)
suffix=[eot if role=='assistant' else end]
result.extend(prefix);mask.extend([0]*len(prefix))
result.extend(content_tokens);mask.extend([1 if role=='assistant' else 0]*len(content_tokens))
result.extend(suffix);mask.extend([1 if role=='assistant' else 0])
if add_assistant_prefix:
tail=[start]+text_ids(tokenizer,'assistant\n')
result.extend(tail);mask.extend([0]*len(tail))
return result,mask
def decode_reply(tokenizer,ids):
spec=require_ids(tokenizer)
clean=[]
for i in ids:
if i in (spec['<|eot_id|>'],spec['<|im_end|>'],spec['<|eos|>']):break
clean.append(i)
return tokenizer.decode(clean,skip_special_tokens=True)
|