Add baichuan-2 chat format (abetlen#936)

caiyesd · web-flow · commit b8f29f4bf041 · 2023-11-22T06:08:06.000-05:00
Signed-off-by: caiyesd &lt;caiyesd@gmail.com&gt;
diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py
@@ -456,6 +456,21 @@ def format_oasst_llama(
     return ChatFormatterResponse(prompt=_prompt)
 
 
+@register_chat_format("baichuan-2")
+def format_baichuan2(
+    messages: List[llama_types.ChatCompletionRequestMessage],
+    **kwargs: Any,
+) -> ChatFormatterResponse:
+    _system_template = "{system_message}"
+    _roles = dict(user="<reserved_106>", assistant="<reserved_107>")
+    _sep = ""
+    system_message = _get_system_message(messages)
+    system_message = _system_template.format(system_message=system_message)
+    _messages = _map_roles(messages, _roles)
+    _messages.append((_roles["assistant"], None))
+    _prompt = _format_no_colon_single(system_message, _messages, _sep)
+    return ChatFormatterResponse(prompt=_prompt)
+
 @register_chat_format("openbuddy")
 def format_openbuddy(
     messages: List[llama_types.ChatCompletionRequestMessage],
