Add EOS for vicuna templates

This commit is contained in:
Kai Zau 2024-03-30 14:54:12 +09:00
parent e0f9d9d732
commit e423aa1adf

View file

@ -15717,15 +15717,12 @@ static int32_t llama_chat_apply_template_internal(
// Vicuna 1.1+, Nous Capybara, etc. // Vicuna 1.1+, Nous Capybara, etc.
for (auto message : chat) { for (auto message : chat) {
std::string role(message->role); std::string role(message->role);
if (role == "user") {
ss << "USER: ";
} else if (role == "assistant") {
ss << "ASSISTANT: ";
}
ss << message->content << "\n";
if (role == "system") { if (role == "system") {
// Extra newline after system message ss << message->content << "\n\n";
ss << "\n"; } else if (role == "user") {
ss << "USER: " << message->content << "\n";
} else if (role == "assistant") {
ss << "ASSISTANT: " << message->content << "</s>\n";
} }
} }
if (add_ass) { if (add_ass) {
@ -15736,13 +15733,12 @@ static int32_t llama_chat_apply_template_internal(
for (auto message : chat) { for (auto message : chat) {
std::string role(message->role); std::string role(message->role);
if (role == "system") { if (role == "system") {
ss << "SYSTEM: "; ss << "SYSTEM: " << message->content << "\n";
} else if (role == "user") { } else if (role == "user") {
ss << "USER: "; ss << "USER: " << message->content << "\n";
} else if (role == "assistant") { } else if (role == "assistant") {
ss << "ASSISTANT: "; ss << "ASSISTANT: " << message->content << "</s>\n";
} }
ss << message->content << "\n";
} }
if (add_ass) { if (add_ass) {
ss << "ASSISTANT:"; ss << "ASSISTANT:";