From b938cccf05ff7b6ae58f7d79f5f7d17591950648 Mon Sep 17 00:00:00 2001
From: chiensen <113402074+chiensen@users.noreply.github.com>
Date: Tue, 12 Dec 2023 09:44:04 +0800
Subject: [PATCH] Add Pygmalion chat format (#986)

---
 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 24ef9e2..62e3783 100644
--- a/llama_cpp/llama_chat_format.py
+++ b/llama_cpp/llama_chat_format.py
@@ -637,6 +637,23 @@ def format_zephyr(
     _prompt = _format_chatml(system_message, _messages, _sep)
     return ChatFormatterResponse(prompt=_prompt, stop=_sep)
 
+
+@register_chat_format("pygmalion")
+def format_pygmalion(
+    messages: List[llama_types.ChatCompletionRequestMessage],
+    **kwargs: Any,
+) -> ChatFormatterResponse:
+    system_template = """<|system|>{system_message}"""
+    system_message = _get_system_message(messages)
+    system_message = system_template.format(system_message=system_message)
+    _roles = dict(user="<|user|>", assistant="<|model|>")
+    _sep = "\n"
+    _messages = _map_roles(messages, _roles)
+    _messages.append((_roles["assistant"], None))
+    _prompt = _format_chatml(system_message, _messages, _sep)
+    return ChatFormatterResponse(prompt=_prompt, stop=_sep)
+
+
 @register_chat_format("chatml")
 def format_chatml(
     messages: List[llama_types.ChatCompletionRequestMessage],