From 4a1c6ae2ced0f85e1a7414cb13ace9f342a13a56 Mon Sep 17 00:00:00 2001
From: YoungJoong Noah Kim <142862540+noah-kim-theori@users.noreply.github.com>
Date: Tue, 7 May 2024 16:18:15 +0900
Subject: [PATCH] Add Cohere Command R chat template (#411)

---
 python/sglang/lang/chat_template.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/python/sglang/lang/chat_template.py b/python/sglang/lang/chat_template.py
index d91dee365..9dc817818 100644
--- a/python/sglang/lang/chat_template.py
+++ b/python/sglang/lang/chat_template.py
@@ -227,6 +227,19 @@ register_chat_template(
     )
 )
 
+register_chat_template(
+    ChatTemplate(
+        name="c4ai-command-r",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>", "<|END_OF_TURN_TOKEN|>"),
+            "user": ("<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", "<|END_OF_TURN_TOKEN|>"),
+            "assistant": ("<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", "<|END_OF_TURN_TOKEN|>"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+    )
+)
+
 
 @register_chat_template_matching_function
 def match_dbrx(model_path: str):
@@ -287,6 +300,13 @@ def match_gemma_it(model_path: str):
         return get_chat_template("gemma-it")
 
 
+@register_chat_template_matching_function
+def match_c4ai_command_r(model_path: str):
+    model_path = model_path.lower()
+    if "c4ai-command-r" in model_path:
+        return get_chat_template("c4ai-command-r")
+
+
 if __name__ == "__main__":
     messages = [
         {"role": "system", "content": None},  # None means default