| import gradio as gr | |
| from transformers import AutoTokenizer | |
| import json | |
| tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta") | |
| demo_conversation = """[ | |
| {"role": "system", "content": "You are a helpful chatbot."}, | |
| {"role": "user", "content": "Hi there!"}, | |
| {"role": "assistant", "content": "Hello, human!"}, | |
| {"role": "user", "content": "Can I ask a question?"} | |
| ]""" | |
| default_template = """{% for message in messages %} | |
| {{ "<|im_start|>" + message["role"] + "\\n" + message["content"] + "<|im_end|>\\n" }} | |
| {% endfor %} | |
| {% if add_generation_prompt %} | |
| {{ "<|im_start|>assistant\\n" }} | |
| {% endif %}""" | |
| description_text = """### This space is a helper app for writing [Chat Templates](https://huggingface.co/docs/transformers/main/en/chat_templating). | |
| ### When you're happy with the outputs from your template, you can use the code block at the end to add it to a PR!""" | |
| def apply_chat_template(template, test_conversation, add_generation_prompt, cleanup_whitespace): | |
| if cleanup_whitespace: | |
| template = "".join([line.strip() for line in template.split('\n')]) | |
| tokenizer.chat_template = template | |
| outputs = [] | |
| conversation = json.loads(test_conversation) | |
| pr_snippet = ( | |
| "tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)", | |
| f"tokenizer.chat_template = \"{template}\"", | |
| tokenizer.push_to_hub(CHECKPOINT, create_pr=True) | |
| ) | |
| pr_snippet = "\n".join(pr_snippet) | |
| formatted = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=add_generation_prompt) | |
| return formatted, pr_snippet | |
| iface = gr.Interface( | |
| description=description_text, | |
| fn=apply_chat_template, | |
| inputs=[ | |
| gr.TextArea(value=default_template, lines=10, max_lines=30, label="Chat Template"), | |
| gr.TextArea(value=demo_conversation, lines=6, label="Conversation"), | |
| gr.Checkbox(value=False, label="Add generation prompt"), | |
| gr.Checkbox(value=True, label="Cleanup template whitespace"), | |
| ], | |
| outputs=[ | |
| gr.TextArea(label="Formatted conversation")), | |
| gr.TextArea(label="Code snippet to create PR", lines=3, show_copy_button=True)) | |
| ] | |
| iface.launch() | |