from transformers import AutoTokenizer
import transformers.utils.sentencepiece_model_pb2 as sp_model
import json

proto = sp_model.ModelProto()

llama_orgin_tokenizer_path = 'Llama-2-7b-hf' # this a local path
llama_tokenizer_with_added_token_path = 'Llama-2-7b-hf-with-added-token'

with open(f'{llama_orgin_tokenizer_path}/tokenizer.model', 'rb') as fin:
    proto_str = fin.read()
    proto.ParseFromString(proto_str)

# append a user-defined token
adding_token = 'aaabbbcc'
proto.pieces.append(sp_model.ModelProto.SentencePiece(piece=adding_token, score=0.0, type=4)) #type 4 means user defined
with open(f'{llama_tokenizer_with_added_token_path}/tokenizer.model', 'wb') as fout:
    fout.write(proto.SerializeToString())


with open(f'{llama_tokenizer_with_added_token_path}/tokenizer_config.json','w', encoding='utf-8') as fout:
    tokenizer_config = {
        "model_max_length":99999999999999,
        "tokenizer_class": "LlamaTokenizer"
    }
    json.dump(tokenizer_config, fout)



tokenizer_fast = AutoTokenizer.from_pretrained(llama_tokenizer_with_added_token_path, use_fast=True)
tokenizer_slow = AutoTokenizer.from_pretrained(llama_tokenizer_with_added_token_path, use_fast=False)

text = 'I love aaabbbcc.'
print('fast:', tokenizer_fast.tokenize(text))
print('slow:', tokenizer_slow.tokenize(text))

# The outputs:
# fast: ['▁I', '▁love', '▁aa', 'ab', 'bb', 'cc', '.']
# slow: ['▁I', '▁love', '▁', 'aaabbbcc', '.']