Merge pull request #104 from roydsouza/fix/moe-memory-and-multimodal-tokens-rebased

Fix: Resolve multimodal BOA/EOA tokens dynamically from config.json

Author: solderzzc

Sources/SwiftLM/Server.swift

@@ -3013,15 +3013,15 @@ public final class ALMModelFactory: ModelFactory, @unchecked Sendable {
     ) async throws -> ModelContext {
         let context = try await LLMModelFactory.shared._load(configuration: configuration, tokenizerLoader: tokenizerLoader)
 
-        let numAudioEmbeddings = OmniModelFactory.extractNumAudioEmbeddings(configuration: configuration)
+        let tokens = OmniModelFactory.extractMultimodalTokens(configuration: configuration)
         let messageGenerator = DefaultMessageGenerator()
         let processor = ALMUserInputProcessor(
             tokenizer: context.tokenizer,
             configuration: context.configuration,
             messageGenerator: messageGenerator,
-            boaToken: 255010,
-            eoaToken: 255011,
-            numAudioEmbeddings: numAudioEmbeddings
+            boaToken: tokens.boa,
+            eoaToken: tokens.eoa,
+            numAudioEmbeddings: tokens.numAudio
         )
 
         return .init(
@@ -3081,10 +3081,12 @@ public final class OmniModelFactory: ModelFactory, @unchecked Sendable {
         tokenizerLoader: any TokenizerLoader
     ) async throws -> ModelContext {
         let vlmContext = try await VLMModelFactory.shared._load(configuration: configuration, tokenizerLoader: tokenizerLoader)
-        let numAudioEmbeddings = OmniModelFactory.extractNumAudioEmbeddings(configuration: configuration)
+        let tokens = OmniModelFactory.extractMultimodalTokens(configuration: configuration)
         let omniProcessor = OmniUserInputProcessor(
             vlmProcessor: vlmContext.processor,
-            numAudioEmbeddings: numAudioEmbeddings
+            boaToken: tokens.boa,
+            eoaToken: tokens.eoa,
+            numAudioEmbeddings: tokens.numAudio
         )
 
         return .init(
@@ -3095,19 +3097,35 @@ public final class OmniModelFactory: ModelFactory, @unchecked Sendable {
         )
     }
 
+    @available(*, deprecated, message: "Use extractMultimodalTokens(configuration:).numAudio instead")
     public static func extractNumAudioEmbeddings(configuration: ResolvedModelConfiguration) -> Int {
+        extractMultimodalTokens(configuration: configuration).numAudio
+    }
+
+    public static func extractMultimodalTokens(configuration: ResolvedModelConfiguration) -> (numAudio: Int, boa: Int, eoa: Int) {
         let configurationURL = configuration.modelDirectory.appending(component: "config.json")
+        var numAudio = 128
+        var boa = 255010
+        var eoa = 255011
+        
         if let data = try? Data(contentsOf: configurationURL),
            let dict = try? JSONSerialization.jsonObject(with: data) as? [String: Any] {
 
+            // Extract num_audio_embeddings
             if let subsampling = dict["subsampling_conv_channels"] as? [Int] {
-                return subsampling.first ?? 128
-            }
-            if let audioConfig = dict["audio_config"] as? [String: Any],
+                numAudio = subsampling.first ?? 128
+            } else if let audioConfig = dict["audio_config"] as? [String: Any],
                let embeddings = audioConfig["num_audio_embeddings"] as? Int {
-                return embeddings
+                numAudio = embeddings
             }
+            
+            // Extract BOA/EOA tokens
+            if let b = dict["boa_token_id"] as? Int { boa = b }
+            else if let b = (dict["audio_config"] as? [String: Any])?["boa_token_id"] as? Int { boa = b }
+            
+            if let e = dict["eoa_token_id"] as? Int { eoa = e }
+            else if let e = (dict["audio_config"] as? [String: Any])?["eoa_token_id"] as? Int { eoa = e }
         }
-        return 128
+        return (numAudio, boa, eoa)
     }
 }

tests/SwiftLMTests/MultimodalTokenExtractionTests.swift

@@ -0,0 +1,68 @@
+import XCTest
+import Foundation
+@testable import SwiftLM
+import MLXLMCommon
+
+final class MultimodalTokenExtractionTests: XCTestCase {
+
+    func testExtractMultimodalTokens_Defaults() throws {
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: tempDir) }
+        
+        let config = ModelConfiguration(directory: tempDir).resolved(modelDirectory: tempDir, tokenizerDirectory: tempDir)
+        
+        let tokens = OmniModelFactory.extractMultimodalTokens(configuration: config)
+        XCTAssertEqual(tokens.numAudio, 128)
+        XCTAssertEqual(tokens.boa, 255010)
+        XCTAssertEqual(tokens.eoa, 255011)
+    }
+
+    func testExtractMultimodalTokens_FromConfig() throws {
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: tempDir) }
+        
+        let jsonDict: [String: Any] = [
+            "subsampling_conv_channels": [256],
+            "boa_token_id": 999990,
+            "eoa_token_id": 999991
+        ]
+        
+        let jsonData = try JSONSerialization.data(withJSONObject: jsonDict)
+        let configURL = tempDir.appendingPathComponent("config.json")
+        try jsonData.write(to: configURL)
+        
+        let config = ModelConfiguration(directory: tempDir).resolved(modelDirectory: tempDir, tokenizerDirectory: tempDir)
+        let tokens = OmniModelFactory.extractMultimodalTokens(configuration: config)
+        
+        XCTAssertEqual(tokens.numAudio, 256)
+        XCTAssertEqual(tokens.boa, 999990)
+        XCTAssertEqual(tokens.eoa, 999991)
+    }
+
+    func testExtractMultimodalTokens_FromAudioConfigFallback() throws {
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: tempDir) }
+        
+        let jsonDict: [String: Any] = [
+            "audio_config": [
+                "num_audio_embeddings": 512,
+                "boa_token_id": 888880,
+                "eoa_token_id": 888881
+            ]
+        ]
+        
+        let jsonData = try JSONSerialization.data(withJSONObject: jsonDict)
+        let configURL = tempDir.appendingPathComponent("config.json")
+        try jsonData.write(to: configURL)
+        
+        let config = ModelConfiguration(directory: tempDir).resolved(modelDirectory: tempDir, tokenizerDirectory: tempDir)
+        let tokens = OmniModelFactory.extractMultimodalTokens(configuration: config)
+        
+        XCTAssertEqual(tokens.numAudio, 512)
+        XCTAssertEqual(tokens.boa, 888880)
+        XCTAssertEqual(tokens.eoa, 888881)
+    }
+}