EmbeddedLLM/chat-template-generation

1from transformers import AutoTokenizer2from fastchat.conversation import get_conv_template3import os4from utils import sanitize_jinja25import difflib6 7def test_llama2_template():8    jinja_lines = []9    with open("../templates/airoboros_v2.jinja2", "r") as f:10        jinja_lines = f.readlines()11 12    print("jinja_lines: ", jinja_lines)13 14    print("sanitized: ", sanitize_jinja2(jinja_lines))15 16    chat = [17    {"role": "system", "content": "You are a helpful assistant."},18    {"role": "user", "content": "Hello, how are you?"},19    {"role": "assistant", "content": "I'm doing great. How can I help you today?"},20    {"role": "user", "content": "I'd like to show off how chat templating works!"},21    ]22 23    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="jondurbin/airoboros-l2-7b-2.2", trust_remote_code=True)24    # f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"25    transformer_prompt = tokenizer.apply_chat_template(chat, tokenize=False)26    print("default template")27    print(transformer_prompt)28    # print(tokenizer.chat_template)29    # tokenizer.eos_token = "<|end_of_turn|>"30    tokenizer.chat_template = sanitize_jinja2(jinja_lines)31 32    transformer_prompt = tokenizer.apply_chat_template(chat, tokenize=False)33    print()34    print("add_generation_prompt False:")35    print(transformer_prompt)36 37    transformer_prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)38    print()39    print("add_generation_prompt True:")40    print(transformer_prompt)41    # transformer_prompt = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True)42    # print(transformer_prompt)43 44 45    print("Fastchat template: ")46    conv = get_conv_template("airoboros_v2")47 48    conv.set_system_message(chat[0]["content"])49    conv.append_message(conv.roles[0], chat[1]["content"])50    conv.append_message(conv.roles[1], chat[2]["content"])51    conv.append_message(conv.roles[0], chat[3]["content"])52    conv.append_message(conv.roles[1], None)53    print(conv.get_prompt())54    matcher = difflib.SequenceMatcher(a=transformer_prompt, b=conv.get_prompt())55    print("Matching Sequences:")56    for match in matcher.get_matching_blocks():57        print("Match             : {}".format(match))58        print("Matching Sequence : {}".format(transformer_prompt[match.a:match.a+match.size]))59    assert transformer_prompt == conv.get_prompt()60 61if __name__  == "__main__":62    test_llama2_template()