@inferrlm/react-native-mlx 0.4.2 → 0.4.8

package/MLXReactNative.podspec

@@ -24,13 +24,19 @@ Pod::Spec.new do |s|
 
   spm_dependency(s,
     url: "https://github.com/ml-explore/mlx-swift-lm.git",
-    requirement: {kind: "branch", branch: "main"},
+    requirement: {kind: "upToNextMinorVersion", minimumVersion: "3.31.3"},
     products: ["MLXLLM", "MLXLMCommon"]
   )
 
+  spm_dependency(s,
+    url: "https://github.com/huggingface/swift-transformers",
+    requirement: {kind: "upToNextMinorVersion", minimumVersion: "1.2.0"},
+    products: ["Tokenizers"]
+  )
+
   spm_dependency(s,
     url: "https://github.com/Blaizzy/mlx-audio-swift.git",
-    requirement: {kind: "branch", branch: "main"},
+    requirement: {kind: "revision", revision: "856e04afb3c6eb931d92bb0d6ae7bbfbdfa89b15"},
     products: ["MLXAudioTTS", "MLXAudioSTT", "MLXAudioCore"]
   )
 

package/app.plugin.js

@@ -0,0 +1,63 @@
+const { withPodfile } = require('@expo/config-plugins');
+
+function disableDeterministicPodUuids(contents) {
+  const line = "install! 'cocoapods', :deterministic_uuids => false";
+
+  if (contents.includes(line)) {
+    return contents;
+  }
+
+  const anchor = 'prepare_react_native_project!';
+
+  if (!contents.includes(anchor)) {
+    return contents;
+  }
+
+  return contents.replace(anchor, `${line}\n\n${anchor}`);
+}
+
+function injectSpmRootFix(contents) {
+  const block = [
+    'if defined?(::SPMManager) && ::SPMManager.instance_methods.include?(:add_spm_to_target)',
+    '  ::SPMManager.class_eval do',
+    '    unless method_defined?(:inferrlm_add_spm_to_target)',
+    '      alias_method :inferrlm_add_spm_to_target, :add_spm_to_target',
+    '',
+    '      def add_spm_to_target(project, target, url, requirement, products)',
+    '        root = project.root_object',
+    '        if root && project.objects_by_uuid[root.uuid] != root',
+    '          root.add_referrer(project)',
+    '        end',
+    '',
+    '        inferrlm_add_spm_to_target(project, target, url, requirement, products)',
+    '      end',
+    '    end',
+    '  end',
+    'end',
+  ].join('\n');
+
+  if (contents.includes('::SPMManager.class_eval do')) {
+    return contents;
+  }
+
+  const anchor = 'prepare_react_native_project!';
+
+  if (!contents.includes(anchor)) {
+    return contents;
+  }
+
+  return contents.replace(anchor, `${block}\n\n${anchor}`);
+}
+
+module.exports = function withMlxIosPods(config) {
+  return withPodfile(config, (modConfig) => {
+    modConfig.modResults.contents = disableDeterministicPodUuids(modConfig.modResults.contents);
+    modConfig.modResults.contents = injectSpmRootFix(modConfig.modResults.contents);
+    return modConfig;
+  });
+};
+
+module.exports._helpers = {
+  disableDeterministicPodUuids,
+  injectSpmRootFix,
+};

package/ios/Sources/HybridLLM.swift

@@ -1,9 +1,149 @@
 import Foundation
 import NitroModules
+import Tokenizers
 internal import MLX
 internal import MLXLLM
 internal import MLXLMCommon
