From bbe64fe0699bf7bac3927e5763c25300c5d9c3bf Mon Sep 17 00:00:00 2001
From: Gabe Goodhart <ghart@us.ibm.com>
Date: Mon, 4 Nov 2024 16:55:18 -0700
Subject: [PATCH] feat(granite*): Add granite chat template

Branch: GraniteThreeSupport

This is a port of the work done in llama.cpp with a slight tweak for the
tool call response:
https://github.com/ggerganov/llama.cpp/pull/10013

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
---
 llama.cpp/llama.cpp | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/llama.cpp/llama.cpp b/llama.cpp/llama.cpp
index 16defdd3c4..60a6a3ec1b 100644
--- a/llama.cpp/llama.cpp
+++ b/llama.cpp/llama.cpp
@@ -19291,6 +19291,19 @@ static int32_t llama_chat_apply_template_internal(
         if (add_ass) {
             ss << "[|assistant|]";
         }
+    } else if (tmpl == "granite" || tmpl == "granitemoe" || tmpl_contains("<|start_of_role|>")) {
+        // IBM Granite template
+        for (const auto & message : chat) {
+            std::string role(message->role);
+            ss << "<|start_of_role|>" << role << "<|end_of_role|>";
+            if (role == "assistant_tool_call") {
+                ss << "<|tool_call|>";
+            }
+            ss << message->content << "<|end_of_text|>\n";
+        }
+        if (add_ass) {
+            ss << "<|start_of_role|>assistant<|end_of_role|>\n";
+        }
     } else {
         // template not supported
         return -1;