Use chat template

Author: DePasqualeOrg

Applications/LLMEval/ContentView.swift

@@ -159,7 +159,8 @@ class LLMEvaluator {
 
     /// This controls which model loads. `phi3_5_4bit` is one of the smaller ones, so this will fit on
     /// more devices.
-    let modelConfiguration = ModelConfiguration.phi3_5_4bit
+    //    let modelConfiguration = ModelConfiguration.phi3_5_4bit
+    let modelConfiguration = ModelConfiguration.mistral7B4bit
 
     /// parameters controlling the output
     let generateParameters = GenerateParameters(temperature: 0.6)
@@ -212,11 +213,9 @@ class LLMEvaluator {
         do {
             let modelContainer = try await load()
 
-            // augment the prompt as needed
-            let prompt = modelConfiguration.prepare(prompt: prompt)
-
-            let promptTokens = await modelContainer.perform { _, tokenizer in
-                tokenizer.encode(text: prompt)
+            let messages = [["role": "user", "content": prompt]]
+            let promptTokens = try await modelContainer.perform { _, tokenizer in
+                try tokenizer.applyChatTemplate(messages: messages)
             }
 
             // each time you generate you will get something new

Libraries/LLM/Models.swift

@@ -130,7 +130,8 @@ extension ModelConfiguration {
 
     public static let mistral7B4bit = ModelConfiguration(
         id: "mlx-community/Mistral-7B-Instruct-v0.3-4bit",
-        defaultPrompt: "Describe the Swift language."
+        defaultPrompt: "Describe the Swift language.",
+        extraEOSTokens: ["<|im_end|>"]  // !! This doesn't seem to work
     ) { prompt in
         "<s>[INST] \(prompt) [/INST] "
     }