bleam 0.0.10 → 0.0.11

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

@@ -0,0 +1,119 @@
+import Foundation
+
+public struct Qwen3ModelConfiguration: Decodable, Sendable {
+    public struct RopeScaling: Decodable, Sendable {
+        public let ropeType: String
+        public let factor: Float
+        public let originalMaximumPositions: Int
+        public let betaFast: Float
+        public let betaSlow: Float
+
+        enum CodingKeys: String, CodingKey {
+            case ropeType = "rope_type"
+            case type
+            case factor
+            case originalMaximumPositions = "original_max_position_embeddings"
+            case betaFast = "beta_fast"
+            case betaSlow = "beta_slow"
+        }
+
+        public init(from decoder: Decoder) throws {
+            let values = try decoder.container(keyedBy: CodingKeys.self)
+            ropeType = try values.decodeIfPresent(String.self, forKey: .ropeType)
+                ?? values.decodeIfPresent(String.self, forKey: .type)
+                ?? "yarn"
+            factor = try values.decode(Float.self, forKey: .factor)
+            originalMaximumPositions = try values.decode(Int.self, forKey: .originalMaximumPositions)
+            betaFast = try values.decodeIfPresent(Float.self, forKey: .betaFast) ?? 32
+            betaSlow = try values.decodeIfPresent(Float.self, forKey: .betaSlow) ?? 1
+        }
+    }
+
+    public struct Quantization: Decodable, Sendable {
+        public let bits: Int
+        public let groupSize: Int
+
+        enum CodingKeys: String, CodingKey {
+            case bits
+            case groupSize = "group_size"
+        }
+    }
+
+    public let architectures: [String]
+    public let modelType: String
+    public let vocabularySize: Int
+    public let hiddenSize: Int
+    public let intermediateSize: Int
+    public let hiddenLayerCount: Int
+    public let attentionHeadCount: Int
+    public let keyValueHeadCount: Int
+    public let maximumPositionCount: Int
+    public let rmsNormEpsilon: Float
+    public let ropeTheta: Float
+    public let tieWordEmbeddings: Bool
+    public let attentionBias: Bool
+    public let headDimension: Int
+    public let ropeScaling: RopeScaling?
+    public let quantization: Quantization?
+
+    enum CodingKeys: String, CodingKey {
+        case architectures
+        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 maximumPositionCount = "max_position_embeddings"
+        case rmsNormEpsilon = "rms_norm_eps"
+        case ropeTheta = "rope_theta"
+        case tieWordEmbeddings = "tie_word_embeddings"
+        case attentionBias = "attention_bias"
+        case headDimension = "head_dim"
+        case ropeScaling = "rope_scaling"
+        case quantization
+    }
+
+    public init(from decoder: Decoder) throws {
+        let values = try decoder.container(keyedBy: CodingKeys.self)
+        architectures = try values.decodeIfPresent([String].self, forKey: .architectures) ?? []
+        modelType = try values.decodeIfPresent(String.self, forKey: .modelType) ?? "qwen3"
+        vocabularySize = try values.decode(Int.self, forKey: .vocabularySize)
+        hiddenSize = try values.decode(Int.self, forKey: .hiddenSize)
+        intermediateSize = try values.decode(Int.self, forKey: .intermediateSize)
+        hiddenLayerCount = try values.decode(Int.self, forKey: .hiddenLayerCount)
+        attentionHeadCount = try values.decode(Int.self, forKey: .attentionHeadCount)
+        keyValueHeadCount = try values.decode(Int.self, forKey: .keyValueHeadCount)
+        maximumPositionCount = try values.decode(Int.self, forKey: .maximumPositionCount)
+        rmsNormEpsilon = try values.decodeIfPresent(Float.self, forKey: .rmsNormEpsilon) ?? 1e-6
+        ropeTheta = try values.decodeIfPresent(Float.self, forKey: .ropeTheta) ?? 1_000_000
+        tieWordEmbeddings = try values.decodeIfPresent(Bool.self, forKey: .tieWordEmbeddings) ?? true
+        attentionBias = try values.decodeIfPresent(Bool.self, forKey: .attentionBias) ?? false
+        headDimension = try values.decodeIfPresent(Int.self, forKey: .headDimension)
+            ?? hiddenSize / attentionHeadCount
+        ropeScaling = try values.decodeIfPresent(RopeScaling.self, forKey: .ropeScaling)
+        quantization = try values.decodeIfPresent(Quantization.self, forKey: .quantization)
+    }
+
+    public static func load(from modelURL: URL) throws -> Qwen3ModelConfiguration {
+        try JSONDecoder().decode(
+            Qwen3ModelConfiguration.self,
+            from: Data(contentsOf: modelURL.appendingPathComponent("config.json"))
+        )
+    }
+
+    public func validate() throws {
+        guard modelType == "qwen3", architectures.isEmpty || architectures.contains("Qwen3ForCausalLM") else {
+            throw Qwen3CheckpointError.unsupportedArchitecture(architectures.first ?? modelType)
+        }
+        guard hiddenSize > 0, intermediateSize > 0, hiddenLayerCount > 0,
+              attentionHeadCount > 0, keyValueHeadCount > 0,
+              attentionHeadCount.isMultiple(of: keyValueHeadCount), headDimension > 0 else {
+            throw Qwen3CheckpointError.invalidConfiguration("invalid model dimensions")
+        }
+        if let ropeScaling, ropeScaling.ropeType != "yarn" {
+            throw Qwen3CheckpointError.invalidConfiguration("unsupported RoPE scaling type \(ropeScaling.ropeType)")
+        }
+    }
+}

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

