From 0407527f46367888de4518df97341ed0ba6c3fa5 Mon Sep 17 00:00:00 2001
From: DrewWalkup <drew.walkup@gmail.com>
Date: Sun, 17 Mar 2024 14:14:24 -0700
Subject: [PATCH] add user-assistant chat format

---
 llama_cpp/llama_chat_format.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py
index 81ca5520b..70c673886 100644
--- 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