@inferrlm/react-native-mlx 0.4.1 → 0.4.2-alpha.1

package/ios/Sources/HybridLLM.swift

@@ -66,6 +66,45 @@ class HybridLLM: HybridLLMSpec {
                      allocatedMB, cacheMB, peakMB)
     }
 
+    private func parseEosTokenIds(from modelDir: URL) -> Set<Int> {
+        let configURL = modelDir.appendingPathComponent("config.json")
+        guard let data = try? Data(contentsOf: configURL),
+              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
+        else { return [] }
+
+        var allIds = Set<Int>()
+
+        if let ids = extractEosIds(from: json) {
+            allIds.formUnion(ids)
+        }
+
+        for key in ["text_config", "language_config", "llm_config"] {
+            if let nested = json[key] as? [String: Any],
+               let ids = extractEosIds(from: nested) {
+                allIds.formUnion(ids)
+            }
+        }
+
+        let genConfigURL = modelDir.appendingPathComponent("generation_config.json")
+        if let genData = try? Data(contentsOf: genConfigURL),
+           let genJson = try? JSONSerialization.jsonObject(with: genData) as? [String: Any],
+           let ids = extractEosIds(from: genJson) {
+            allIds = ids
+        }
+
+        return allIds
+    }
+
+    private func extractEosIds(from dict: [String: Any]) -> Set<Int>? {
+        if let id = dict["eos_token_id"] as? Int {
+            return [id]
+        }
+        if let ids = dict["eos_token_id"] as? [Int], !ids.isEmpty {
+            return Set(ids)
+        }
+        return nil
+    }
+
     private func buildToolSchema(from tool: ToolDefinition) -> ToolSpec {
         var properties: [String: [String: Any]] = [:]
         var required: [String] = []
@@ -125,6 +164,33 @@ class HybridLLM: HybridLLMSpec {
 
                 try Task.checkCancellation()
 
+                /*
+                 mlx-swift-lm only reads top-level eos_token_id from config.json.
+                 Models like Qwen3.5 nest it inside text_config, leaving the stop
+                 set empty. Parse it ourselves and patch the container.
+                 Also add common chat stop tokens as extraEOSTokens.
+                */
+                let containerEos = await loadedContainer.configuration.eosTokenIds
+                let containerExtra = await loadedContainer.configuration.extraEOSTokens
+                log("EOS state after load - ids: \(containerEos), extra: \(containerExtra)")
+
+                let parsed = self.parseEosTokenIds(from: modelDir)
+                let chatStopTokens: Set<String> = ["<|endoftext|>", "<|im_end|>", "<|im_start|>"]
+                let needsIdPatch = containerEos.isEmpty && !parsed.isEmpty
+                let missingExtra = chatStopTokens.subtracting(containerExtra)
+
+                if needsIdPatch || !missingExtra.isEmpty {
+                    await loadedContainer.update { ctx in
+                        if needsIdPatch {
+                            ctx.configuration.eosTokenIds = parsed
+                        }
+                        ctx.configuration.extraEOSTokens.formUnion(chatStopTokens)
+                    }
+                    let updated = await loadedContainer.configuration.eosTokenIds
+                    let updatedExtra = await loadedContainer.configuration.extraEOSTokens
+                    log("EOS patched - ids: \(updated), extra: \(updatedExtra)")
+                }
+
                 let memoryAfterContainer = self.getMemoryUsage()
                 let gpuAfterContainer = self.getGPUMemoryUsage()
                 log("Model loaded - Host: \(memoryAfterContainer), GPU: \(gpuAfterContainer)")

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@inferrlm/react-native-mlx",
   "description": "MLX Swift integration for React Native - InferrLM fork with enhanced features",
-  "version": "0.4.1",
+  "version": "0.4.2-alpha.1",
   "main": "./lib/module/index.js",
   "module": "./lib/module/index.js",
   "types": "./lib/typescript/src/index.d.ts",