File size: 1,026 Bytes
ca7abfa | 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 | import os
from tokenizers import Tokenizer,models,trainers,pre_tokenizers,decoders,processors
def build_tokenizer(vocab_size=32000,files=None,save_dir="tokenizer"):
tok=Tokenizer(models.BPE(unk_token="<unk>"))
tok.pre_tokenizer=pre_tokenizers.ByteLevel(add_prefix_space=False)
tok.decoder=decoders.ByteLevel()
tok.post_processor=processors.ByteLevel(trim_offsets=False)
specials=["<pad>","<s>","</s>","<unk>"]
tr=trainers.BpeTrainer(vocab_size=vocab_size,min_frequency=2,
special_tokens=specials,show_progress=True)
if files:
tok.train(files,tr)
else:
tmp="/tmp/_meno_tok.txt"
with open(tmp,'w') as f: f.write("Tokenizer data placeholder. "*500+"\n")
tok.train([tmp],tr)
os.remove(tmp)
os.makedirs(save_dir,exist_ok=True)
tok.save(os.path.join(save_dir,"tokenizer.json"))
return tok
def load_tokenizer(path="tokenizer"):
return Tokenizer.from_file(os.path.join(path,"tokenizer.json"))
|