-internal import Tokenizers
+
+private typealias ToolSpec = [String: Any]
+
+private func isJANGModel(at dir: URL) -> Bool {
+    FileManager.default.fileExists(atPath: dir.appendingPathComponent("jang_config.json").path)
+}
+
+private protocol ChatTemplateConfigurableTokenizer: MLXLMCommon.Tokenizer {
+    func applyChatTemplate(
+        messages: [[String: any Sendable]],
+        chatTemplate: Tokenizers.ChatTemplateArgument?,
+        addGenerationPrompt: Bool,
+        truncation: Bool,
+        maxLength: Int?,
+        tools: [[String: any Sendable]]?,
+        additionalContext: [String: any Sendable]?
+    ) throws -> [Int]
+}
+
+private extension MLXLMCommon.Tokenizer {
+    func applyChatTemplate(
+        messages: [[String: any Sendable]],
+        chatTemplate: Tokenizers.ChatTemplateArgument?,
+        addGenerationPrompt: Bool,
+        truncation: Bool,
+        maxLength: Int?,
+        tools: [[String: any Sendable]]?,
+        additionalContext: [String: any Sendable]?
+    ) throws -> [Int] {
+        guard let tokenizer = self as? any ChatTemplateConfigurableTokenizer else {
+            throw MLXLMCommon.TokenizerError.missingChatTemplate
+        }
+
+        return try tokenizer.applyChatTemplate(
+            messages: messages,
+            chatTemplate: chatTemplate,
+            addGenerationPrompt: addGenerationPrompt,
+            truncation: truncation,
+            maxLength: maxLength,
+            tools: tools,
+            additionalContext: additionalContext
+        )
+    }
+}
+
+private struct TransformersTokenizerBridge: ChatTemplateConfigurableTokenizer {
+    private let upstream: any Tokenizers.Tokenizer
+
+    init(_ upstream: any Tokenizers.Tokenizer) {
+        self.upstream = upstream
+    }
+
+    func encode(text: String, addSpecialTokens: Bool) -> [Int] {
+        upstream.encode(text: text, addSpecialTokens: addSpecialTokens)
+    }
+
+    func decode(tokenIds: [Int], skipSpecialTokens: Bool) -> String {
+        upstream.decode(tokens: tokenIds, skipSpecialTokens: skipSpecialTokens)
+    }
+
+    func convertTokenToId(_ token: String) -> Int? {
+        upstream.convertTokenToId(token)
+    }
+
+    func convertIdToToken(_ id: Int) -> String? {
+        upstream.convertIdToToken(id)
+    }
+
+    var bosToken: String? { upstream.bosToken }
+    var eosToken: String? { upstream.eosToken }
+    var unknownToken: String? { upstream.unknownToken }
+
+    func applyChatTemplate(
+        messages: [[String: any Sendable]],
+        tools: [[String: any Sendable]]?,
+        additionalContext: [String: any Sendable]?
+    ) throws -> [Int] {
+        do {
+            return try upstream.applyChatTemplate(
+                messages: messages,
+                tools: tools,
+                additionalContext: additionalContext
+            )
+        } catch Tokenizers.TokenizerError.missingChatTemplate {
+            throw MLXLMCommon.TokenizerError.missingChatTemplate
+        }
+    }
+
+    func applyChatTemplate(
+        messages: [[String: any Sendable]],
+        chatTemplate: Tokenizers.ChatTemplateArgument?,
+        addGenerationPrompt: Bool,
+        truncation: Bool,
+        maxLength: Int?,
+        tools: [[String: any Sendable]]?,
+        additionalContext: [String: any Sendable]?
+    ) throws -> [Int] {
+        do {
+            return try upstream.applyChatTemplate(
+                messages: messages,
+                chatTemplate: chatTemplate,
+                addGenerationPrompt: addGenerationPrompt,
+                truncation: truncation,
+                maxLength: maxLength,
+                tools: tools,
+                additionalContext: additionalContext
+            )
+        } catch Tokenizers.TokenizerError.missingChatTemplate {
+            throw MLXLMCommon.TokenizerError.missingChatTemplate
+        }
+    }
+}
+
+private struct TransformersTokenizerLoader: MLXLMCommon.TokenizerLoader {
+    func load(from directory: URL) async throws -> any MLXLMCommon.Tokenizer {
+        let upstream = try await Tokenizers.AutoTokenizer.from(modelFolder: directory)
+        return TransformersTokenizerBridge(upstream)
+    }
+}
+
+@MainActor
+private final class LoadProgress {
+    private let callback: ((Double) -> Void)?
+    private(set) var value: Double = 0
+
+    init(callback: ((Double) -> Void)?) {
+        self.callback = callback
+    }
+
+    func set(_ nextValue: Double) {
+        let clamped = min(1.0, max(value, nextValue))
+        guard clamped > value else { return }
+        value = clamped
+        callback?(clamped)
+    }
+
+    func tick(toward upperBound: Double, step: Double) {
+        set(min(upperBound, value + step))
+    }
+}
 
 class HybridLLM: HybridLLMSpec {
     private var session: ChatSession?
@@ -16,7 +156,6 @@ class HybridLLM: HybridLLMSpec {
         totalTime: 0,
         toolExecutionTime: 0
     )
-    private var modelFactory: ModelFactory = LLMModelFactory.shared
     private var manageHistory: Bool = false
     private var messageHistory: [LLMMessage] = []
     private var loadTask: Task<Void, Error>?
@@ -142,6 +281,9 @@ class HybridLLM: HybridLLMSpec {
 
         return Promise.async { [self] in
             let task = Task { @MainActor in
+                let progress = LoadProgress(callback: options?.onProgress)
+                progress.set(0.02)
+
                 Memory.cacheLimit = 2000000
 
                 self.currentTask?.cancel()
@@ -151,6 +293,7 @@ class HybridLLM: HybridLLMSpec {
                 self.tools = []
                 self.toolSchemas = []
                 Memory.clearCache()
+                progress.set(0.12)
 
                 let memoryAfterCleanup = self.getMemoryUsage()
                 let gpuAfterCleanup = self.getGPUMemoryUsage()
@@ -158,13 +301,30 @@ class HybridLLM: HybridLLMSpec {
 
                 let modelDir = await ModelDownloader.shared.getModelDirectory(modelId: modelId)
                 log("Loading from directory: \(modelDir.path)")
+                progress.set(0.18)
 
-                let config = ModelConfiguration(directory: modelDir)
-                let loadedContainer = try await self.modelFactory.loadContainer(
-                    configuration: config
-                ) { progress in
-                    options?.onProgress?(progress.fractionCompleted)
+                if isJANGModel(at: modelDir) {
+                    throw LLMError.unsupportedModel("JANG model format is not supported by the current MLX dependency set")
+                }
+
+                var loadingPulse: Task<Void, Never>?
+                if options?.onProgress != nil {
+                    loadingPulse = Task { @MainActor in
+                        while !Task.isCancelled {
+                            try? await Task.sleep(nanoseconds: 250_000_000)
+                            progress.tick(toward: 0.78, step: 0.04)
+                        }
+                    }
                 }
+                defer { loadingPulse?.cancel() }
+
+                progress.set(0.22)
+
+                let loadedContainer = try await LLMModelFactory.shared.loadContainer(
+                    from: modelDir,
+                    using: TransformersTokenizerLoader()
+                )
+                progress.set(0.86)
 
                 try Task.checkCancellation()
 
@@ -194,6 +354,9 @@ class HybridLLM: HybridLLMSpec {
                     let updatedExtra = await loadedContainer.configuration.extraEOSTokens
                     log("EOS patched - ids: \(updated), extra: \(updatedExtra)")
                 }
+                progress.set(0.92)
+
+                progress.set(0.95)
 
                 let memoryAfterContainer = self.getMemoryUsage()
                 let gpuAfterContainer = self.getGPUMemoryUsage()
@@ -205,10 +368,14 @@ class HybridLLM: HybridLLMSpec {
                     log("Loaded \(self.tools.count) tools: \(self.tools.map { $0.name })")
                 }
 
-                let additionalContextDict: [String: Any]? = if let messages = options?.additionalContext {
-                    ["messages": messages.map { ["role": $0.role, "content": $0.content] }]
+                let additionalContextDict: [String: any Sendable]?
+                if let messages = options?.additionalContext {
+                    let contextMessages: [[String: String]] = messages.map {
+                        ["role": $0.role, "content": $0.content]
+                    }
+                    additionalContextDict = ["messages": contextMessages]
                 } else {
-                    nil
+                    additionalContextDict = nil
                 }
 
                 self.container = loadedContainer
@@ -221,6 +388,8 @@ class HybridLLM: HybridLLMSpec {
                 if self.manageHistory {
                     log("History management enabled with \(self.messageHistory.count) initial messages")
                 }
+
+                progress.set(1.0)
             }
 
             self.loadTask = task
@@ -471,7 +640,7 @@ class HybridLLM: HybridLLMSpec {
                     additionalContext: additionalCtx
                 )
                 self.log("template_applied token_count=\(result.count)")
-                let decoded = context.tokenizer.decode(tokens: Array(result.suffix(60)))
+                let decoded = context.tokenizer.decode(tokenIds: Array(result.suffix(60)))
                 self.log("input_tail_decoded: \(decoded)")
                 self.lastInputContainedThinkTag = decoded.contains("<think>")
                 return result
@@ -487,7 +656,7 @@ class HybridLLM: HybridLLMSpec {
                     additionalContext: additionalCtx
                 )
                 self.log("fallback_template_applied token_count=\(result.count)")
-                let decoded = context.tokenizer.decode(tokens: Array(result.suffix(60)))
+                let decoded = context.tokenizer.decode(tokenIds: Array(result.suffix(60)))
                 self.log("fallback_input_tail_decoded: \(decoded)")
                 self.lastInputContainedThinkTag = decoded.contains("<think>")
                 return result

package/ios/Sources/LLMError.swift

@@ -3,4 +3,5 @@ import Foundation
 public enum LLMError: Error {
     case notLoaded
     case generationFailed(String)
+    case unsupportedModel(String)
 }

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.2",
+  "version": "0.4.8",
   "main": "./lib/module/index.js",
   "module": "./lib/module/index.js",
   "types": "./lib/typescript/src/index.d.ts",