From e23b5c8689d5aea56e1f6ecad2eaee66e3aa8081 Mon Sep 17 00:00:00 2001
From: HanishKVC <hanishkvc@gmail.com>
Date: Sat, 20 Apr 2024 23:26:16 +0530
Subject: [PATCH] ChatOn+Main: Cleanup the Requested ChatOn ReversePrompt
 handling

Avoid the use of the seperate vector, which inturn is copied to
the main vector on return. Now directly pass the main reverse
prompt vector and inturn directly add to passed vector.

Also keep data and return status seperate. Explicitly identify
a unknown template_id situation and return failure status.
---
 common/chaton.hpp      | 18 ++++++++++--------
 examples/main/main.cpp |  7 +------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/common/chaton.hpp b/common/chaton.hpp
index 91b3d4802..b33027669 100644
--- a/common/chaton.hpp
+++ b/common/chaton.hpp
@@ -36,16 +36,18 @@ inline bool llama_chat_apply_template_simple(
 }
 
 // return what should be the reverse prompt for the given template id
-// ie possible end text tag(s) of specified model type's chat query response
-inline std::vector<std::string> llama_chat_reverse_prompt(std::string &template_id) {
-    std::vector<std::string> rends;
-
+// ie possible end text tag(s) of specified model type's chat query response.
+// Note that It adds these reverse prompts to any that may already exist in the passed vector.
+inline bool llama_chat_reverse_prompt(std::string &template_id, std::vector<std::string> &rprompts) {
     if (template_id == "chatml") {
-        rends.push_back("<|im_start|>user\n");
+        rprompts.push_back("<|im_start|>user\n");
     } else if (template_id == "llama2") {
-        rends.push_back("</s>");
+        rprompts.push_back("</s>");
     } else if (template_id == "llama3") {
-        rends.push_back("<|eot_id|>");
+        rprompts.push_back("<|eot_id|>");
+    } else {
+        LOG_TEELN("WARN:%s:Unknown template [%s] requested", __func__, template_id.c_str());
+        return false;
     }
-    return rends;
+    return true;
 }
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index dfd16670e..32bcee9c4 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -373,15 +373,10 @@ int main(int argc, char ** argv) {
     // handle chaton mode, it adds on to any reverse prompt specified explicitly by the user
     if (params.chaton) {
         params.interactive_first = true;
-        std::vector<std::string> resp_ends = llama_chat_reverse_prompt(params.chaton_template_id);
-        if (resp_ends.size() == 0) {
+        if (!llama_chat_reverse_prompt(params.chaton_template_id, params.antiprompt)) {
             LOG_TEELN("ERRR:%s:ChatOn:Unsupported ChatTemplateType:%s", __func__, params.chaton_template_id.c_str());
             exit(1);
         }
-        for (size_t i = 0; i < resp_ends.size(); i++)
-        {
-            params.antiprompt.emplace_back(resp_ends[i]);
-        }
     }
 
     // enable interactive mode if interactive start is specified