Use chat template for Qwen 2 VL

Author: DePasqualeOrg

Applications/VLMEval/ContentView.swift

@@ -331,7 +331,16 @@ class VLMEvaluator {
             MLXRandom.seed(UInt64(Date.timeIntervalSinceReferenceDate * 1000))
 
             let result = try await modelContainer.perform { context in
-                var userInput = UserInput(prompt: prompt, images: [.ciImage(image)])
+                var userInput = UserInput(
+                    messages: [
+                        [
+                            "role": "user",
+                            "content": [
+                                ["type": "text", "text": prompt],
+                                ["type": "image"],
+                            ],
+                        ]
+                    ], images: [.ciImage(image)])
                 userInput.processing.resize = .init(width: 448, height: 448)
 
                 let input = try await context.processor.prepare(input: userInput)

Libraries/MLXLMCommon/UserInput.swift

@@ -4,18 +4,19 @@ import CoreImage
 import Foundation
 import MLX
 
+public typealias Message = [String: Any]
+
 /// Container for raw user input.
 ///
 /// A ``UserInputProcessor`` can convert this to ``LMInput``.
 /// See also ``ModelContext``.
 public struct UserInput: Sendable {
-
     /// Representation of a prompt or series of messages (conversation).
     public enum Prompt: Sendable, CustomStringConvertible {
         case text(String)
-        case messages([[String: String]])
+        case messages([Message])
 
-        public func asMessages() -> [[String: String]] {
+        public func asMessages() -> [Message] {
             switch self {
             case .text(let text):
                 return [["role": "user", "content": text]]
@@ -116,7 +117,7 @@ public struct UserInput: Sendable {
         self.images = images
     }
 
-    public init(messages: [[String: String]], images: [Image] = [Image]()) {
+    public init(messages: [Message], images: [Image] = [Image]()) {
         self.prompt = .messages(messages)
         self.images = images
     }

Libraries/MLXVLM/Models/Idefics3.swift

@@ -805,7 +805,7 @@ public class Idefics3Processor: UserInputProcessor {
     }
 
     public func prepare(input: UserInput) throws -> LMInput {
-        let prompt = input.prompt.asMessages().last?["content"] ?? ""
+        let prompt = input.prompt.asMessages().last?["content"] as? String ?? ""
 
         if input.images.isEmpty {
             // No image scenario

Libraries/MLXVLM/Models/Paligemma.swift

@@ -478,7 +478,7 @@ public class PaligGemmaProcessor: UserInputProcessor {
         }
 
         // this doesn't have a chat template so just use the last message.
-        var prompt = input.prompt.asMessages().last?["content"] ?? ""
+        var prompt = input.prompt.asMessages().last?["content"] as? String ?? ""
 
         // based on transformers/processing_paligemma
         let count = input.images.count * config.imageSequenceLength

Libraries/MLXVLM/Models/Qwen2VL.swift

@@ -686,69 +686,65 @@ public class Qwen2VLProcessor: UserInputProcessor {
         return (flattenedPatches, .init(gridT, gridH, gridW))
     }
 
-    public func prepare(prompt: UserInput.Prompt, imageTHW: [THW]?) -> String {
-        // the tokenizer does have a chat template and it expects messages
-        // like this:
-        //
-        // [{'role': 'user', 'content': [{'type': 'text', 'text': 'What are these?'},
-        //  {'type': 'image'}, {'type': 'image'}, {'type': 'image'}]}]
-        //
-        // The output of the prompt template is fed into
-        // image_processing_qwen2_vl.preprocess where it is further augmented
-        // by replacing tokens according to imageTHW.
-        //
-        // Neither the structured content nor the postprocessing of the template
-        // are supported in current Tokenizer/Jinja (swift) so handle that here.
-
-        var messages = prompt.asMessages()
-        if messages[0]["role"] != "system" {
+    private func prepareMessages(_ messages: [Message]) -> [Message] {
+        var messages = messages
+        print(messages)
+        // Add system message if not present
+        if let role = messages[0]["role"] as? String, role != "system" {
             messages.insert(["role": "system", "content": "You are a helpful assistant."], at: 0)
         }
 
-        let lastIndex = messages.count - 1
-        var lastMessage = messages[lastIndex]["content"] ?? ""
-
-        // image_processing_qwen2_vl.preprocess -- inject image_pad tokens for each image
-        let mergeLength = config.mergeSize * config.mergeSize
-        for thw in imageTHW ?? [] {
-            lastMessage += "<|vision_start|>"
-            lastMessage += Array(repeating: "<|image_pad|>", count: thw.product / mergeLength)
-                .joined()
-            lastMessage += "<|vision_end|>"
-        }
-
-        messages[lastIndex]["content"] = lastMessage
-
-        return
-            messages
-            .map {
-                "<|im_start|>\($0["role"] ?? "user")\n\($0["content"] ?? "")<|im_end|>"
-            }
-            .joined(separator: "\n")
-            + "\n<|im_start|>assistant\n"
+        return messages
     }
 
+    //    public func prepare(prompt: UserInput.Prompt, imageTHW: [THW]?) throws -> String {
+    //        let messages = prepareMessages(prompt.asMessages())
+    //        let tokens = try tokenizer.applyChatTemplate(messages: messages)
+    //        return tokenizer.decode(tokens: tokens)
+    //    }
+
     public func prepare(input: UserInput) throws -> LMInput {
+        // Text-only input
         if input.images.isEmpty {
-            // just a straight text prompt
-            let prompt = prepare(prompt: input.prompt, imageTHW: nil)
-            let promptTokens = try tokenizer.encode(text: prompt)
+            let messages = input.prompt.asMessages()
+            let promptTokens = try tokenizer.applyChatTemplate(messages: messages)
             return LMInput(tokens: MLXArray(promptTokens))
         }
 
-        // image_processing_qwen2_vl.preprocess
+        // Input with images
         let images = try input.images.map {
             try preprocess(images: [$0.asCIImage()], processing: input.processing)
         }
         let pixels = concatenated(images.map { $0.0 })
         let image = LMInput.ProcessedImage(pixels: pixels, imageGridThw: images.map { $0.1 })
 
-        // processing_qwen2_vl.Qwen2VLProcessor
-        let prompt = prepare(prompt: input.prompt, imageTHW: image.imageGridThw)
-        let promptTokens = try tokenizer.encode(text: prompt)
+        // Get tokens from messages
+        let messages = prepareMessages(input.prompt.asMessages())
+        var promptTokens = try tokenizer.applyChatTemplate(messages: messages)
+
+        // Replace single image pad token with correct number for each image
+        let imagePadToken = try tokenizer.encode(text: "<|image_pad|>").first!
+        let mergeLength = config.mergeSize * config.mergeSize
+
+        // TODO: This assumes that there is only one image. A better solution is needed for the case when multiple images are included.
+        if let imageGridThw = image.imageGridThw {
+            for thw in imageGridThw {
+                if let padIndex = promptTokens.firstIndex(of: imagePadToken) {
+                    let paddingCount = thw.product / mergeLength
+                    promptTokens.replaceSubrange(
+                        padIndex ... (padIndex),
+                        with: Array(repeating: imagePadToken, count: paddingCount)
+                    )
+                }
+            }
+        }
+
+        // TODO: For debugging. Remove later.
+        let promptTokensDecoded = try tokenizer.decode(tokens: promptTokens)
+        print(promptTokensDecoded)
+
         let promptArray = MLXArray(promptTokens).expandedDimensions(axis: 0)
         let mask = ones(like: promptArray).asType(.int8)
-
         return LMInput(text: .init(tokens: promptArray, mask: mask), image: image)
     }
 

Package.swift

@@ -28,7 +28,9 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/ml-explore/mlx-swift", from: "0.21.2"),
-        .package(url: "https://github.com/huggingface/swift-transformers", from: "0.1.13"),
+        //        .package(url: "https://github.com/huggingface/swift-transformers", from: "0.1.13"),
+        .package(
+            url: "https://github.com/DePasqualeOrg/swift-transformers", branch: "images-and-tools"),
         .package(url: "https://github.com/1024jp/GzipSwift", "6.0.1" ... "6.0.1"),
         .package(url: "https://github.com/apple/swift-async-algorithms", from: "1.0.0"),
     ],