feat: catch correct antiprompt according to the template, ref #65

Author: pminev

ac-local-plugin/code/LocalLlama.cpp

@@ -53,6 +53,7 @@ class ChatSession {
     IoEndpoint& m_io;
 
     std::string m_roleUser;
+    std::string m_userPrefix;
     std::string m_roleAsistant;
     std::unique_ptr<llama::ChatFormat> m_chatFormat;
     std::vector<llama::ChatMsg> m_chatMessages;
@@ -77,6 +78,7 @@ class ChatSession {
             m_chatFormat = std::make_unique<llama::ChatFormat>(modelChatParams.chatTemplate);
         } else {
             modelChatParams.chatTemplate = chatTemplate;
+            modelChatParams.roleAssistant = params.roleAssistant.value();
             m_chatFormat = std::make_unique<llama::ChatFormat>(std::move(modelChatParams));
         }
 
@@ -85,6 +87,22 @@ class ChatSession {
 
         m_roleUser = params.roleUser;
         m_roleAsistant = params.roleAssistant;
+
+        auto trim = [](const std::string& str) {
+            auto begin = std::find_if_not(str.begin(), str.end(), [](unsigned char ch) {
+                return std::isspace(ch);
+            });
+
+            auto end = std::find_if_not(str.rbegin(), str.rend(), [](unsigned char ch) {
+                return std::isspace(ch);
+            }).base();
+
+            return (begin < end) ? std::string(begin, end) : "";
+        };
+
+        // user prefix should a substr before stop
+        m_userPrefix = m_chatFormat->formatMsg({.role = m_roleUser, .text = "stop"}, {}, false);
+        m_userPrefix = trim(m_userPrefix.substr(0, m_userPrefix.find("stop")));
     }
 
     ~ChatSession() {
@@ -132,8 +150,7 @@ class ChatSession {
         }
 
         ac::llama::AntipromptManager antiprompt;
-        auto userPrefix = "\n" + m_roleUser + ": ";
-        antiprompt.addAntiprompt(userPrefix);
+        antiprompt.addAntiprompt(m_userPrefix);
 
         std::string fullResponse;
         Schema::OpGetChatResponse::Return ret;

ac-local-plugin/example/ep-chat.cpp

@@ -38,10 +38,18 @@ int main() try {
 
     const std::string roleUser = "user";
     const std::string roleAssistant = "assistant";
+    const std::string chatTemplate =
+                        "{% for message in messages %}"
+                        "{{ '<|' + message['role'] + '|>\\n' + message['content'] + '<|end|>' + '\\n' }}"
+                        "{% endfor %}"
+                        "{% if add_generation_prompt %}"
+                        "{{ '<|' + assistant_role + '|>\\n' }}"
+                        "{% endif %}";
 
     sid = llama.call<schema::StateModelLoaded::OpStartInstance>({
         .instanceType = "chat",
         .setup = "A chat between a human user and a helpful AI assistant.",
+        .chatTemplate = chatTemplate,
         .roleUser = roleUser,
         .roleAssistant = roleAssistant,
     });

ac-local-plugin/schema/ac/schema/LlamaCpp.hpp

@@ -96,10 +96,11 @@ struct StateModelLoaded {
             v(ubatchSize, "ubatch_size", "Size of the context");
             v(ctrlVectorPaths, "ctrl_vectors", "Paths to the control vectors.");
             v(setup, "setup", "Initial setup prompt for the chat session");
-            v(chatTemplate, "chat_template", "Chat template to use. If empty will use the model default");
+            v(chatTemplate, "chat_template", "Valid Jinja chat template to use. If empty will use the model default");
             v(bosOverride, "bos_override", "BOS token to use with the custom template. If empty will use the model default");
             v(eosOverride, "eos_override", "EOS token to use with the custom template. If empty will use the model default");
             v(roleUser, "role_user", "Role name for the user");
+            v(roleAssistant, "role_assistant", "Role name for the assistant");
         }
     };
 

code/ac/llama/ChatFormat.cpp

@@ -111,6 +111,7 @@ class JinjaImpl final : public ChatFormat::impl {
     JinjaImpl(ChatFormat::Params params)
     {
         m_templateStr = std::move(params.chatTemplate);
+        m_assistantRole = std::move(params.roleAssistant);
 
         try {
             m_minjaTemplate = std::make_unique<minja::chat_template>(m_templateStr, params.bosToken, params.eosToken);
@@ -121,9 +122,9 @@ class JinjaImpl final : public ChatFormat::impl {
 
     ~JinjaImpl() {}
 
-    virtual std::string formatChat(std::span<const ChatMsg> chat, bool /*addAssistantPrompt*/) const override {
+    virtual std::string formatChat(std::span<const ChatMsg> chat, bool addAssistantPrompt) const override {
         auto [jChat, size] = ac2jsonChatMessages(chat);
-        return size == 0 ? std::string{} : applyJinja(jChat);
+        return size == 0 ? std::string{} : applyJinja(jChat, addAssistantPrompt);
     }
 
     virtual std::string formatMsg(const ChatMsg& msg, std::span<const ChatMsg> history, bool addAssistantPrompt) const override {
@@ -132,10 +133,10 @@ class JinjaImpl final : public ChatFormat::impl {
         }
 
         auto [jchat, size] = ac2jsonChatMessages(history);
-        auto fmtHistory = applyJinja(jchat);
+        auto fmtHistory = applyJinja(jchat, addAssistantPrompt);
 
         jchat.push_back({{"role", msg.role}, {"content", msg.text}});
-        auto fmtNew = applyJinja(jchat);
+        auto fmtNew = applyJinja(jchat, addAssistantPrompt);
 
         return fmtNew.substr(fmtHistory.size());
     }
@@ -156,19 +157,22 @@ class JinjaImpl final : public ChatFormat::impl {
         return {messages, size};
     }
 
-    std::string applyJinja(acnl::json jChat) const {
+    std::string applyJinja(acnl::json jChat, bool addAssistantPrompt) const {
         auto startsWith = [](const std::string& str, const std::string& prefix) {
             return str.rfind(prefix, 0) == 0;
         };
 
         minja::chat_template_inputs tmpl_inputs;
         tmpl_inputs.messages = jChat;
+        tmpl_inputs.add_generation_prompt = addAssistantPrompt;
+        tmpl_inputs.extra_context = {
+            {"assistant_role",  m_assistantRole}
+        };
 
-        minja::chat_template_options tmpl_opts;
         // To avoid double BOS / EOS tokens, we're manually removing begining / trailing tokens
         // instead of using `chat_template_options.use_bos_token = false`, since these tokens
         // may be needed inside the template / between messages too.
-        auto result = m_minjaTemplate->apply(tmpl_inputs, tmpl_opts);
+        auto result = m_minjaTemplate->apply(tmpl_inputs);
         if (startsWith(result, m_minjaTemplate->bos_token())) {
             result = result.substr(m_minjaTemplate->bos_token().size());
         }
@@ -180,6 +184,7 @@ class JinjaImpl final : public ChatFormat::impl {
 
     std::unique_ptr<minja::chat_template> m_minjaTemplate;
     std::string m_templateStr;
+    std::string m_assistantRole;
 };
 
 

code/ac/llama/ChatFormat.hpp

@@ -22,6 +22,7 @@ class AC_LLAMA_EXPORT ChatFormat {
         std::string chatTemplate;
         std::string bosToken;
         std::string eosToken;
+        std::string roleAssistant = "";
     };
 
     explicit ChatFormat(std::string templateStr);