From 3fcaf199671a70b0293362e962c03bd665fe49dc Mon Sep 17 00:00:00 2001
From: HanishKVC <hanishkvc@gmail.com>
Date: Tue, 14 May 2024 00:58:16 +0530
Subject: [PATCH] ChatON+:Multi4Single: applyGlobalIfAny flag wrt templating
 api

Given that now the multi chat templating logic itself is used to
apply chat templating/tagging to a single chat message, so give
flexibility of deciding whether global tags if any should be
applied or not wrt the core tagging logic.

examples/main inturn updated to not apply global tags if any wrt
the system message. Also the user messages already dont apply
global tags if any, as its currently implemented to build on the
existing in-prefix/suffix and anitprompt flow.
---
 common/chaton.hpp      | 23 +++++++++++++++--------
 examples/main/main.cpp |  2 +-
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/common/chaton.hpp b/common/chaton.hpp
index 94c8b4c8c9..3e69c5fdaa 100644
--- a/common/chaton.hpp
+++ b/common/chaton.hpp
@@ -531,6 +531,7 @@ inline bool chaton_tmpl_apply_ex(
         const std::string &tmpl,
         const std::vector<const llama_chat_message *> &msgs,
         bool alertAssistantAtEnd,
+        bool applyGlobalIfAny,
         std::string &tagged,
         std::string &types,
         std::vector<int32_t> &lens
@@ -539,8 +540,10 @@ inline bool chaton_tmpl_apply_ex(
         return false;
     }
     ChatParts cp = {};
-    std::string globalBegin = chaton_tmpl_role_getkeys(tmpl, K_GLOBAL, {K_BEGIN});
-    cp.add_part(ChatParts::S, globalBegin);
+    if (applyGlobalIfAny) {
+        std::string globalBegin = chaton_tmpl_role_getkeys(tmpl, K_GLOBAL, {K_BEGIN});
+        cp.add_part(ChatParts::S, globalBegin);
+    }
     int cntSystem = 0;
     int cntUser = 0;
     int cntOthers = 0;
@@ -590,8 +593,10 @@ inline bool chaton_tmpl_apply_ex(
         auto assistantBeginPrefix = chaton_tmpl_role_getkeys(tmpl, K_ASSISTANT, {K_BEGIN, K_PREFIX});
         cp.add_part(ChatParts::S, assistantBeginPrefix);
     }
-    auto globalEnd = chaton_tmpl_role_getkeys(tmpl, K_GLOBAL, {K_END});
-    cp.add_part(ChatParts::S, globalEnd);
+    if (applyGlobalIfAny) {
+        auto globalEnd = chaton_tmpl_role_getkeys(tmpl, K_GLOBAL, {K_END});
+        cp.add_part(ChatParts::S, globalEnd);
+    }
     cp.dump();
     tagged = cp.str();
     LOGLN("DBUG:%s:%s:%s", __func__, tmpl.c_str(), tagged.c_str());
@@ -608,11 +613,12 @@ inline int32_t chaton_tmpl_apply(
         const std::string &tmpl,
         const std::vector<const llama_chat_message *> &msgs,
         bool alertAssistantAtEnd,
+        bool applyGlobalIfAny,
         std::string &tagged
         ) {
     std::string types;
     std::vector<int32_t> lens;
-    if (!chaton_tmpl_apply_ex(tmpl, msgs, alertAssistantAtEnd, tagged, types, lens)) {
+    if (!chaton_tmpl_apply_ex(tmpl, msgs, alertAssistantAtEnd, applyGlobalIfAny, tagged, types, lens)) {
         return -1;
     }
     return tagged.size();
@@ -629,12 +635,13 @@ inline size_t chaton_tmpl_apply_single(
         const std::string &role,
         const std::string &content,
         bool alertAssistantAtEnd,
+        bool applyGlobalIfAny,
         std::string &tagged
         ) {
     std::string types;
     std::vector<int32_t> lens;
     llama_chat_message cm {role.c_str(), content.c_str()};
-    if (!chaton_tmpl_apply_ex(tmpl, {&cm}, alertAssistantAtEnd, tagged, types, lens)) {
+    if (!chaton_tmpl_apply_ex(tmpl, {&cm}, alertAssistantAtEnd, applyGlobalIfAny, tagged, types, lens)) {
         return -1;
     }
     return tagged.size();
@@ -669,7 +676,7 @@ inline int32_t chaton_tmpl_apply_capi(
         vMsgs.push_back(&msgs[i]);
     }
     std::string taggedMsgs;
-    int32_t taggedLength = chaton_tmpl_apply(tmpl, vMsgs, alertAssistantAtEnd, taggedMsgs);
+    int32_t taggedLength = chaton_tmpl_apply(tmpl, vMsgs, alertAssistantAtEnd, true, taggedMsgs);
     if (taggedLength < 0) {
         return taggedLength;
     }
@@ -714,7 +721,7 @@ inline int32_t chaton_tmpl_apply_ex_capi(
     std::string taggedMsgs;
     std::string types;
     std::vector<int32_t> lens;
-    if (!chaton_tmpl_apply_ex(tmpl, vMsgs, alertAssistantAtEnd, taggedMsgs, types, lens)) {
+    if (!chaton_tmpl_apply_ex(tmpl, vMsgs, alertAssistantAtEnd, true, taggedMsgs, types, lens)) {
         return -1;
     }
     int32_t taggedLength = taggedMsgs.size();
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index 1a982d4ce3..f86592fc7d 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -266,7 +266,7 @@ int main(int argc, char ** argv) {
             params.prompt = "<|im_start|>system\n" + params.prompt + "<|im_end|>";
         }
         if (params.chaton) {
-            chaton_tmpl_apply_single(params.chaton_template_id, K_SYSTEM, params.prompt, params.prompt);
+            chaton_tmpl_apply_single(params.chaton_template_id, K_SYSTEM, params.prompt, false, false, params.prompt);
         }
         embd_inp = ::llama_tokenize(ctx, params.prompt, true, true);
     } else {