From 6e5c58c3e97530c4597490f583f28cc436c8ef86 Mon Sep 17 00:00:00 2001
From: Ray Strode <rstrode@redhat.com>
Date: Thu, 5 Dec 2024 08:37:23 -0500
Subject: [PATCH] AutocompleteTemplate: Use modelfile template for granite3

granite3 is currently using the "hole filler" template since it
lacks FIM support. This template is getting sent to ollama raw,
so the template from the model file is not getting applied to it.

This commit ensures the proper format is fed to the model, by
moving the hole filler system message to be a global constant, and
using it in the existing holeFillerTemplate and a new
granteHoleFillerTemplate that also applies the appropiate model
specific formatting.
---
 .../templating/AutocompleteTemplate.ts        | 28 ++++++++++++++-----
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/core/autocomplete/templating/AutocompleteTemplate.ts b/core/autocomplete/templating/AutocompleteTemplate.ts
index 9ade88e46a..1547a413e6 100644
--- a/core/autocomplete/templating/AutocompleteTemplate.ts
+++ b/core/autocomplete/templating/AutocompleteTemplate.ts
@@ -237,10 +237,8 @@ Fill in the blank to complete the code block. Your response should include only
   completionOptions: { stop: ["\n"] },
 };
 
-const holeFillerTemplate: AutocompleteTemplate = {
-  template: (prefix: string, suffix: string) => {
-    // From https://github.com/VictorTaelin/AI-scripts
-    const SYSTEM_MSG = `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'. Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.  All completions MUST be truthful, accurate, well-written and correct.
+// From https://github.com/VictorTaelin/AI-scripts
+const HOLE_FILLER_SYSTEM_MSG = `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'. Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.  All completions MUST be truthful, accurate, well-written and correct.
 
 ## EXAMPLE QUERY:
 
@@ -324,8 +322,10 @@ function hypothenuse(a, b) {
 
 <COMPLETION>a ** 2 + </COMPLETION>`;
 
+const holeFillerTemplate: AutocompleteTemplate = {
+  template: (prefix: string, suffix: string) => {
     const fullPrompt =
-      SYSTEM_MSG +
+      HOLE_FILLER_SYSTEM_MSG +
       `\n\n<QUERY>\n${prefix}{{FILL_HERE}}${suffix}\n</QUERY>\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n<COMPLETION>`;
     return fullPrompt;
   },
@@ -334,6 +334,17 @@ function hypothenuse(a, b) {
   },
 };
 
+const graniteHoleFillerTemplate: AutocompleteTemplate = {
+  template: (prefix: string, suffix: string) => {
+    const request = `<QUERY>\n${prefix}{{FILL_HERE}}${suffix}\n</QUERY>\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE.`;
+
+    return `<|start_of_role|>system<|end_of_role|>${HOLE_FILLER_SYSTEM_MSG}<|end_of_text|>\n<|start_of_role|>user<|end_of_role|>${request}<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>`;
+  },
+  completionOptions: {
+    stop: ["<|end_of_text|>", "<|start_of_role|>", "</COMPLETION>"],
+  },
+};
+
 export function getTemplateForModel(model: string): AutocompleteTemplate {
   const lowerCaseModel = model.toLowerCase();
 
@@ -380,11 +391,14 @@ export function getTemplateForModel(model: string): AutocompleteTemplate {
   if (
     lowerCaseModel.includes("gpt") ||
     lowerCaseModel.includes("davinci-002") ||
-    lowerCaseModel.includes("claude") ||
-    lowerCaseModel.includes("granite3")
+    lowerCaseModel.includes("claude")
   ) {
     return holeFillerTemplate;
   }
 
+  if (lowerCaseModel.includes("granite3")) {
+    return graniteHoleFillerTemplate;
+  }
+
   return stableCodeFimTemplate;
 }