@@ -0,0 +1,84 @@
+import MLX
+
+public final class Qwen3InferenceSession: @unchecked Sendable {
+    private let model: Qwen3CausalLM
+    private let endTokenIDs: Set<Int>
+
+    public init(model: Qwen3CausalLM, endTokenIDs: Set<Int> = [151643, 151645]) {
+        self.model = model
+        self.endTokenIDs = endTokenIDs
+    }
+
+    public func generate(
+        inputIDs: [Int],
+        parameters: GenerationParameters = .balanced,
+        isCancelled: (() -> Bool)? = nil,
+        onToken: ((Int) -> Bool)? = nil
+    ) throws -> TokenGenerationResult {
+        try TokenSampler.validate(parameters)
+        guard !inputIDs.isEmpty else { throw TextGenerationError.emptyPrompt }
+        for token in inputIDs where token < 0 || token >= model.vocabularySize {
+            throw TextGenerationError.invalidTokenID(token)
+        }
+        let requestedLength = inputIDs.count + parameters.maximumNewTokens
+        guard requestedLength <= model.maximumContextLength else {
+            throw TextGenerationError.contextExceeded(
+                requested: requestedLength,
+                maximum: model.maximumContextLength
+            )
+        }
+        guard parameters.maximumNewTokens > 0 else {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .length)
+        }
+        if isCancelled?() == true {
+            return TokenGenerationResult(promptTokenCount: inputIDs.count, tokenIDs: [], finishReason: .cancelled)
+        }
+        if let seed = parameters.seed { MLXRandom.seed(seed) }
+
+        let cache = model.makeCache()
+        defer {
+            cache.forEach { $0.clear() }
+            MLX.Memory.clearCache()
+        }
+        var logits = model.logits(
+            for: MLXArray(inputIDs.map(Int32.init)).reshaped([1, inputIDs.count]),
+            cache: cache
+        )
+        var generated: [Int] = []
+        var history = inputIDs
+        var finishReason: GenerationFinishReason = .length
+
+        for _ in 0..<parameters.maximumNewTokens {
+            if isCancelled?() == true {
+                finishReason = .cancelled
+                break
+            }
+            let sampled = try TokenSampler.sample(
+                logits: logits,
+                parameters: parameters,
+                tokenHistory: history
+            )
+            MLX.eval(sampled)
+            let token = Int(sampled.item(Int32.self))
+            if endTokenIDs.contains(token) {
+                finishReason = .endOfSequence
+                break
+            }
+            generated.append(token)
+            history.append(token)
+            if onToken?(token) == false {
+                finishReason = .callbackStopped
+                break
+            }
+            logits = model.logits(
+                for: MLXArray([Int32(token)]).reshaped([1, 1]),
+                cache: cache
+            )
+        }
+        return TokenGenerationResult(
+            promptTokenCount: inputIDs.count,
+            tokenIDs: generated,
+            finishReason: finishReason
+        )
+    }
+}

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

