"""Train Ares's first inspectable byte-level BPE tokenizer.
Run from a notebook after installing: pip install tokenizers
"""
import argparse, json
from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers

def main():
    p=argparse.ArgumentParser(); p.add_argument('--data',default='corpus.txt'); p.add_argument('--out',default='ares_tokenizer.json'); p.add_argument('--vocab-size',type=int,default=8192); args=p.parse_args()
    tok=Tokenizer(models.BPE(unk_token='[UNK]'))
    tok.pre_tokenizer=pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder=decoders.ByteLevel()
    trainer=trainers.BpeTrainer(vocab_size=args.vocab_size,min_frequency=2,special_tokens=['[PAD]','[BOS]','[EOS]','[UNK]'],show_progress=True)
    tok.train([args.data],trainer); tok.save(args.out)
    sample='Ares can write code, explain an algorithm, and continue a fictional scene.'
    enc=tok.encode(sample); assert tok.decode(enc.ids)==sample, repr(tok.decode(enc.ids))
    print(json.dumps({'vocab_size':tok.get_vocab_size(),'sample_tokens':len(enc.ids),'round_trip':'passed','output':args.out},indent=2))
if __name__=='__main__': main()