Use chat template for Qwen 2 VL

Author: DePasqualeOrg

Libraries/MLXVLM/Models/Qwen2VL.swift

@@ -686,69 +686,51 @@ 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()
+    private func prepareMessages(_ messages: [[String: String]], imageTHW: [THW]?) -> [[String: String]] {
+        var messages = messages
+        // Add system message if not present
         if messages[0]["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|>"
+        // Add image markers to last message if needed
+        if let imageTHW {
+            let lastIndex = messages.count - 1
+            var content = messages[lastIndex]["content"] ?? ""
+            let mergeLength = config.mergeSize * config.mergeSize
+            for thw in imageTHW {
+                content += "<|vision_start|>"
+                content += Array(repeating: "<|image_pad|>", count: thw.product / mergeLength).joined()
+                content += "<|vision_end|>"
+            }
+            messages[lastIndex]["content"] = content
         }
+        return messages
+    }
 
-        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"
+    public func prepare(prompt: UserInput.Prompt, imageTHW: [THW]?) throws -> String {
+        let messages = prepareMessages(prompt.asMessages(), imageTHW: imageTHW)
+        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)
+        // Prepare messages with image markers
+        let messages = prepareMessages(input.prompt.asMessages(), imageTHW: image.imageGridThw)
+        let promptTokens = try tokenizer.applyChatTemplate(messages: messages)
         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"),
     ],