@@ -0,0 +1,197 @@
+import Foundation
+import MLX
+import MLXNN
+
+public final class Qwen3RMSNorm: Module, UnaryLayer {
+    @ParameterInfo public var weight: MLXArray
+    public let epsilon: Float
+
+    public init(dimensions: Int, epsilon: Float) {
+        _weight = ParameterInfo(wrappedValue: MLXArray.ones([dimensions]))
+        self.epsilon = epsilon
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        MLXFast.rmsNorm(x, weight: weight, eps: epsilon)
+    }
+}
+
+public final class Qwen3RotaryEmbedding: Module {
+    private let dimensions: Int
+    private let base: Float
+    private let frequencies: MLXArray?
+    private let magnitudeScale: Float
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        dimensions = configuration.headDimension
+        base = configuration.ropeTheta
+        if let scaling = configuration.ropeScaling {
+            precondition(configuration.headDimension.isMultiple(of: 2))
+            func correctionDimension(_ rotations: Float) -> Float {
+                Float(configuration.headDimension)
+                    * log(Float(scaling.originalMaximumPositions) / (rotations * 2 * .pi))
+                    / (2 * log(configuration.ropeTheta))
+            }
+            let low = max(Int(floor(correctionDimension(scaling.betaFast))), 0)
+            let high = min(Int(ceil(correctionDimension(scaling.betaSlow))), configuration.headDimension - 1)
+            let safeHigh = low == high ? Float(high) + 0.001 : Float(high)
+            let ramp = clip(
+                (MLXArray(0..<(configuration.headDimension / 2)).asType(.float32) - Float(low))
+                    / (safeHigh - Float(low)),
+                min: 0,
+                max: 1
+            )
+            let exponent = MLXArray(stride(from: 0, to: configuration.headDimension, by: 2))
+                .asType(.float32) / Float(configuration.headDimension)
+            let unscaled = pow(configuration.ropeTheta, exponent)
+            let interpolated = scaling.factor * unscaled
+            let mask = 1 - ramp
+            frequencies = (interpolated * unscaled) / (interpolated * mask + unscaled * (1 - mask))
+            magnitudeScale = scaling.factor <= 1 ? 1 : 0.1 * log(scaling.factor) + 1
+        } else {
+            frequencies = nil
+            magnitudeScale = 1
+        }
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray, offset: Int) -> MLXArray {
+        MLXFast.RoPE(
+            magnitudeScale == 1 ? x : x * magnitudeScale,
+            dimensions: dimensions,
+            traditional: false,
+            base: frequencies == nil ? base : nil,
+            scale: 1,
+            offset: offset,
+            freqs: frequencies
+        )
+    }
+}
+
+public final class Qwen3SelfAttention: Module {
+    private let headCount: Int
+    private let keyValueHeadCount: Int
+    private let headDimension: Int
+    private let scale: Float
+
+    @ModuleInfo public var q_proj: Linear
+    @ModuleInfo public var k_proj: Linear
+    @ModuleInfo public var v_proj: Linear
+    @ModuleInfo public var o_proj: Linear
+    @ModuleInfo public var q_norm: Qwen3RMSNorm
+    @ModuleInfo public var k_norm: Qwen3RMSNorm
+    @ModuleInfo public var rope: Qwen3RotaryEmbedding
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        headCount = configuration.attentionHeadCount
+        keyValueHeadCount = configuration.keyValueHeadCount
+        headDimension = configuration.headDimension
+        scale = 1 / sqrt(Float(configuration.headDimension))
+        _q_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.hiddenSize,
+            configuration.attentionHeadCount * configuration.headDimension,
+            bias: configuration.attentionBias
+        ))
+        _k_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.hiddenSize,
+            configuration.keyValueHeadCount * configuration.headDimension,
+            bias: configuration.attentionBias
+        ))
+        _v_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.hiddenSize,
+            configuration.keyValueHeadCount * configuration.headDimension,
+            bias: configuration.attentionBias
+        ))
+        _o_proj = ModuleInfo(wrappedValue: Linear(
+            configuration.attentionHeadCount * configuration.headDimension,
+            configuration.hiddenSize,
+            bias: configuration.attentionBias
+        ))
+        _q_norm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.headDimension,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        _k_norm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.headDimension,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        _rope = ModuleInfo(wrappedValue: Qwen3RotaryEmbedding(configuration: configuration))
+        super.init()
+    }
+
+    public func callAsFunction(_ hiddenStates: MLXArray, mask: MLXArray?, cache: TextKVCache?) -> MLXArray {
+        let batchSize = hiddenStates.dim(0)
+        let sequenceLength = hiddenStates.dim(1)
+        var queries = q_proj(hiddenStates).reshaped([batchSize, sequenceLength, headCount, headDimension])
+        var keys = k_proj(hiddenStates).reshaped([batchSize, sequenceLength, keyValueHeadCount, headDimension])
+        var values = v_proj(hiddenStates).reshaped([batchSize, sequenceLength, keyValueHeadCount, headDimension])
+        queries = q_norm(queries).transposed(0, 2, 1, 3)
+        keys = k_norm(keys).transposed(0, 2, 1, 3)
+        values = values.transposed(0, 2, 1, 3)
+        let offset = cache?.length ?? 0
+        queries = rope(queries, offset: offset)
+        keys = rope(keys, offset: offset)
+        if let cache {
+            (keys, values) = cache.update(keys: keys, values: values)
+        }
+        let attended = MLXFast.scaledDotProductAttention(
+            queries: queries,
+            keys: keys,
+            values: values,
+            scale: scale,
+            mask: mask
+        )
+        return o_proj(attended.transposed(0, 2, 1, 3).reshaped([
+            batchSize, sequenceLength, headCount * headDimension,
+        ]))
+    }
+}
+
+public final class Qwen3FeedForward: Module, @unchecked Sendable {
+    @ModuleInfo public var gate_proj: Linear
+    @ModuleInfo public var up_proj: Linear
+    @ModuleInfo public var down_proj: Linear
+
+    nonisolated(unsafe) private static let activate: (MLXArray, MLXArray) -> MLXArray = compile {
+        silu($0) * $1
+    }
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        _gate_proj = ModuleInfo(wrappedValue: Linear(configuration.hiddenSize, configuration.intermediateSize, bias: false))
+        _up_proj = ModuleInfo(wrappedValue: Linear(configuration.hiddenSize, configuration.intermediateSize, bias: false))
+        _down_proj = ModuleInfo(wrappedValue: Linear(configuration.intermediateSize, configuration.hiddenSize, bias: false))
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        down_proj(Self.activate(gate_proj(x), up_proj(x)))
+    }
+}
+
+public final class Qwen3DecoderBlock: Module {
+    @ModuleInfo public var self_attn: Qwen3SelfAttention
+    @ModuleInfo public var mlp: Qwen3FeedForward
+    @ModuleInfo public var input_layernorm: Qwen3RMSNorm
+    @ModuleInfo public var post_attention_layernorm: Qwen3RMSNorm
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        _self_attn = ModuleInfo(wrappedValue: Qwen3SelfAttention(configuration: configuration))
+        _mlp = ModuleInfo(wrappedValue: Qwen3FeedForward(configuration: configuration))
+        _input_layernorm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.hiddenSize,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        _post_attention_layernorm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.hiddenSize,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        super.init()
+    }
+
+    public func callAsFunction(_ hiddenStates: MLXArray, mask: MLXArray?, cache: TextKVCache?) -> MLXArray {
+        var hidden = hiddenStates + self_attn(input_layernorm(hiddenStates), mask: mask, cache: cache)
+        hidden = hidden + mlp(post_attention_layernorm(hidden))
+        return hidden
+    }
+}

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

