From 07a14c82c8ca2d03cc94b9bc7048a9776d2067f3 Mon Sep 17 00:00:00 2001
From: Fedor Moiseev <femoiseev@gmail.com>
Date: Thu, 28 Dec 2023 07:35:31 +0000
Subject: [PATCH] Add Saiga chat format.

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

diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py
index 037f96a2dd..8b7a74204a 100644
--- a/llama_cpp/llama_chat_format.py
+++ b/llama_cpp/llama_chat_format.py
@@ -704,6 +704,28 @@ def format_openchat(
     return ChatFormatterResponse(prompt=_prompt, stop=_sep)
 
 
+# Chat format for Saiga models, see more details and available models:
+# https://huggingface.co/collections/IlyaGusev/saiga2-saigamistral-6505d4ccc3d1e53166b636cd
+@register_chat_format("saiga")
+def format_saiga(
+    messages: list[llama_types.ChatCompletionRequestMessage],
+    **kwargs,
+) -> ChatFormatterResponse:
+    _message_template = "<s>{role}\n{content}</s>"
+    _roles = dict(user="user", bot="bot", system="system")
+    _messages = _map_roles(messages, _roles)
+
+    _prompt = ""
+    for role, content in _messages:
+        if content:
+            _prompt += _message_template.format(role=role, content=content)
+        else:
+            _prompt += f"<s>{role}\n"
+    # Response template
+    _prompt += "<s>bot"
+    return ChatFormatterResponse(prompt=_prompt.strip())
+
+
 @register_chat_completion_handler("functionary")
 def functionary_chat_handler(
     llama: llama.Llama,