bleam 0.0.10 → 0.0.12

package/templates/native/ios/Bleam/AI/TextGeneration/QwenHybrid/QwenHybridLinearAttention.swift

@@ -0,0 +1,111 @@
+import Foundation
+import MLX
+import MLXNN
+
+public final class QwenHybridLinearAttention: Module, @unchecked Sendable {
+    private let numValueHeads: Int
+    private let numKeyHeads: Int
+    private let keyHeadDim: Int
+    private let valueHeadDim: Int
+    private let keyDimension: Int
+    private let valueDimension: Int
+    private let convolutionDimension: Int
+    private let kernelSize: Int
+    private let keyNormWeight: MLXArray
+
+    @ModuleInfo public var in_proj_qkv: Linear
+    @ModuleInfo public var in_proj_z: Linear
+    @ModuleInfo public var in_proj_b: Linear
+    @ModuleInfo public var in_proj_a: Linear
+    @ModuleInfo public var out_proj: Linear
+    @ModuleInfo public var norm: QwenHybridGatedRMSNorm
+    @ModuleInfo public var conv1d: QwenHybridDepthwiseConv1D
+    @ParameterInfo public var dt_bias: MLXArray
+    @ParameterInfo public var A_log: MLXArray
+
+    public init(config: any QwenHybridConfiguration) {
+        numValueHeads = config.linearNumValueHeads
+        numKeyHeads = config.linearNumKeyHeads
+        keyHeadDim = config.linearKeyHeadDim
+        valueHeadDim = config.linearValueHeadDim
+        keyDimension = config.linearNumKeyHeads * config.linearKeyHeadDim
+        valueDimension = config.linearNumValueHeads * config.linearValueHeadDim
+        convolutionDimension = keyDimension * 2 + valueDimension
+        kernelSize = config.linearConvKernelDim
+        keyNormWeight = MLX.ones([config.linearKeyHeadDim])
+        _in_proj_qkv = ModuleInfo(wrappedValue: Linear(config.hiddenSize, convolutionDimension, bias: false))
+        _in_proj_z = ModuleInfo(wrappedValue: Linear(config.hiddenSize, valueDimension, bias: false))
+        _in_proj_b = ModuleInfo(wrappedValue: Linear(config.hiddenSize, numValueHeads, bias: false))
+        _in_proj_a = ModuleInfo(wrappedValue: Linear(config.hiddenSize, numValueHeads, bias: false))
+        _out_proj = ModuleInfo(wrappedValue: Linear(valueDimension, config.hiddenSize, bias: false))
+        _norm = ModuleInfo(wrappedValue: QwenHybridGatedRMSNorm(dimensions: valueHeadDim, epsilon: config.rmsNormEps))
+        _conv1d = ModuleInfo(wrappedValue: QwenHybridDepthwiseConv1D(channels: convolutionDimension, kernelSize: kernelSize))
+        _dt_bias = ParameterInfo(wrappedValue: MLX.ones([numValueHeads]))
+        _A_log = ParameterInfo(wrappedValue: MLX.zeros([numValueHeads]))
+        super.init()
+    }
+
+    public func callAsFunction(_ input: MLXArray, cache: QwenHybridLayerCache) -> MLXArray {
+        let batch = input.dim(0)
+        let length = input.dim(1)
+        let mixed = causalConvolution(in_proj_qkv(input), cache: cache)
+        let pieces = mixed.split(indices: [keyDimension, keyDimension * 2], axis: -1)
+        var query = pieces[0].reshaped([batch, length, numKeyHeads, keyHeadDim])
+        var key = pieces[1].reshaped([batch, length, numKeyHeads, keyHeadDim])
+        let value = pieces[2].reshaped([batch, length, numValueHeads, valueHeadDim]).asType(.float32)
+        query = MLXFast.rmsNorm(query.asType(.float32), weight: keyNormWeight, eps: 1e-6) / Float(keyHeadDim)
+        key = MLXFast.rmsNorm(key.asType(.float32), weight: keyNormWeight, eps: 1e-6) / sqrt(Float(keyHeadDim))
+        let gate = MLX.exp(-MLX.exp(A_log.asType(.float32)) * MLXNN.softplus(in_proj_a(input).asType(.float32) + dt_bias))
+        let beta = MLX.sigmoid(in_proj_b(input)).asType(.float32)
+        let initialState = cache.recurrentState ?? MLX.zeros([batch, numValueHeads, valueHeadDim, keyHeadDim], dtype: .float32)
+        let updated = QwenHybridGatedDelta.update(query: query, key: key, value: value, gate: gate, beta: beta, state: initialState)
+        cache.recurrentState = updated.state
+        let normalized = norm(
+            updated.output.asType(input.dtype).reshaped([-1, valueHeadDim]),
+            gate: in_proj_z(input).reshaped([-1, valueHeadDim])
+        )
+        return out_proj(normalized.reshaped([batch, length, valueDimension]))
+    }
+
+    private func causalConvolution(_ input: MLXArray, cache: QwenHybridLayerCache) -> MLXArray {
+        let length = input.dim(1)
+        if let previous = cache.convolutionState {
+            let combined = MLX.concatenated([previous, input], axis: 1)
+            cache.convolutionState = combined[0..., (combined.dim(1) - kernelSize + 1)..., 0...]
+            return MLXNN.silu(MLX.conv1d(combined, conv1d.weight, padding: 0, groups: convolutionDimension))
+        }
+        let result = MLXNN.silu(MLX.conv1d(input, conv1d.weight, padding: kernelSize - 1, groups: convolutionDimension))
+        let paddedInput = padded(input, widths: [[0, 0], [max(0, kernelSize - 1 - length), 0], [0, 0]])
+        cache.convolutionState = paddedInput[0..., (paddedInput.dim(1) - kernelSize + 1)..., 0...]
+        return result[0..., 0..<length, 0...]
+    }
+}
+
+public enum QwenHybridGatedDelta {
+    public static func update(
+        query: MLXArray,
+        key: MLXArray,
+        value: MLXArray,
+        gate: MLXArray,
+        beta: MLXArray,
+        state initialState: MLXArray
+    ) -> (output: MLXArray, state: MLXArray) {
+        let repeats = value.dim(2) / key.dim(2)
+        let query = repeats == 1 ? query : MLX.repeated(query, count: repeats, axis: 2)
+        let key = repeats == 1 ? key : MLX.repeated(key, count: repeats, axis: 2)
+        var state = initialState
+        var outputs: [MLXArray] = []
+        outputs.reserveCapacity(query.dim(1))
+        for index in 0..<query.dim(1) {
+            let q = query[0..., index, 0..., 0...]
+            let k = key[0..., index, 0..., 0...]
+            let v = value[0..., index, 0..., 0...]
+            state = state * gate[0..., index, 0...].expandedDimensions(axes: [2, 3])
+            let remembered = (state * k.expandedDimensions(axis: 2)).sum(axis: -1)
+            let delta = (v - remembered) * beta[0..., index, 0...].expandedDimensions(axis: -1)
+            state = state + k.expandedDimensions(axis: 2) * delta.expandedDimensions(axis: -1)
+            outputs.append((state * q.expandedDimensions(axis: 2)).sum(axis: -1))
+        }
+        return (MLX.stacked(outputs, axis: 1), state)
+    }
+}