@@ -0,0 +1,130 @@
+import MLX
+import MLXNN
+
+public final class Qwen3TransformerModel: Module, HiddenStateModel {
+    public let configuration: Qwen3ModelConfiguration
+
+    @ModuleInfo public var embed_tokens: Embedding
+    public var layers: [Qwen3DecoderBlock]
+    @ModuleInfo public var norm: Qwen3RMSNorm
+
+    public var vocabularySize: Int { configuration.vocabularySize }
+    public var maximumContextLength: Int { configuration.maximumPositionCount }
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        self.configuration = configuration
+        _embed_tokens = ModuleInfo(wrappedValue: Embedding(
+            embeddingCount: configuration.vocabularySize,
+            dimensions: configuration.hiddenSize
+        ))
+        layers = (0..<configuration.hiddenLayerCount).map { _ in
+            Qwen3DecoderBlock(configuration: configuration)
+        }
+        _norm = ModuleInfo(wrappedValue: Qwen3RMSNorm(
+            dimensions: configuration.hiddenSize,
+            epsilon: configuration.rmsNormEpsilon
+        ))
+        super.init()
+    }
+
+    public func callAsFunction(_ inputIDs: MLXArray, cache: [TextKVCache]? = nil) -> MLXArray {
+        var hidden = embed_tokens(inputIDs)
+        let mask = causalMask(sequenceLength: inputIDs.dim(1), offset: cache?.first?.length ?? 0)
+        for (index, layer) in layers.enumerated() {
+            hidden = layer(hidden, mask: mask, cache: cache?[index])
+        }
+        return norm(hidden)
+    }
+
+    public func hiddenStates(
+        for inputIDs: MLXArray,
+        at layerIndices: Set<Int>,
+        attentionMask: MLXArray? = nil
+    ) -> [Int: MLXArray] {
+        var hidden = embed_tokens(inputIDs)
+        let mask = causalMask(sequenceLength: inputIDs.dim(1), offset: 0, attentionMask: attentionMask)
+        var result: [Int: MLXArray] = [:]
+        if layerIndices.contains(0) {
+            MLX.eval(hidden)
+            result[0] = hidden
+        }
+        for (index, layer) in layers.enumerated() {
+            hidden = layer(hidden, mask: mask, cache: nil)
+            let layerIndex = index + 1
+            if layerIndices.contains(layerIndex) {
+                let state = layerIndex == layers.count ? norm(hidden) : hidden
+                MLX.eval(state)
+                result[layerIndex] = state
+            }
+        }
+        return result
+    }
+
+    private func causalMask(
+        sequenceLength: Int,
+        offset: Int,
+        attentionMask: MLXArray? = nil
+    ) -> MLXArray? {
+        if sequenceLength == 1, attentionMask == nil { return nil }
+        let totalLength = sequenceLength + offset
+        let rows = MLXArray.arange(sequenceLength, dtype: .float32).expandedDimensions(axis: 1)
+        let columns = MLXArray.arange(totalLength, dtype: .float32).expandedDimensions(axis: 0)
+        var mask = MLX.where(
+            columns .<= rows + Float(offset),
+            MLXArray(Float(0)),
+            MLXArray(-Float.infinity)
+        ).reshaped([1, 1, sequenceLength, totalLength])
+        if let attentionMask {
+            let padding = MLX.where(
+                attentionMask .== Int32(1),
+                MLXArray(Float(0)),
+                MLXArray(Float(-1e9))
+            ).reshaped([attentionMask.dim(0), 1, 1, attentionMask.dim(1)])
+            mask = mask + padding
+        }
+        return mask
+    }
+}
+
+public final class Qwen3CausalLM: Module, CausalLanguageModel {
+    public let configuration: Qwen3ModelConfiguration
+    @ModuleInfo public var model: Qwen3TransformerModel
+    @ModuleInfo public var lm_head: Linear
+    public var usesTiedOutputHead: Bool
+
+    public var vocabularySize: Int { configuration.vocabularySize }
+    public var maximumContextLength: Int { configuration.maximumPositionCount }
+
+    public init(configuration: Qwen3ModelConfiguration) {
+        self.configuration = configuration
+        usesTiedOutputHead = configuration.tieWordEmbeddings
+        _model = ModuleInfo(wrappedValue: Qwen3TransformerModel(configuration: configuration))
+        _lm_head = ModuleInfo(wrappedValue: Linear(configuration.hiddenSize, configuration.vocabularySize, bias: false))
+        super.init()
+    }
+
+    public func logits(for inputIDs: MLXArray, cache: [TextKVCache]) -> MLXArray {
+        project(model(inputIDs, cache: cache))
+    }
+
+    public func makeCache() -> [TextKVCache] {
+        (0..<configuration.hiddenLayerCount).map { _ in TextKVCache() }
+    }
+
+    private func project(_ hiddenStates: MLXArray) -> MLXArray {
+        guard usesTiedOutputHead else { return lm_head(hiddenStates) }
+        if let embedding = model.embed_tokens as? QuantizedEmbedding {
+            return MLX.quantizedMM(
+                hiddenStates,
+                embedding.weight,
+                scales: embedding.scales,
+                biases: embedding.biases,
+                transpose: true,
+                groupSize: embedding.groupSize,
+                bits: embedding.bits,
+                mode: .affine
+            )
+        }
+        return matmul(hiddenStates, model.embed_tokens.weight.T)
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Qwen35/Qwen35Checkpoint.swift

@@ -0,0 +1,102 @@
+import Foundation
+import MLX
+import MLXNN
+
+public enum Qwen35Checkpoint {
+    private struct Index: Decodable {
+        let weightMap: [String: String]
+
+        private enum CodingKeys: String, CodingKey { case weightMap = "weight_map" }
+    }
+
+    public static func load(
+        from directory: URL,
+        expectedQuantization: Qwen35Quantization? = nil
+    ) throws -> Qwen35ForCausalLM {
+        let config = try Qwen35Configuration.load(from: directory)
+        try config.validate(expectedQuantization: expectedQuantization)
+        let index = try checkpointIndex(in: directory)
+        try validateFiles(index: index, directory: directory, config: config)
+        let model = Qwen35ForCausalLM(config: config)
+        guard config.quantization.mode.lowercased() == "affine" else {
+            throw Qwen35Error.invalidCheckpoint("unsupported quantization mode \(config.quantization.mode)")
+        }
+        MLXNN.quantize(model: model, groupSize: config.quantization.groupSize, bits: config.quantization.bits, mode: .affine)
+
+        var weights: [String: MLXArray] = [:]
+        var usesUpstreamNormOffsets = false
+        for file in Set(index.weightMap.values) {
+            let arrays = try MLX.loadArrays(url: directory.appendingPathComponent(file))
+            usesUpstreamNormOffsets = usesUpstreamNormOffsets || arrays.contains { name, value in
+                name.hasSuffix("conv1d.weight") && value.ndim == 3 && value.dim(2) != 1
+            }
+            for (name, var value) in arrays {
+                guard let key = modelKey(from: name) else { continue }
+                if key.hasSuffix("conv1d.weight"), value.ndim == 3, value.dim(2) != 1 {
+                    value = value.transposed(0, 2, 1)
+                }
+                weights[key] = value
+            }
+        }
+        if usesUpstreamNormOffsets {
+            for key in Array(weights.keys) where isOffsetNorm(key) && weights[key]?.ndim == 1 {
+                weights[key] = weights[key]! + 1
+            }
+        }
+        guard !weights.isEmpty else { throw Qwen35Error.invalidCheckpoint("no language-model tensors") }
+        do {
+            try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.noUnusedKeys, .shapeMismatch])
+        } catch {
+            throw Qwen35Error.invalidCheckpoint(error.localizedDescription)
+        }
+        MLX.eval(model)
+        MLX.Memory.clearCache()
+        return model
+    }
+
+    private static func checkpointIndex(in directory: URL) throws -> Index {
+        let indexURL = directory.appendingPathComponent("model.safetensors.index.json")
+        if FileManager.default.fileExists(atPath: indexURL.path) {
+            return try JSONDecoder().decode(Index.self, from: Data(contentsOf: indexURL))
+        }
+        let single = directory.appendingPathComponent("model.safetensors")
+        guard FileManager.default.fileExists(atPath: single.path) else { throw Qwen35Error.missingFile(single.lastPathComponent) }
+        let names = try MLX.loadArrays(url: single).keys
+        return Index(weightMap: Dictionary(uniqueKeysWithValues: names.map { ($0, single.lastPathComponent) }))
+    }
+
+    private static func validateFiles(index: Index, directory: URL, config: Qwen35Configuration) throws {
+        var required = [
+            "language_model.model.embed_tokens.weight",
+            "language_model.model.layers.0.linear_attn.in_proj_qkv.weight",
+            "language_model.model.layers.0.mlp.gate_proj.weight",
+            "language_model.model.norm.weight",
+        ]
+        if let fullLayer = config.layerTypes.firstIndex(of: "full_attention") {
+            required.append("language_model.model.layers.\(fullLayer).self_attn.q_proj.weight")
+        }
+        if !config.tieWordEmbeddings { required.append("language_model.lm_head.weight") }
+        for key in required where index.weightMap[key] == nil {
+            throw Qwen35Error.invalidCheckpoint("missing tensor \(key)")
+        }
+        for file in Set(index.weightMap.values) where !FileManager.default.fileExists(atPath: directory.appendingPathComponent(file).path) {
+            throw Qwen35Error.missingFile(file)
+        }
+    }
+
+    private static func modelKey(from checkpointKey: String) -> String? {
+        let modelPrefix = "language_model.model."
+        if checkpointKey.hasPrefix(modelPrefix) { return "model." + checkpointKey.dropFirst(modelPrefix.count) }
+        let headPrefix = "language_model.lm_head."
+        if checkpointKey.hasPrefix(headPrefix) { return "lm_head." + checkpointKey.dropFirst(headPrefix.count) }
+        return nil
+    }
+
+    private static func isOffsetNorm(_ key: String) -> Bool {
+        key.hasSuffix(".input_layernorm.weight")
+            || key.hasSuffix(".post_attention_layernorm.weight")
+            || key.hasSuffix(".q_norm.weight")
+            || key.hasSuffix(".k_norm.weight")
+            || key == "model.norm.weight"
+    }
+}