Add defaults

app.py

@@ -11,13 +11,88 @@ demo_conversation = """[
     {"role": "user", "content": "Can I ask a question?"}
 ]"""
 
-default_template = """{% for message in messages %}
+chat_templates = {
+    "chatml": """{% for message in messages %}
     {{ "<|im_start|>" + message["role"] + "\\n" + message["content"] + "<|im_end|>\\n" }}
 {% endfor %}
 {% if add_generation_prompt %}
     {{ "<|im_start|>assistant\\n" }}
+{% endif %}""",
+    "zephyr": """{% for message in messages %}
+{% if message['role'] == 'user' %}
+{{ '<|user|>\n' + message['content'] + eos_token }}
+{% elif message['role'] == 'system' %}
+{{ '<|system|>\n' + message['content'] + eos_token }}
+{% elif message['role'] == 'assistant' %}
+{{ '<|assistant|>\n'  + message['content'] + eos_token }}
+{% endif %}
+{% if loop.last and add_generation_prompt %}
+{{ '<|assistant|>' }}
+{% endif %}
+{% endfor %}""",
+    "llama": """{% if messages[0]['role'] == 'system' %}
+{% set loop_messages = messages[1:] %}
+{% set system_message = messages[0]['content'] %}
+{% else %}
+{% set loop_messages = messages %}
+{% set system_message = false %}
+{% endif %}
+{% for message in loop_messages %}
+{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
+{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
+{% endif %}
+{% if loop.index0 == 0 and system_message != false %}
+{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}
+{% else %}
+{% set content = message['content'] %}
+{% endif %}
+{% if message['role'] == 'user' %}
+{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}
+{% elif message['role'] == 'assistant' %}
+{{ ' '  + content.strip() + ' ' + eos_token }}
+{% endif %}
+{% endfor %}"""
+    "alpaca": """{% for message in messages %}
+{% if message['role'] == 'system' %}
+{{ message['content'] + '\n\n' }}
+{% elif message['role'] == 'user' %}
+{{ '### Instruction:\n' + message['content'] + '\n\n' }}
+{% elif message['role'] == 'assistant' %}
+{{ '### Response:\n'  + message['content'] + '\n\n' }}
+{% endif %}
+{% if loop.last and add_generation_prompt %}
+{{ '### Response:\n' }}
+{% endif %}
+{% endfor %}""",
+    "vicuna": """{% for message in messages %}
+{% if message['role'] == 'system' %}
+{{ message['content'] + '\n' }}
+{% elif message['role'] == 'user' %}
+{{ 'USER:\n' + message['content'] + '\n' }}
+{% elif message['role'] == 'assistant' %}
+{{ 'ASSISTANT:\n'  + message['content'] + '\n' }}
+{% endif %}
+{% if loop.last and add_generation_prompt %}
+{{ 'ASSISTANT:\n' }}
+{% endif %}
+{% endfor %}""",
+    "falcon": """{% for message in messages %}
+{% if not loop.first %}
+{{ '\n' }}
+{% endif %}
+{% if message['role'] == 'system' %}
+{{ 'System: ' }}
+{% elif message['role'] == 'user' %}
+{{ 'User: ' }}
+{% elif message['role'] == 'assistant' %}
+{{ 'Falcon: ' }}
+{% endif %}
+{{ message['content'] }}
+{% endfor %}
+{% if add_generation_prompt %}
+{{ '\n' + 'Falcon:' }}
 {% endif %}"""
-
+}
 description_text = """# Chat Template Creator
 
 ### This space is a helper app for writing [Chat Templates](https://huggingface.co/docs/transformers/main/en/chat_templating).
@@ -40,6 +115,9 @@ def apply_chat_template(template, test_conversation, add_generation_prompt, clea
     formatted = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=add_generation_prompt)
     return formatted, pr_snippet
 
+def load_template(template_name):
+    template_in.value = chat_templates[template_name]
+
 with gr.Blocks() as demo:
 
     gr.Markdown(description_text)
@@ -47,13 +125,13 @@ with gr.Blocks() as demo:
     with gr.Row():
         gr.Markdown("### Pick an existing template to start:")
     with gr.Row():
-        gr.Button("ChatML")
-        gr.Button("Zephyr")
-        gr.Button("LLaMA")
+        load_chatml = gr.Button("ChatML")
+        load_zephyr = gr.Button("Zephyr")
+        load_llama = gr.Button("LLaMA")
     with gr.Row():
-        gr.Button("Alpaca")
-        gr.Button("Vicuna")
-        gr.Button("Something else")
+        load_alpaca = gr.Button("Alpaca")
+        load_vicuna = gr.Button("Vicuna")
+        load_falcon = gr.Button("Falcon")
     with gr.Row():
         with gr.Column():
             template_in = gr.TextArea(value=default_template, lines=10, max_lines=30, label="Chat Template")
@@ -64,10 +142,17 @@ with gr.Blocks() as demo:
         with gr.Column():
             formatted_out = gr.TextArea(label="Formatted conversation")
             code_snippet_out = gr.TextArea(label="Code snippet to create PR", lines=3, show_label=True, show_copy_button=True)
-    submit.click(fn=apply_chat_template, 
-              inputs=[template_in, conversation_in, generation_prompt_check, cleanup_whitespace_check],
-              outputs=[formatted_out, code_snippet_out]
-              )
+submit.click(fn=apply_chat_template, 
+            inputs=[template_in, conversation_in, generation_prompt_check, cleanup_whitespace_check],
+            outputs=[formatted_out, code_snippet_out]
+            )
+load_chatml.click(fn=load_template, "chatml")
+load_zephyr.click(fn=load_template, "zephyr")
+load_llama.click(fn=load_template, "llama")
+load_alpaca.click(fn=load_template, "alpaca")
+load_vicuna.click(fn=load_template, "vicuna")
+load_falcon.click(fn=load_template, "falcon")
+
 
 demo.launch()