bleam 0.0.10 → 0.0.11

package/templates/native/ios/Bleam/AI/TextGeneration/Bonsai/BonsaiTextSnapshotValidator.swift

@@ -0,0 +1,201 @@
+import Foundation
+
+public enum BonsaiTextSnapshotValidator {
+    private struct Quantization: Decodable, Equatable {
+        let groupSize: Int
+        let bits: Int
+        let mode: String
+
+        private enum CodingKeys: String, CodingKey {
+            case groupSize = "group_size"
+            case bits
+            case mode
+        }
+
+        init(groupSize: Int, bits: Int, mode: String = "affine") {
+            self.groupSize = groupSize
+            self.bits = bits
+            self.mode = mode
+        }
+
+        init(from decoder: Decoder) throws {
+            let values = try decoder.container(keyedBy: CodingKeys.self)
+            groupSize = try values.decode(Int.self, forKey: .groupSize)
+            bits = try values.decode(Int.self, forKey: .bits)
+            mode = try values.decodeIfPresent(String.self, forKey: .mode) ?? "affine"
+        }
+    }
+
+    private struct Qwen3Identity: Decodable {
+        struct RopeScaling: Decodable {
+            let ropeType: String
+            let factor: Float
+            let originalMaximumPositions: Int
+
+            private enum CodingKeys: String, CodingKey {
+                case ropeType = "rope_type"
+                case factor
+                case originalMaximumPositions = "original_max_position_embeddings"
+            }
+        }
+
+        let hiddenActivation: String
+        let layerTypes: [String]
+        let ropeScaling: RopeScaling
+        let quantization: Quantization
+
+        private enum CodingKeys: String, CodingKey {
+            case hiddenActivation = "hidden_act"
+            case layerTypes = "layer_types"
+            case ropeScaling = "rope_scaling"
+            case quantization
+        }
+    }
+
+    private struct Qwen35TextIdentity: Decodable {
+        let modelType: String
+        let vocabularySize: Int
+        let hiddenSize: Int
+        let intermediateSize: Int
+        let hiddenLayerCount: Int
+        let attentionHeadCount: Int
+        let keyValueHeadCount: Int
+        let headDimension: Int
+        let maximumPositionCount: Int
+        let quantization: Quantization
+
+        private enum CodingKeys: String, CodingKey {
+            case modelType = "model_type"
+            case vocabularySize = "vocab_size"
+            case hiddenSize = "hidden_size"
+            case intermediateSize = "intermediate_size"
+            case hiddenLayerCount = "num_hidden_layers"
+            case attentionHeadCount = "num_attention_heads"
+            case keyValueHeadCount = "num_key_value_heads"
+            case headDimension = "head_dim"
+            case maximumPositionCount = "max_position_embeddings"
+            case quantization
+        }
+    }
+
+    private struct Qwen35Identity: Decodable {
+        let architectures: [String]
+        let modelType: String
+        let textConfiguration: Qwen35TextIdentity
+        let quantization: Quantization
+
+        private enum CodingKeys: String, CodingKey {
+            case architectures
+            case modelType = "model_type"
+            case textConfiguration = "text_config"
+            case quantization
+        }
+    }
+
+    public static func validate(modelURL: URL, profile: BonsaiTextModelProfile) throws {
+        var isDirectory: ObjCBool = false
+        guard FileManager.default.fileExists(atPath: modelURL.path, isDirectory: &isDirectory),
+              isDirectory.boolValue else {
+            throw BonsaiTextError.invalidModelPath(modelURL.path)
+        }
+
+        do {
+            switch profile.backend {
+            case .qwen3:
+                try validateQwen3(modelURL: modelURL, profile: profile)
+            case .qwen35:
+                try validateQwen35(modelURL: modelURL, profile: profile)
+            }
+        } catch let error as BonsaiTextError {
+            throw error
+        } catch {
+            throw BonsaiTextError.snapshotMismatch(model: profile.id, detail: error.localizedDescription)
+        }
+    }
+
+    private static func validateQwen3(modelURL: URL, profile: BonsaiTextModelProfile) throws {
+        let configuration = try Qwen3ModelConfiguration.load(from: modelURL)
+        try configuration.validate()
+        let data = try Data(contentsOf: modelURL.appendingPathComponent("config.json"))
+        let identity = try JSONDecoder().decode(Qwen3Identity.self, from: data)
+        guard configuration.architectures == ["Qwen3ForCausalLM"], configuration.modelType == "qwen3" else {
+            throw mismatch(profile, "expected Qwen3ForCausalLM/qwen3 architecture")
+        }
+        guard configuration.vocabularySize == profile.vocabularySize,
+              configuration.hiddenSize == profile.hiddenSize,
+              configuration.intermediateSize == profile.intermediateSize,
+              configuration.hiddenLayerCount == profile.hiddenLayerCount,
+              configuration.attentionHeadCount == profile.attentionHeadCount,
+              configuration.keyValueHeadCount == profile.keyValueHeadCount,
+              configuration.headDimension == profile.headDimension,
+              configuration.maximumPositionCount == profile.contextLength else {
+            throw mismatch(profile, "Qwen3 dimensions do not match the stable model ID")
+        }
+        let expectedRopeTheta: Float = switch profile.id {
+        case .ternaryBonsai4B: 5_000_000
+        default: 1_000_000
+        }
+        guard configuration.rmsNormEpsilon == 1e-6,
+              configuration.ropeTheta == expectedRopeTheta,
+              configuration.attentionBias == false,
+              configuration.tieWordEmbeddings == (profile.id != .ternaryBonsai8B),
+              identity.hiddenActivation == "silu",
+              identity.layerTypes.count == profile.hiddenLayerCount,
+              identity.layerTypes.allSatisfy({ $0 == "full_attention" }),
+              identity.ropeScaling.ropeType == "yarn",
+              identity.ropeScaling.factor == 4,
+              identity.ropeScaling.originalMaximumPositions == profile.contextLength / 4 else {
+            throw mismatch(profile, "Qwen3 architecture settings do not match the stable model ID")
+        }
+        guard configuration.quantization?.bits == 2,
+              configuration.quantization?.groupSize == 128,
+              identity.quantization == Quantization(groupSize: 128, bits: 2) else {
+            throw mismatch(profile, "expected packed 2-bit/group-128 quantization")
+        }
+    }
+
+    private static func validateQwen35(modelURL: URL, profile: BonsaiTextModelProfile) throws {
+        let data = try Data(contentsOf: modelURL.appendingPathComponent("config.json"))
+        let identity = try JSONDecoder().decode(Qwen35Identity.self, from: data)
+        let text = identity.textConfiguration
+        guard identity.architectures == ["Qwen3_5ForConditionalGeneration"],
+              identity.modelType == "qwen3_5", text.modelType == "qwen3_5_text" else {
+            throw mismatch(profile, "expected Qwen3.5 conditional-generation/text architecture")
+        }
+        guard text.vocabularySize == profile.vocabularySize,
+              text.hiddenSize == profile.hiddenSize,
+              text.intermediateSize == profile.intermediateSize,
+              text.hiddenLayerCount == profile.hiddenLayerCount,
+              text.attentionHeadCount == profile.attentionHeadCount,
+              text.keyValueHeadCount == profile.keyValueHeadCount,
+              text.headDimension == profile.headDimension,
+              text.maximumPositionCount == profile.contextLength else {
+            throw mismatch(profile, "Qwen3.5 dimensions do not match the stable model ID")
+        }
+        let requiredQuantization = Quantization(groupSize: 128, bits: 2)
+        guard identity.quantization == requiredQuantization,
+              text.quantization == requiredQuantization else {
+            throw mismatch(profile, "Qwen3.5 requires 2-bit/group-128 at both snapshot levels")
+        }
+
+        let backendConfiguration = try Qwen35Configuration.load(from: modelURL)
+        try backendConfiguration.validate(expectedQuantization: Qwen35Quantization(groupSize: 128, bits: 2))
+        guard backendConfiguration.rmsNormEps == 1e-6,
+              backendConfiguration.attentionBias == false,
+              backendConfiguration.tieWordEmbeddings == false,
+              backendConfiguration.fullAttentionInterval == 4,
+              backendConfiguration.linearConvKernelDim == 4,
+              backendConfiguration.linearKeyHeadDim == 128,
+              backendConfiguration.linearValueHeadDim == 128,
+              backendConfiguration.linearNumKeyHeads == 16,
+              backendConfiguration.linearNumValueHeads == 48,
+              backendConfiguration.ropeTheta == 10_000_000,
+              backendConfiguration.partialRotaryFactor == 0.25 else {
+            throw mismatch(profile, "Qwen3.5 architecture settings do not match the stable model ID")
+        }
+    }
+
+    private static func mismatch(_ profile: BonsaiTextModelProfile, _ detail: String) -> BonsaiTextError {
+        .snapshotMismatch(model: profile.id, detail: detail)
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/GenerationTypes.swift

@@ -0,0 +1,85 @@
+import Foundation
+
+public struct GenerationParameters: Sendable {
+    public var maximumNewTokens: Int
+    public var temperature: Float
+    public var topP: Float
+    public var topK: Int?
+    public var repetitionPenalty: Float
+    public var repetitionContextSize: Int
+    public var seed: UInt64?
+
+    public init(
+        maximumNewTokens: Int = 512,
+        temperature: Float = 0.7,
+        topP: Float = 0.95,
+        topK: Int? = nil,
+        repetitionPenalty: Float = 1,
+        repetitionContextSize: Int = 64,
+        seed: UInt64? = nil
+    ) {
+        self.maximumNewTokens = maximumNewTokens
+        self.temperature = temperature
+        self.topP = topP
+        self.topK = topK
+        self.repetitionPenalty = repetitionPenalty
+        self.repetitionContextSize = repetitionContextSize
+        self.seed = seed
+    }
+
+    public static let balanced = GenerationParameters()
+}
+
+public enum GenerationFinishReason: String, Codable, Sendable {
+    case endOfSequence
+    case length
+    case cancelled
+    case callbackStopped
+}
+
+public struct TokenGenerationResult: Sendable {
+    public let promptTokenCount: Int
+    public let tokenIDs: [Int]
+    public let finishReason: GenerationFinishReason
+
+    public init(promptTokenCount: Int, tokenIDs: [Int], finishReason: GenerationFinishReason) {
+        self.promptTokenCount = promptTokenCount
+        self.tokenIDs = tokenIDs
+        self.finishReason = finishReason
+    }
+}
+
+public enum TextGenerationError: LocalizedError {
+    case emptyPrompt
+    case invalidMaximumNewTokens(Int)
+    case contextExceeded(requested: Int, maximum: Int)
+    case invalidTokenID(Int)
+    case invalidTemperature(Float)
+    case invalidTopP(Float)
+    case invalidTopK(Int)
+    case invalidRepetitionPenalty(Float)
+    case invalidRepetitionContextSize(Int)
+
+    public var errorDescription: String? {
+        switch self {
+        case .emptyPrompt:
+            "The pre-tokenized prompt is empty"
+        case .invalidMaximumNewTokens(let value):
+            "Maximum new tokens must be nonnegative, got \(value)"
+        case .contextExceeded(let requested, let maximum):
+            "Requested context length \(requested) exceeds model maximum \(maximum)"
+        case .invalidTokenID(let value):
+            "Token ID \(value) is outside the model vocabulary"
+        case .invalidTemperature(let value):
+            "Temperature must be nonnegative, got \(value)"
+        case .invalidTopP(let value):
+            "Top-p must be in (0, 1], got \(value)"
+        case .invalidTopK(let value):
+            "Top-k must be positive, got \(value)"
+        case .invalidRepetitionPenalty(let value):
+            "Repetition penalty must be positive, got \(value)"
+        case .invalidRepetitionContextSize(let value):
+            "Repetition context size must be nonnegative, got \(value)"
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/ModelCapabilities.swift

@@ -0,0 +1,21 @@
+import MLX
+
+/// Encoding is an independent capability; callers do not need a causal LM head.
+public protocol HiddenStateModel: AnyObject {
+    var vocabularySize: Int { get }
+    var maximumContextLength: Int { get }
+
+    func hiddenStates(
+        for inputIDs: MLXArray,
+        at layerIndices: Set<Int>,
+        attentionMask: MLXArray?
+    ) -> [Int: MLXArray]
+}
+
+public protocol CausalLanguageModel: AnyObject {
+    var vocabularySize: Int { get }
+    var maximumContextLength: Int { get }
+
+    func makeCache() -> [TextKVCache]
+    func logits(for inputIDs: MLXArray, cache: [TextKVCache]) -> MLXArray
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/TextKVCache.swift

@@ -0,0 +1,23 @@
+import MLX
+
+public final class TextKVCache {
+    private var keys: MLXArray?
+    private var values: MLXArray?
+
+    public private(set) var length = 0
+
+    public init() {}
+
+    public func update(keys newKeys: MLXArray, values newValues: MLXArray) -> (MLXArray, MLXArray) {
+        keys = keys.map { concatenated([$0, newKeys], axis: 2) } ?? newKeys
+        values = values.map { concatenated([$0, newValues], axis: 2) } ?? newValues
+        length += newKeys.dim(2)
+        return (keys!, values!)
+    }
+
+    public func clear() {
+        keys = nil
+        values = nil
+        length = 0
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/TokenSampler.swift

@@ -0,0 +1,79 @@
+import MLX
+
+public enum TokenSampler {
+    public static func validate(_ parameters: GenerationParameters) throws {
+        guard parameters.maximumNewTokens >= 0 else {
+            throw TextGenerationError.invalidMaximumNewTokens(parameters.maximumNewTokens)
+        }
+        guard parameters.temperature >= 0 else {
+            throw TextGenerationError.invalidTemperature(parameters.temperature)
+        }
+        guard parameters.topP > 0, parameters.topP <= 1 else {
+            throw TextGenerationError.invalidTopP(parameters.topP)
+        }
+        if let topK = parameters.topK, topK <= 0 {
+            throw TextGenerationError.invalidTopK(topK)
+        }
+        guard parameters.repetitionPenalty > 0 else {
+            throw TextGenerationError.invalidRepetitionPenalty(parameters.repetitionPenalty)
+        }
+        guard parameters.repetitionContextSize >= 0 else {
+            throw TextGenerationError.invalidRepetitionContextSize(parameters.repetitionContextSize)
+        }
+    }
+
+    public static func sample(
+        logits: MLXArray,
+        parameters: GenerationParameters,
+        tokenHistory: [Int]
+    ) throws -> MLXArray {
+        try validate(parameters)
+        let selected: MLXArray
+        switch logits.ndim {
+        case 3: selected = logits[0, -1]
+        case 2: selected = logits[0]
+        default: selected = logits
+        }
+        let values = selected
+
+        if parameters.repetitionPenalty != 1,
+           parameters.repetitionContextSize > 0,
+           !tokenHistory.isEmpty {
+            let recent = Set(tokenHistory.suffix(parameters.repetitionContextSize))
+            let valid = recent.filter { $0 >= 0 && $0 < values.dim(-1) }
+            if !valid.isEmpty {
+                let indices = MLXArray(valid.sorted().map(Int32.init))
+                let repeated = MLX.take(values, indices, axis: -1)
+                values[indices] = MLX.where(
+                    repeated .> 0,
+                    repeated / parameters.repetitionPenalty,
+                    repeated * parameters.repetitionPenalty
+                )
+            }
+        }
+
+        if parameters.temperature == 0 {
+            return argMax(values).asType(Int32.self)
+        }
+
+        // Match mlx-lm ordering: select top-p from untempered probabilities.
+        let probabilities = softmax(values, axis: -1)
+        let candidateIndices: MLXArray
+        if let topK = parameters.topK, topK < probabilities.dim(-1) {
+            candidateIndices = argPartition(probabilities, kth: -topK, axis: -1)[(-topK)...]
+        } else {
+            candidateIndices = MLXArray.arange(probabilities.dim(-1))
+        }
+        let candidates = take(probabilities, candidateIndices, axis: -1)
+        let order = argSort(-candidates, axis: -1)
+        let sortedIndices = take(candidateIndices, order, axis: -1)
+        let sortedProbabilities = take(candidates, order, axis: -1)
+        let cumulativeBefore = cumsum(sortedProbabilities, axis: -1) - sortedProbabilities
+        let filtered = MLX.where(
+            cumulativeBefore .< parameters.topP,
+            log(sortedProbabilities) / parameters.temperature,
+            MLX.full(sortedProbabilities.shape, values: -Float.infinity)
+        )
+        return sortedIndices[MLXRandom.categorical(filtered)].asType(Int32.self)
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen3/Qwen3CheckpointLoader.swift

@@ -0,0 +1,161 @@
+import Foundation
+import MLX
+import MLXNN
+
+public enum Qwen3CheckpointError: LocalizedError {
+    case missingFile(String)
+    case unsupportedArchitecture(String)
+    case unsupportedQuantization(bits: Int?, groupSize: Int?)
+    case invalidConfiguration(String)
+    case invalidCheckpoint(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .missingFile(let file): "Qwen3 checkpoint is missing \(file)"
+        case .unsupportedArchitecture(let architecture): "Unsupported Qwen3 architecture \(architecture)"
+        case .unsupportedQuantization(let bits, let groupSize):
+            "Bonsai Qwen3 requires 2-bit/group-128 affine weights, got bits=\(bits.map(String.init) ?? "nil") group=\(groupSize.map(String.init) ?? "nil")"
+        case .invalidConfiguration(let message): "Invalid Qwen3 configuration: \(message)"
+        case .invalidCheckpoint(let message): "Invalid Qwen3 checkpoint: \(message)"
+        }
+    }
+}
+
+public enum Qwen3CheckpointLoader {
+    private struct SafetensorsIndex: Decodable {
+        let weightMap: [String: String]
+
+        enum CodingKeys: String, CodingKey {
+            case weightMap = "weight_map"
+        }
+    }
+
+    public static func load(from modelURL: URL) throws -> Qwen3CausalLM {
+        let (configuration, quantization, tensors) = try loadSnapshot(from: modelURL, requireLanguageHead: true)
+
+        let model = Qwen3CausalLM(configuration: configuration)
+        MLXNN.quantize(
+            model: model,
+            groupSize: quantization.groupSize,
+            bits: quantization.bits,
+            mode: .affine
+        )
+        do {
+            try model.update(
+                parameters: ModuleParameters.unflattened(tensors),
+                verify: [.noUnusedKeys, .shapeMismatch]
+            )
+        } catch {
+            throw Qwen3CheckpointError.invalidCheckpoint(error.localizedDescription)
+        }
+        model.usesTiedOutputHead = configuration.tieWordEmbeddings && tensors["lm_head.weight"] == nil
+        MLX.eval(model)
+        MLX.Memory.clearCache()
+        return model
+    }
+
+    public static func loadTransformer(from modelURL: URL) throws -> Qwen3TransformerModel {
+        let (configuration, quantization, tensors) = try loadSnapshot(from: modelURL, requireLanguageHead: false)
+        let model = Qwen3TransformerModel(configuration: configuration)
+        MLXNN.quantize(
+            model: model,
+            groupSize: quantization.groupSize,
+            bits: quantization.bits,
+            mode: .affine
+        )
+        let transformerTensors: [String: MLXArray] = Dictionary(uniqueKeysWithValues: tensors.compactMap { name, value in
+            guard name.hasPrefix("model.") else { return nil }
+            return (String(name.dropFirst("model.".count)), value)
+        })
+        do {
+            try model.update(
+                parameters: ModuleParameters.unflattened(transformerTensors),
+                verify: [.noUnusedKeys, .shapeMismatch]
+            )
+        } catch {
+            throw Qwen3CheckpointError.invalidCheckpoint(error.localizedDescription)
+        }
+        MLX.eval(model)
+        MLX.Memory.clearCache()
+        return model
+    }
+
+    private static func loadSnapshot(
+        from modelURL: URL,
+        requireLanguageHead: Bool
+    ) throws -> (Qwen3ModelConfiguration, Qwen3ModelConfiguration.Quantization, [String: MLXArray]) {
+        let configuration = try Qwen3ModelConfiguration.load(from: modelURL)
+        try configuration.validate()
+        guard let quantization = configuration.quantization,
+              quantization.bits == 2, quantization.groupSize == 128 else {
+            throw Qwen3CheckpointError.unsupportedQuantization(
+                bits: configuration.quantization?.bits,
+                groupSize: configuration.quantization?.groupSize
+            )
+        }
+        let files = try checkpointFiles(at: modelURL)
+        var tensors: [String: MLXArray] = [:]
+        for file in files {
+            for (name, value) in try loadArrays(url: modelURL.appendingPathComponent(file)) {
+                tensors[name] = value
+            }
+        }
+        try validateTensors(
+            tensors,
+            configuration: configuration,
+            requireLanguageHead: requireLanguageHead
+        )
+        return (configuration, quantization, tensors)
+    }
+
+    private static func checkpointFiles(at modelURL: URL) throws -> [String] {
+        let fileManager = FileManager.default
+        let indexURL = modelURL.appendingPathComponent("model.safetensors.index.json")
+        if fileManager.fileExists(atPath: indexURL.path) {
+            let index = try JSONDecoder().decode(SafetensorsIndex.self, from: Data(contentsOf: indexURL))
+            let files = Set(index.weightMap.values).sorted()
+            for file in files where !fileManager.fileExists(atPath: modelURL.appendingPathComponent(file).path) {
+                throw Qwen3CheckpointError.missingFile(file)
+            }
+            return files
+        }
+        let file = "model.safetensors"
+        guard fileManager.fileExists(atPath: modelURL.appendingPathComponent(file).path) else {
+            throw Qwen3CheckpointError.missingFile(file)
+        }
+        return [file]
+    }
+
+    private static func validateTensors(
+        _ tensors: [String: MLXArray],
+        configuration: Qwen3ModelConfiguration,
+        requireLanguageHead: Bool
+    ) throws {
+        var requiredModules = ["model.embed_tokens"]
+        for index in 0..<configuration.hiddenLayerCount {
+            let prefix = "model.layers.\(index)"
+            requiredModules += [
+                "\(prefix).self_attn.q_proj", "\(prefix).self_attn.k_proj",
+                "\(prefix).self_attn.v_proj", "\(prefix).self_attn.o_proj",
+                "\(prefix).mlp.gate_proj", "\(prefix).mlp.up_proj", "\(prefix).mlp.down_proj",
+            ]
+            for norm in ["input_layernorm", "post_attention_layernorm", "self_attn.q_norm", "self_attn.k_norm"] {
+                let key = "\(prefix).\(norm).weight"
+                guard tensors[key] != nil else { throw Qwen3CheckpointError.invalidCheckpoint("missing tensor \(key)") }
+            }
+        }
+        if requireLanguageHead, !configuration.tieWordEmbeddings { requiredModules.append("lm_head") }
+        for module in requiredModules {
+            for suffix in ["weight", "scales", "biases"] {
+                let key = "\(module).\(suffix)"
+                guard tensors[key] != nil else { throw Qwen3CheckpointError.invalidCheckpoint("missing tensor \(key)") }
+            }
+            guard tensors["\(module).weight"]?.dtype == .uint32 else {
+                throw Qwen3CheckpointError.invalidCheckpoint("\(module).weight is not packed uint32")
+            }
+        }
+        guard tensors["model.norm.weight"] != nil else {
+            throw Qwen3CheckpointError.invalidCheckpoint("missing tensor model.norm.weight")
+        }
+    }
+}