add user-assistant chat format

DrewWalkup · DrewWalkup · commit 0407527f4636 · 2024-03-17T14:14:24.000-07:00
diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py
@@ -1037,6 +1037,23 @@ def format_gemma(
     return ChatFormatterResponse(prompt=_prompt, stop=_sep)
 
 
+# Chat format for Nous-Capybara models, see more details:
+# https://huggingface.co/NousResearch/Nous-Capybara-34B
+@register_chat_format("user-assistant")
+def format_user_assistant(
+    messages: List[llama_types.ChatCompletionRequestMessage],
+    **kwargs: Any,
+) -> ChatFormatterResponse:
+    stop_str = "</s>"
+    system_message = _get_system_message(messages)
+    _roles = dict(user="USER:", assistant="ASSISTANT:")
+    _sep = "\n"
+    _messages = _map_roles(messages, _roles)
+    _messages.append((_roles["assistant"], None))
+    _prompt = _format_chatml(system_message, _messages, _sep)
+    return ChatFormatterResponse(prompt=_prompt, stop=stop_str)
+
+
 # Tricky chat formats that require custom chat handlers
 
 