package/templates/native/ios/Bleam/AppDelegate.swift

@@ -4,7 +4,6 @@ import ReactAppDependencyProvider
 import UIKit
 import CryptoKit
 import Darwin
-import Security
 
 func launchArgument(_ name: String) -> String? {
   let arguments = ProcessInfo.processInfo.arguments
@@ -22,17 +21,6 @@ func launchArgument(_ name: String) -> String? {
 
 @main
 class AppDelegate: ExpoAppDelegate {
-  public override func application(
-    _ application: UIApplication,
-    didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey: Any]? = nil
-  ) -> Bool {
-    let launched = super.application(application, didFinishLaunchingWithOptions: launchOptions)
-    if launched {
-      PlatformHelperInstaller.refreshIfInstalledApp()
-    }
-    return launched
-  }
-
   public func application(
     _ application: UIApplication,
     configurationForConnecting connectingSceneSession: UISceneSession,
@@ -63,200 +51,6 @@ class AppDelegate: ExpoAppDelegate {
   }
 }
 
-private enum PlatformHelperInstaller {
-  static func refreshIfInstalledApp() {
-#if targetEnvironment(macCatalyst)
-    let homeDirectory = URL(fileURLWithPath: NSHomeDirectory(), isDirectory: true)
-    let homeApplications = homeDirectory
-      .appendingPathComponent("Applications", isDirectory: true)
-      .standardizedFileURL.path + "/"
-    guard Bundle.main.bundleURL.standardizedFileURL.path.hasPrefix(homeApplications) else {
-      return
-    }
-    guard let bundleIdentifier = Bundle.main.bundleIdentifier else { return }
-
-    let serviceName = "PlatformHelper"
-    let bundledHelper = Bundle.main.bundleURL
-      .appendingPathComponent("Contents/Helpers/\(serviceName)")
-    guard FileManager.default.isExecutableFile(atPath: bundledHelper.path) else {
-      NSLog("Bleam helper unavailable: bundled executable is missing")
-      return
-    }
-    guard
-      let teamIdentifier = Bundle.main.object(forInfoDictionaryKey: "BleamTeamIdentifier") as? String,
-      let platformVersion = Bundle.main.object(forInfoDictionaryKey: "BleamRuntimeVersion") as? String,
-      validSignature(at: bundledHelper, teamIdentifier: teamIdentifier)
-    else {
-      NSLog("Bleam helper unavailable: bundled executable has an invalid identity")
-      return
-    }
-
-    let label = "\(bundleIdentifier).PlatformHelper"
-    let domain = "gui/\(getuid())"
-    let helperRoot = homeDirectory
-      .appendingPathComponent("Library/Application Support/Bleam/Helpers", isDirectory: true)
-      .appendingPathComponent(bundleIdentifier, isDirectory: true)
-    let installedHelper = helperRoot.appendingPathComponent(serviceName)
-    let stagedHelper = helperRoot.appendingPathComponent(".\(serviceName).\(UUID().uuidString)")
-    let backupHelper = helperRoot.appendingPathComponent(".\(serviceName).previous.\(UUID().uuidString)")
-    let failedHelper = helperRoot.appendingPathComponent(".\(serviceName).failed.\(UUID().uuidString)")
-    let pendingApply = homeDirectory
-      .appendingPathComponent("Library/Application Support/Bleam/Updates", isDirectory: true)
-      .appendingPathComponent(bundleIdentifier, isDirectory: true)
-      .appendingPathComponent("platform-apply.pending")
-    let launchAgents = homeDirectory
-      .appendingPathComponent("Library/LaunchAgents", isDirectory: true)
-    let plistURL = launchAgents.appendingPathComponent("\(label).plist")
-    let propertyList: [String: Any] = [
-      "Label": label,
-      "ProgramArguments": [installedHelper.path],
-      "ProcessType": "Background",
-      "MachServices": [label: true],
-      "KeepAlive": [
-        "SuccessfulExit": false,
-        "PathState": [pendingApply.path: true],
-      ],
-      "EnvironmentVariables": [
-        "XPC_SERVICE_NAME": label,
-        "BLEAM_BUNDLE_IDENTIFIER": bundleIdentifier,
-        "BLEAM_TEAM_IDENTIFIER": teamIdentifier,
-        "BLEAM_HELPER_VERSION": platformVersion,
-        "BLEAM_PLATFORM_VERSION": platformVersion,
-      ],
-    ]
-
-    do {
-      defer { try? FileManager.default.removeItem(at: stagedHelper) }
-      try FileManager.default.createDirectory(
-        at: helperRoot,
-        withIntermediateDirectories: true
-      )
-      try FileManager.default.copyItem(at: bundledHelper, to: stagedHelper)
-      guard
-        FileManager.default.isExecutableFile(atPath: stagedHelper.path),
-        validSignature(at: stagedHelper, teamIdentifier: teamIdentifier)
-      else {
-        throw HelperInstallError.invalidSignature
-      }
-      try FileManager.default.createDirectory(
-        at: launchAgents,
-        withIntermediateDirectories: true
-      )
-      let data = try PropertyListSerialization.data(
-        fromPropertyList: propertyList,
-        format: .xml,
-        options: 0
-      )
-      let previousPlist = try? Data(contentsOf: plistURL)
-      let wasRegistered = runLaunchctl(["print", "\(domain)/\(label)"]) == 0
-      if wasRegistered, runLaunchctl(["bootout", "\(domain)/\(label)"]) != 0 {
-        throw HelperInstallError.registrationFailed
-      }
-
-      do {
-        if FileManager.default.fileExists(atPath: installedHelper.path),
-          rename(installedHelper.path, backupHelper.path) != 0
-        {
-          throw HelperInstallError.replacementFailed
-        }
-        guard rename(stagedHelper.path, installedHelper.path) == 0 else {
-          throw HelperInstallError.replacementFailed
-        }
-        try data.write(to: plistURL, options: .atomic)
-        let bootstrapStatus = runLaunchctl(["bootstrap", domain, plistURL.path])
-        guard
-          bootstrapStatus == 0 || runLaunchctl(["print", "\(domain)/\(label)"]) == 0
-        else {
-          throw HelperInstallError.registrationFailed
-        }
-        try? FileManager.default.removeItem(at: backupHelper)
-        try? FileManager.default.removeItem(
-          at: helperRoot.appendingPathComponent("PlatformHelper.xpc", isDirectory: true)
-        )
-      } catch let installError {
-        do {
-          let newJobRegistered = runLaunchctl(["print", "\(domain)/\(label)"]) == 0
-          if newJobRegistered,
-            runLaunchctl(["bootout", "\(domain)/\(label)"]) != 0
-          {
-            throw HelperInstallError.rollbackFailed
-          }
-          if FileManager.default.fileExists(atPath: installedHelper.path),
-            rename(installedHelper.path, failedHelper.path) != 0
-          {
-            throw HelperInstallError.rollbackFailed
-          }
-          if FileManager.default.fileExists(atPath: backupHelper.path),
-            rename(backupHelper.path, installedHelper.path) != 0
-          {
-            throw HelperInstallError.rollbackFailed
-          }
-          if let previousPlist {
-            try previousPlist.write(to: plistURL, options: .atomic)
-            if wasRegistered,
-              runLaunchctl(["bootstrap", domain, plistURL.path]) != 0
-            {
-              throw HelperInstallError.rollbackFailed
-            }
-          } else {
-            try? FileManager.default.removeItem(at: plistURL)
-          }
-          try? FileManager.default.removeItem(at: failedHelper)
-        } catch {
-          NSLog("Bleam helper rollback failed: %@", error.localizedDescription)
-          throw error
-        }
-        throw installError
-      }
-    } catch {
-      NSLog("Bleam helper install failed: %@", error.localizedDescription)
-      return
-    }
-#endif
-  }
-
-  private static func validSignature(at url: URL, teamIdentifier: String) -> Bool {
-    var staticCode: SecStaticCode?
-    guard
-      SecStaticCodeCreateWithPath(url as CFURL, [], &staticCode) == errSecSuccess,
-      let staticCode,
-      SecStaticCodeCheckValidity(
-        staticCode,
-        SecCSFlags(rawValue: kSecCSStrictValidate),
-        nil
-      ) == errSecSuccess
-    else { return false }
-
-    var signingInfo: CFDictionary?
-    guard
-      SecCodeCopySigningInformation(
-        staticCode,
-        SecCSFlags(rawValue: kSecCSSigningInformation),
-        &signingInfo
-      ) == errSecSuccess
-    else { return false }
-    return (signingInfo as? [String: Any])?[kSecCodeInfoTeamIdentifier as String] as? String == teamIdentifier
-  }
-
-  private static func runLaunchctl(_ arguments: [String]) -> Int32 {
-    let values = ["launchctl"] + arguments
-    var argv = values.map { strdup($0) }
-    argv.append(nil)
-    defer { argv.compactMap { $0 }.forEach { free($0) } }
-
-    var pid: pid_t = 0
-    let result = posix_spawn(&pid, "/bin/launchctl", nil, nil, &argv, environ)
-    guard result == 0 else { return Int32(result) }
-    var status: Int32 = 0
-    guard waitpid(pid, &status, 0) == pid else { return -1 }
-    return status
-  }
-}
-
-private enum HelperInstallError: Error {
-  case invalidSignature, registrationFailed, replacementFailed, rollbackFailed
-}
-
 class ReactNativeDelegate: ExpoReactNativeFactoryDelegate {
   private let bundleOverrideURL: URL?