diff --git a/src/llama.cpp b/src/llama.cpp index 5a7976900..863a915d7 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -21867,21 +21867,21 @@ static int32_t llama_chat_apply_template_internal( // See: https://github.com/mistralai/cookbook/blob/main/concept-deep-dive/tokenization/templates.md std::string leading_space = (tmpl == "mistral-v1" ? " " : ""); std::string trailing_space = (tmpl != "mistral-v3-tekken" ? " " : ""); - std::string system_message = ""; + bool is_inside_turn = false; for (auto message : chat) { + if (!is_inside_turn) { + ss << leading_space << "[INST]" << trailing_space; + is_inside_turn = true; + } std::string role(message->role); std::string content = trim(message->content); if (role == "system") { - system_message = content; + ss << system_message << "\n\n"; } else if (role == "user") { - ss << leading_space << "[INST]" << trailing_space; - if (!system_message.empty()) { - ss << system_message << "\n\n"; - system_message = ""; - } ss << content << leading_space << "[/INST]"; } else { ss << trailing_space << content << ""; + is_inside_turn = false; } } } else if (tmpl == "llama2" || tmpl == "mistral" || tmpl_contains("[INST]")) { diff --git a/tests/test-chat-template.cpp b/tests/test-chat-template.cpp index 03e897e66..a7608f3dd 100644 --- a/tests/test-chat-template.cpp +++ b/tests/test-chat-template.cpp @@ -154,6 +154,10 @@ int main(void) { return output; }; assert(fmt_sys("chatml") == "<|im_start|>system\nYou are a helpful assistant<|im_end|>\n"); + assert(fmt_sys("mistral-v1") == " [INST] You are a helpful assistant\n\n"); + assert(fmt_sys("mistral-v2") == "[INST] You are a helpful assistant\n\n"); + assert(fmt_sys("mistral-v3") == "[INST] You are a helpful assistant\n\n"); + assert(fmt_sys("mistral-v3-tekken") == "[INST]You are a helpful assistant\n\n"); assert(fmt_sys("llama2") == "[INST] You are a helpful assistant\n"); assert(fmt_sys("gemma") == ""); // for gemma, system message is merged with user message assert(fmt_sys("llama3") == "<|start_header_id|>system<|end_header_id|>\n\nYou are a helpful assistant<|eot_id|>"); @@ -173,6 +177,10 @@ int main(void) { return output; }; assert(fmt_single("chatml") == "\n<|im_start|>user\nHow are you<|im_end|>\n<|im_start|>assistant\n"); + assert(fmt_single("mistral-v1") == " [INST] How are you [/INST]"); + assert(fmt_single("mistral-v2") == "[INST] How are you[/INST]"); + assert(fmt_single("mistral-v3") == "[INST] How are you[/INST]"); + assert(fmt_single("mistral-v3-tekken") == "[INST]How are you[/INST]"); assert(fmt_single("llama2") == "[INST] How are you [/INST]"); assert(fmt_single("gemma") == "\nuser\nHow are you\nmodel\n"); assert(fmt_single("llama3") == "<|start_header_id|>user<|end_header_id|>\n\nHow are you<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n");