Skip to content

Commit

Permalink
Add Cohere Command R chat template (#411)
Browse files Browse the repository at this point in the history
  • Loading branch information
noah-kim-theori authored May 7, 2024
1 parent 14522e6 commit 4a1c6ae
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions python/sglang/lang/chat_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,19 @@ def get_chat_template_by_model_path(model_path):
)
)

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):
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 4a1c6ae

Please sign in to comment.