From cfcbc7adf3fb7b06a9ed9a33f13176c1710913b7 Mon Sep 17 00:00:00 2001
From: Kai Zau <kaizau@users.noreply.github.com>
Date: Mon, 1 Apr 2024 19:35:20 +0800
Subject: [PATCH] Match openchat template with jinja output

---
 llama.cpp                    | 11 ++++-------
 tests/test-chat-template.cpp |  2 +-
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/llama.cpp b/llama.cpp
index 13dafda4a..07af46444 100644
--- a/llama.cpp
+++ b/llama.cpp
@@ -15701,14 +15701,11 @@ static int32_t llama_chat_apply_template_internal(
         // openchat/openchat-3.5-0106,
         for (auto message : chat) {
             std::string role(message->role);
-            if (role == "user") {
-                ss << "GPT4 Correct User: ";
-            } else if (role == "assistant") {
-                ss << "GPT4 Correct Assistant: ";
+            if (message == chat.front()) {
+                ss << "<s>";
             }
-            // Not documented, but apparently the system message is prepended without prefix:
-            // https://huggingface.co/openchat/openchat_3.5/discussions/5#65448109b4a3f3a2f486fd9d
-            ss << message->content << "<|end_of_turn|>";
+            role[0] = toupper(role[0]);
+            ss << "GPT4 Correct " << role << ": " << message->content << "<|end_of_turn|>";
         }
         if (add_ass) {
             ss << "GPT4 Correct Assistant:";
diff --git a/tests/test-chat-template.cpp b/tests/test-chat-template.cpp
index 440ffe8aa..f25d3b891 100644
--- a/tests/test-chat-template.cpp
+++ b/tests/test-chat-template.cpp
@@ -60,7 +60,7 @@ int main(void) {
         // OrionStarAI/Orion-14B-Chat
         "Human: You are a helpful assistant\n\nHello\n\nAssistant: </s>Hi there</s>Human: Who are you\n\nAssistant: </s>   I am an assistant   </s>Human: Another question\n\nAssistant: </s>",
         // openchat/openchat-3.5-0106
-        "You are a helpful assistant<|end_of_turn|>GPT4 Correct User: Hello<|end_of_turn|>GPT4 Correct Assistant: Hi there<|end_of_turn|>GPT4 Correct User: Who are you<|end_of_turn|>GPT4 Correct Assistant:    I am an assistant   <|end_of_turn|>GPT4 Correct User: Another question<|end_of_turn|>GPT4 Correct Assistant:",
+        "<s>GPT4 Correct System: You are a helpful assistant<|end_of_turn|>GPT4 Correct User: Hello<|end_of_turn|>GPT4 Correct Assistant: Hi there<|end_of_turn|>GPT4 Correct User: Who are you<|end_of_turn|>GPT4 Correct Assistant:    I am an assistant   <|end_of_turn|>GPT4 Correct User: Another question<|end_of_turn|>GPT4 Correct Assistant:",
         // eachadea/vicuna-13b-1.1
         "You are a helpful assistant\n\nUSER: Hello\nASSISTANT: Hi there</s>\nUSER: Who are you\nASSISTANT:    I am an assistant   </s>\nUSER: Another question\nASSISTANT:",
         // Orca-Vicuna