Unverified Commit 4a1c6ae2 authored by YoungJoong Noah Kim's avatar YoungJoong Noah Kim Committed by GitHub
Browse files

Add Cohere Command R chat template (#411)

parent 14522e6a
...@@ -227,6 +227,19 @@ register_chat_template( ...@@ -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 @register_chat_template_matching_function
def match_dbrx(model_path: str): def match_dbrx(model_path: str):
...@@ -287,6 +300,13 @@ def match_gemma_it(model_path: str): ...@@ -287,6 +300,13 @@ def match_gemma_it(model_path: str):
return get_chat_template("gemma-it") 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__": if __name__ == "__main__":
messages = [ messages = [
{"role": "system", "content": None}, # None means default {"role": "system", "content": None}, # None means default
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment