diff --git a/families/llama/runtime/chat_templates.cpp b/families/llama/runtime/chat_templates.cpp index cd8e2e5095..43ffea48ac 100644 --- a/families/llama/runtime/chat_templates.cpp +++ b/families/llama/runtime/chat_templates.cpp @@ -25,6 +25,11 @@ std::string apply_phi(const std::string& prompt, bool /*enable_thinking*/) { return "<|user|>\n" + prompt + "<|end|>\n<|assistant|>\n"; } +// Zephyr / TinyLlama share role tags with Phi but end turns with eos (), not <|end|>. +std::string apply_zephyr(const std::string& prompt, bool /*enable_thinking*/) { + return "<|user|>\n" + prompt + "\n<|assistant|>\n"; +} + std::string apply_gemma(const std::string& prompt, bool /*enable_thinking*/) { return "user\n" + prompt + "\nmodel\n"; } @@ -58,9 +63,14 @@ std::string llama_detect_chat_template_format(const std::string& jinja_template) return "chatml"; if (jinja_template.find("[INST]") != std::string::npos) return "mistral"; + // Phi and Zephyr/TinyLlama both use <|user|>/<|assistant|>. Phi is distinguished by + // literal <|end|> turn separators; Zephyr/TinyLlama concatenate eos_token () instead. if (jinja_template.find("<|user|>") != std::string::npos || - jinja_template.find("<|assistant|>") != std::string::npos) - return "phi"; + jinja_template.find("<|assistant|>") != std::string::npos) { + if (jinja_template.find("<|end|>") != std::string::npos) + return "phi"; + return "zephyr"; + } if (jinja_template.find("") != std::string::npos) return "gemma"; if (jinja_template.find("<|start_header_id|>") != std::string::npos) @@ -82,6 +92,8 @@ std::string llama_apply_chat_template(const std::string& format, const std::stri return apply_mistral(prompt, enable_thinking); if (format == "phi") return apply_phi(prompt, enable_thinking); + if (format == "zephyr") + return apply_zephyr(prompt, enable_thinking); if (format == "gemma") return apply_gemma(prompt, enable_thinking); if (format == "llama3") diff --git a/families/llama/tests/cpp/test_llama_chat_template.cpp b/families/llama/tests/cpp/test_llama_chat_template.cpp index 432b3ddf11..dd668c9966 100644 --- a/families/llama/tests/cpp/test_llama_chat_template.cpp +++ b/families/llama/tests/cpp/test_llama_chat_template.cpp @@ -48,6 +48,26 @@ static void test_detect_phi() { check(fmt == "phi", "phi detection"); } +// TinyLlama-1.1B-Chat / Zephyr: same role tags as Phi, but turns end with eos_token (), +// not Phi's literal <|end|>. Must not be classified as phi (issue #1271). +static void test_detect_zephyr_tinyllama() { + std::string tpl = "{% for message in messages %}\n" + "{% if message['role'] == 'user' %}\n" + "{{ '<|user|>\n' + message['content'] + eos_token }}\n" + "{% elif message['role'] == 'system' %}\n" + "{{ '<|system|>\n' + message['content'] + eos_token }}\n" + "{% elif message['role'] == 'assistant' %}\n" + "{{ '<|assistant|>\n' + message['content'] + eos_token }}\n" + "{% endif %}\n" + "{% if loop.last and add_generation_prompt %}\n" + "{{ '<|assistant|>' }}\n" + "{% endif %}\n" + "{% endfor %}"; + auto fmt = trtmc::llama_detect_chat_template_format(tpl); + check(fmt == "zephyr", "zephyr/tinyllama detection"); + check(fmt != "phi", "zephyr/tinyllama must not be phi"); +} + static void test_detect_gemma() { std::string tpl = "{% for message in messages %}{{ message.role }}\n{{ " "message.content }}\n{% endfor %}"; @@ -86,6 +106,18 @@ static void test_apply_phi() { check(result == "<|user|>\nhello<|end|>\n<|assistant|>\n", "phi application"); } +static void test_apply_zephyr_tinyllama() { + auto result = trtmc::llama_apply_chat_template("zephyr", "hello"); + check(result == "<|user|>\nhello\n<|assistant|>\n", "zephyr/tinyllama application"); + // Issue #1271 e2e prompt: must use , not Phi <|end|>. + auto e2e = trtmc::llama_apply_chat_template( + "zephyr", "What is the capital of France? Answer in one word."); + check(e2e == "<|user|>\nWhat is the capital of France? Answer in one word.\n" + "<|assistant|>\n", + "zephyr/tinyllama e2e prompt render"); + check(e2e.find("<|end|>") == std::string::npos, "zephyr render must not inject <|end|>"); +} + static void test_apply_gemma() { auto result = trtmc::llama_apply_chat_template("gemma", "hello"); check(result == "user\nhello\nmodel\n", @@ -104,12 +136,14 @@ int main() { test_detect_chatml(); test_detect_mistral(); test_detect_phi(); + test_detect_zephyr_tinyllama(); test_detect_gemma(); test_detect_llama3(); test_detect_nemotron_h(); test_apply_chatml_no_thinking(); test_apply_mistral_no_thinking_ignored(); test_apply_phi(); + test_apply_zephyr_tinyllama(); test_apply_gemma(); test_apply_llama3(); test_apply_nemotron_h_no_thinking();