from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("./llama2", use_fast=False)

# add three extra different type tokens:
tok.add_tokens(['|<im_start1>|'], special_tokens=False)
tok.add_tokens(['|<im_start2>|'], special_tokens=True)
tok.add_special_tokens({"pad_token": '|<im_start3>|'})

print(tok)
# output:
"""
LlamaTokenizer(name_or_path='./llama2/', vocab_size=32000, model_max_length=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '|<im_start3>|', 'additional_special_tokens': ['|<im_start2>|']}, clean_up_tokenization_spaces=False),  added_tokens_decoder={
        0: AddedToken("<unk>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
        1: AddedToken("<s>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
        2: AddedToken("</s>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
        32000: AddedToken("|<im_start1>|", rstrip=False, lstrip=False, single_word=False, normalized=True, special=False),
        32001: AddedToken("|<im_start2>|", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
        32002: AddedToken("|<im_start3>|", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
}
"""

print(tok.tokenize("|<im_start1>|system"))
# ['▁', '|<im_start1>|', 'system']

print(tok.tokenize("|<im_start2>|system"))
# ['|<im_start2>|', 'system']

print(tok.tokenize("|<im_start3>|system"))
# ['|<im_start3>|', 'system']

print(tok.decode(tok("|<im_start1>|system").input_ids))
# '<s>  |<im_start1>| system'

print(tok.decode(tok("|<im_start2>|system").input_ids))
# '<s> |<im_start2>| system'

print(tok.decode(tok("|<im_start3>|system").input_ids))
# '<s>|<im_start3>|system'

   tok.add_tokens(['|<im_start2>|'], special_tokens=True)
   tok.add_special_tokens({"pad_token": '|<im_start3>|'})