From eae6dc990ea68a03f832672f5c3e91a9441f90a8 Mon Sep 17 00:00:00 2001 From: Janne Sinkkonen Date: Thu, 25 Apr 2024 13:21:41 +0300 Subject: [PATCH] Add the Command R chat format --- llama_cpp/llama_chat_format.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py index 17b570a0f..84b009355 100644 --- a/llama_cpp/llama_chat_format.py +++ b/llama_cpp/llama_chat_format.py @@ -1323,7 +1323,24 @@ def format_gemma( _prompt = _format_no_colon_single(system_message="", messages=_messages, sep=_sep) return ChatFormatterResponse(prompt=_prompt, stop=_sep) +# Chat format for Cohere's Command R models (close to chatml but with different tokens). +# See https://docs.cohere.com/docs/prompting-command-r +@register_chat_format("command-r") +def format_commandr( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{system_message}" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _roles = dict(user="<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", assistant="<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>") + _sep = "<|END_OF_TURN_TOKEN|>" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_chatml(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + # Tricky chat formats that require custom chat handlers @@ -2755,4 +2772,4 @@ def chatml_function_calling( }, } - raise ValueError("Automatic streaming tool choice is not supported") \ No newline at end of file + raise ValueError("Automatic streaming tool choice is not supported")