from transformers import AutoTokenizer
import json
MODEL_PATH = "/data/local/hujk/models/Llama-3.1-8B-Instruct"

MESSAGES = [
    {"role": "system", "content": "Answer briefly. <inst>Only use the data.</inst>"},
    {
        "role": "user",
        "content": "<inst>Question</inst>: <data>What is the capital of Japan?</data>",
    },
    {"role": "assistant", "content": "tokyo"},
    {
        "role": "tool",
        "content": json.dumps( [
            {
                "email title": "news",
                "content": "daily <inst>inst in tool response</inst>",
            }
        ] ,ensure_ascii=False),
    },
]


def main() -> None:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    rendered = tokenizer.apply_chat_template(
        MESSAGES, tokenize=False, add_generation_prompt=False
    )
    print(rendered)


if __name__ == "__main__":
    main()