Merge branch 'main' of https://github.com/ml-explore/mlx-swift-examples

* some fixes for gemma2

* format

* fixes

* format

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 - repo: https://github.com/slessans/pre-commit-swift-format
-  rev: ""
+  rev: "fd627de92bdf84a75c924ed95691336d14e94cf1"
   hooks:
     - id: swift-format
       args: ["--configuration", ".swift-format"]

Applications/LLMEval/ContentView.swift

@@ -159,7 +159,7 @@ class LLMEvaluator {
 
     /// this controls which model loads -- phi4bit is one of the smaller ones so this will fit on
     /// more devices
-    let modelConfiguration = ModelConfiguration.phi34bit
+    let modelConfiguration = ModelConfiguration.gemma_2_9b_it_4bit
 
     /// parameters controlling the output
     let generateParameters = GenerateParameters(temperature: 0.6)

Libraries/LLM/Configuration.swift

@@ -32,6 +32,7 @@ public enum ModelType: String, Codable {
     case phi
     case phi3
     case gemma
+    case gemma2
     case qwen2
     case starcoder2
     case cohere
@@ -55,6 +56,10 @@ public enum ModelType: String, Codable {
             let configuration = try JSONDecoder().decode(
                 GemmaConfiguration.self, from: Data(contentsOf: configuration))
             return GemmaModel(configuration)
+        case .gemma2:
+            let configuration = try JSONDecoder().decode(
+                Gemma2Configuration.self, from: Data(contentsOf: configuration))
+            return Gemma2Model(configuration)
         case .qwen2:
             let configuration = try JSONDecoder().decode(
                 Qwen2Configuration.self, from: Data(contentsOf: configuration))

Libraries/LLM/Evaluate.swift

@@ -183,19 +183,10 @@ public func generate(
     var start = Date.timeIntervalSinceReferenceDate
     var promptTime: TimeInterval = 0
 
-    // build a set of additional stop tokens
     let additionalEOSTokenIds = Set(
         (extraEOSTokens ?? [])
-            .map {
+            .compactMap {
-                tokenizer.encode(text: $0)
+                tokenizer.convertTokenToId($0)
-            }
-            .filter {
-                // discard anything that is not a single token.  sometimes
-                // the tokenizer will insert a <s> token, so accept that too
-                $0.count == 1 || ($0.count == 2 && $0[0] == 1)
-            }
-            .map {
-                $0.last!
             })
 
     var tokens = [Int]()

Libraries/LLM/Gemma2.swift

@@ -0,0 +1,309 @@
+// Copyright © 2024 Apple Inc.
+
+import Foundation
+import MLX
+import MLXFast
+import MLXNN
+
+// Port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/gemma2.py
+
+// specialized norm for gemma
+private class RMSNorm: Module, UnaryLayer {
+    let weight: MLXArray
+    let eps: Float
+
+    public init(dimensions: Int, eps: Float = 1e-5) {
+        self.weight = MLXArray.ones([dimensions])
+        self.eps = eps
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        return MLXFast.rmsNorm(x, weight: 1.0 + self.weight, eps: self.eps)
+    }
+}
+
+private class Attention: Module {
+
+    let args: Gemma2Configuration
+    let scale: Float
+    let logitSoftCap: Float
+    let headDim: Int
+
+    @ModuleInfo(key: "q_proj") var wq: Linear
+    @ModuleInfo(key: "k_proj") var wk: Linear
+    @ModuleInfo(key: "v_proj") var wv: Linear
+    @ModuleInfo(key: "o_proj") var wo: Linear
+
+    let rope: RoPE
+
+    public init(_ args: Gemma2Configuration) {
+        self.args = args
+
+        let dim = args.hiddenSize
+        let heads = args.attentionHeads
+        let kvHeads = args.kvHeads
+
+        let headDim = args.headDimensions
+        self.headDim = headDim
+        self.scale = pow(Float(args.queryPreAttnScalar), -0.5)
+        self.logitSoftCap = args.attnLogitSoftcapping
+
+        self._wq.wrappedValue = Linear(dim, heads * headDim, bias: false)
+        self._wk.wrappedValue = Linear(dim, kvHeads * headDim, bias: false)
+        self._wv.wrappedValue = Linear(dim, kvHeads * headDim, bias: false)
+        self._wo.wrappedValue = Linear(heads * headDim, dim, bias: false)
+
+        self.rope = RoPE(
+            dimensions: headDim, traditional: args.ropeTraditional, base: args.ropeTheta)
+    }
+
+    public func callAsFunction(
+        _ x: MLXArray, mask: MLXArray? = nil, cache: (MLXArray, MLXArray)? = nil
+    ) -> (MLXArray, (MLXArray, MLXArray)) {
+        let (B, L) = (x.dim(0), x.dim(1))
+
+        var queries = wq(x)
+        var keys = wk(x)
+        var values = wv(x)
+
+        // prepare the queries, keys and values for the attention computation
+        queries = queries.reshaped(B, L, args.attentionHeads, -1).transposed(0, 2, 1, 3)
+        keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
+        values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
+
+        if let (keyCache, valueCache) = cache {
+            queries = rope(queries, offset: keyCache.dim(2))
+            keys = rope(keys, offset: keyCache.dim(2))
+            keys = concatenated([keyCache, keys], axis: 2)
+            values = concatenated([valueCache, values], axis: 2)
+        } else {
+            queries = rope(queries)
+            keys = rope(keys)
+        }
+
+        let newCache = (keys, values)
+
+        let repeats = self.args.attentionHeads / self.args.kvHeads
+        if repeats > 1 {
+            queries = queries.reshaped(
+                [B, self.args.kvHeads, repeats, L, self.headDim]
+            )
+            keys = expandedDimensions(keys, axes: [2])
+            values = expandedDimensions(values, axes: [2])
+        }
+
+        var scores = matmul(queries, keys.swappedAxes(-1, -2))
+        scores = tanh(scores / self.logitSoftCap) * self.logitSoftCap
+
+        if mask != nil {
+            scores = scores + mask!
+        }
+        scores = softmax(scores, axis: -1, precise: true)
+        var output = matmul(scores, values)
+        if repeats > 1 {
+            output = output.reshaped([B, self.args.attentionHeads, L, self.headDim])
+        }
+        output = output.transposed(0, 2, 1, 3).reshaped(B, L, -1)
+        return (wo(output), newCache)
+    }
+}
+
+private class MLP: Module, UnaryLayer {
+
+    @ModuleInfo(key: "gate_proj") var gate: Linear
+    @ModuleInfo(key: "down_proj") var down: Linear
+    @ModuleInfo(key: "up_proj") var up: Linear
+
+    public init(dimensions: Int, hiddenDimensions: Int) {
+        self._gate.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false)
+        self._down.wrappedValue = Linear(hiddenDimensions, dimensions, bias: false)
+        self._up.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false)
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        down(gelu(gate(x)) * up(x))
+    }
+}
+
+// Minimal changes from Gemma TransformerBlock
+private class TransformerBlock: Module {
+
+    @ModuleInfo(key: "self_attn") var attention: Attention
+    let mlp: MLP
+
+    @ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm
+    @ModuleInfo(key: "pre_feedforward_layernorm") var preFeedforwardLayerNorm: RMSNorm
+    @ModuleInfo(key: "post_feedforward_layernorm") var postFeedforwardLayerNorm: RMSNorm
+    @ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm
+
+    public init(_ args: Gemma2Configuration) {
+        self._attention.wrappedValue = Attention(args)
+        self.mlp = MLP(dimensions: args.hiddenSize, hiddenDimensions: args.intermediateSize)
+        self._inputLayerNorm.wrappedValue = RMSNorm(
+            dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        self._preFeedforwardLayerNorm.wrappedValue = RMSNorm(
+            dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        self._postFeedforwardLayerNorm.wrappedValue = RMSNorm(
+            dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        self._postAttentionLayerNorm.wrappedValue = RMSNorm(
+            dimensions: args.hiddenSize, eps: args.rmsNormEps)
+    }
+
+    public func callAsFunction(
+        _ x: MLXArray, mask: MLXArray? = nil, cache: (MLXArray, MLXArray)? = nil
+    ) -> (MLXArray, (MLXArray, MLXArray)) {
+        var (r, cache) = attention(inputLayerNorm(x), mask: mask, cache: cache)
+        let h = x + postAttentionLayerNorm(r)
+        r = mlp(preFeedforwardLayerNorm(h))
+        let out = h + postFeedforwardLayerNorm(r)
+        return (out, cache)
+    }
+}
+
+// Uses Gemma2TransformerBlock, otherwise same as GemmaModelInner
+public class ModelInner: Module {
+
+    @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
+
+    fileprivate let layers: [TransformerBlock]
+    fileprivate let norm: RMSNorm
+
+    let hiddenScale: Float
+
+    public init(_ args: Gemma2Configuration) {
+        precondition(args.vocabularySize > 0)
+
+        self._embedTokens.wrappedValue = Embedding(
+            embeddingCount: args.vocabularySize, dimensions: args.hiddenSize)
+
+        self.hiddenScale = pow(Float(args.hiddenSize), 0.5)
+
+        self.layers = (0 ..< args.hiddenLayers)
+            .map { _ in
+                TransformerBlock(args)
+            }
+        self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+    }
+
+    public func callAsFunction(_ inputs: MLXArray, cache: [(MLXArray, MLXArray)]? = nil) -> (
+        MLXArray, [(MLXArray, MLXArray)]
+    ) {
+        var h = embedTokens(inputs)
+        h = h * hiddenScale
+
+        var mask: MLXArray? = nil
+        if h.dim(1) > 1 {
+            mask = MultiHeadAttention.createAdditiveCausalMask(h.dim(1))
+            mask = mask?.asType(h.dtype)
+        }
+
+        var newCache = [(MLXArray, MLXArray)]()
+
+        for (i, layer) in layers.enumerated() {
+            var cacheUpdate: (MLXArray, MLXArray)
+            (h, cacheUpdate) = layer(h, mask: mask, cache: cache?[i])
+            newCache.append(cacheUpdate)
+        }
+
+        return (norm(h), newCache)
+    }
+}
+
+// Uses Gemma2ModelInner, otherwise same as GemmaModel
+public class Gemma2Model: Module, LLMModel {
+
+    public let vocabularySize: Int
+    let model: ModelInner
+    let logitSoftCap: Float
+
+    public init(_ args: Gemma2Configuration) {
+        self.vocabularySize = args.vocabularySize
+        self.model = ModelInner(args)
+        self.logitSoftCap = args.finalLogitSoftcapping
+    }
+
+    public func callAsFunction(_ inputs: MLXArray, cache: [(MLXArray, MLXArray)]?) -> (
+        MLXArray, [(MLXArray, MLXArray)]
+    ) {
+        var (out, cache) = model(inputs, cache: cache)
+        out = model.embedTokens.asLinear(out)
+        out = tanh(out / self.logitSoftCap) * self.logitSoftCap
+        return (out, cache)
+    }
+}
+
+public struct Gemma2Configuration: Codable {
+
+    var hiddenSize: Int
+    var hiddenLayers: Int
+    var intermediateSize: Int
+    var attentionHeads: Int
+    var headDimensions: Int
+    var rmsNormEps: Float
+    var vocabularySize: Int
+    var kvHeads: Int
+    var ropeTheta: Float = 10_000
+    var ropeTraditional: Bool = false
+    var attnLogitSoftcapping: Float = 50.0
+    var finalLogitSoftcapping: Float = 30.0
+    var queryPreAttnScalar: Int = 256
+
+    enum CodingKeys: String, CodingKey {
+        case hiddenSize = "hidden_size"
+        case hiddenLayers = "num_hidden_layers"
+        case intermediateSize = "intermediate_size"
+        case attentionHeads = "num_attention_heads"
+        case headDimensions = "head_dim"
+        case rmsNormEps = "rms_norm_eps"
+        case vocabularySize = "vocab_size"
+        case kvHeads = "num_key_value_heads"
+        case ropeTheta = "rope_theta"
+        case ropeTraditional = "rope_traditional"
+        case attnLogitSoftcapping = "attn_logit_softcapping"
+        case finalLogitSoftcapping = "final_logit_softcapping"
+        case queryPreAttnScalar = "query_pre_attn_scalar"
+    }
+
+    public init(from decoder: Decoder) throws {
+        // custom implementation to handle optional keys with required values
+        let container: KeyedDecodingContainer<CodingKeys> = try decoder.container(
+            keyedBy: CodingKeys.self)
+
+        self.hiddenSize = try container.decode(
+            Int.self, forKey: CodingKeys.hiddenSize)
+        self.hiddenLayers = try container.decode(
+            Int.self, forKey: CodingKeys.hiddenLayers)
+        self.intermediateSize = try container.decode(
+            Int.self, forKey: CodingKeys.intermediateSize)
+        self.attentionHeads = try container.decode(
+            Int.self, forKey: CodingKeys.attentionHeads)
+        self.headDimensions = try container.decode(
+            Int.self, forKey: CodingKeys.headDimensions)
+        self.rmsNormEps = try container.decode(
+            Float.self, forKey: CodingKeys.rmsNormEps)
+        self.vocabularySize = try container.decode(
+            Int.self, forKey: CodingKeys.vocabularySize)
+        self.kvHeads = try container.decode(Int.self, forKey: CodingKeys.kvHeads)
+        self.ropeTheta =
+            try container.decodeIfPresent(Float.self, forKey: CodingKeys.ropeTheta)
+            ?? 10_000
+        self.ropeTraditional =
+            try container.decodeIfPresent(
+                Bool.self, forKey: CodingKeys.ropeTraditional) ?? false
+        self.attnLogitSoftcapping = try container.decode(
+            Float.self, forKey: CodingKeys.attnLogitSoftcapping)
+        self.finalLogitSoftcapping = try container.decode(
+            Float.self, forKey: CodingKeys.finalLogitSoftcapping)
+        self.queryPreAttnScalar = try container.decode(
+            Int.self, forKey: CodingKeys.queryPreAttnScalar)
+    }
+}
+
+// MARK: - LoRA
+
+extension Gemma2Model: LoRAModel {
+    public func loraLinearLayers() -> LoRALinearLayers {
+        model.layers.map { ($0.attention, ["q_proj", "v_proj"]) }
+    }
+}

Libraries/LLM/Llama.swift

@@ -7,6 +7,86 @@ import MLXNN
 
 // port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/llama.py
 
+func computeBaseFrequency(
+    base: Float, dims: Int, ropeType: String, ropeScaling: [String: StringOrNumber]?
+)
+    -> Float
+{
+    if ropeType != "llama3" {
+        return base
+    }
+
+    guard let ropeScaling = ropeScaling else {
+        return base
+    }
+
+    guard case .float(let factor) = ropeScaling["factor"],
+        case .float(let lowFreqFactor) = ropeScaling["low_freq_factor"] ?? .float(1.0),
+        case .float(let highFreqFactor) = ropeScaling["high_freq_factor"] ?? .float(4.0),
+        case .float(let oldContextLen) = ropeScaling["original_max_position_embeddings"]
+            ?? .float(8192)
+    else {
+        return base
+    }
+
+    let lowFreqWavelen = oldContextLen / lowFreqFactor
+    let highFreqWavelen = oldContextLen / highFreqFactor
+
+    let freqs = (0 ..< dims).compactMap { index -> Float? in
+        if index % 2 == 0 {
+            return pow(base, Float(index) / Float(dims))
+        }
+        return nil
+    }
+
+    let newBaseFreqs = freqs.map { freq -> Float in
+        let wavelen = 2 * .pi / freq
+        let smooth = max(
+            0, min(1, (wavelen - highFreqWavelen) / (lowFreqWavelen - highFreqWavelen)))
+        return freq * ((1 - smooth) * factor + smooth)
+    }
+
+    return newBaseFreqs.reduce(0, +) / Float(newBaseFreqs.count)
+}
+
+private class DynamicNTKScalingRoPE: Module {
+    let dims: Int
+    let maxPositionEmbeddings: Int?
+    let traditional: Bool
+    let base: Float
+    var scale: Float
+    let ropeType: String
+    let ropeScaling: [String: StringOrNumber]?
+
+    init(
+        dims: Int, maxPositionEmbeddings: Int?, traditional: Bool = false,
+        base: Float = 10000, scale: Float = 1.0, ropeType: String = "default",
+        ropeScaling: [String: StringOrNumber]? = nil
+    ) {
+        self.dims = dims
+        self.maxPositionEmbeddings = maxPositionEmbeddings
+        self.traditional = traditional
+        self.base = computeBaseFrequency(
+            base: base, dims: dims, ropeType: ropeType, ropeScaling: ropeScaling)
+        self.scale = scale
+        self.ropeType = ropeType
+        self.ropeScaling = ropeScaling
+    }
+
+    func callAsFunction(_ x: MLXArray, offset: Int = 0) -> MLXArray {
+        let seqLen = x.dim(1) + offset
+        var base = self.base
+        if let maxPositionEmbeddings, seqLen > maxPositionEmbeddings {
+            let factorAdjustment = Float(seqLen) / Float(maxPositionEmbeddings) - 1
+            let dimensionRatio = Float(dims) / Float(Float(dims) - 2)
+            let adjustedScale = scale * pow(1 + factorAdjustment, dimensionRatio)
+            base *= adjustedScale
+        }
+        return MLXFast.RoPE(
+            x, dimensions: dims, traditional: traditional, base: base, scale: scale, offset: offset)
+    }
+}
+
 private class Attention: Module {
 
     let args: LlamaConfiguration
@@ -17,43 +97,40 @@ private class Attention: Module {
     @ModuleInfo(key: "v_proj") var wv: Linear
     @ModuleInfo(key: "o_proj") var wo: Linear
 
-    let rope: RoPE
+    let rope: DynamicNTKScalingRoPE
 
-    public init(_ args: LlamaConfiguration) {
+    init(_ args: LlamaConfiguration) {
         self.args = args
 
         let dim = args.hiddenSize
         let heads = args.attentionHeads
         let kvHeads = args.kvHeads
 
-        let headDim = args.hiddenSize / heads
+        let headDim = args.headDimensions ?? (args.hiddenSize / heads)
         self.scale = pow(Float(headDim), -0.5)
 
-        self._wq.wrappedValue = Linear(dim, heads * headDim, bias: false)
+        self._wq.wrappedValue = Linear(dim, heads * headDim, bias: args.attentionBias)
-        self._wk.wrappedValue = Linear(dim, kvHeads * headDim, bias: false)
+        self._wk.wrappedValue = Linear(dim, kvHeads * headDim, bias: args.attentionBias)
-        self._wv.wrappedValue = Linear(dim, kvHeads * headDim, bias: false)
+        self._wv.wrappedValue = Linear(dim, kvHeads * headDim, bias: args.attentionBias)
-        self._wo.wrappedValue = Linear(heads * headDim, dim, bias: false)
+        self._wo.wrappedValue = Linear(heads * headDim, dim, bias: args.attentionBias)
 
-        let ropeScale: Float
+        self.rope = DynamicNTKScalingRoPE(
-        if let ropeScaling = args.ropeScaling, ropeScaling["type"] == .string("linear"),
+            dims: headDim,
-            let factor = ropeScaling["factor"]
+            maxPositionEmbeddings: args.maxPositionEmbeddings,
-        {
+            traditional: args.ropeTraditional,
-            switch factor {
+            base: args.ropeTheta,
-            case .string:
+            scale: 1.0,
-                fatalError("ropeScaling.factor must be a float")
+            ropeType: {
-            case .float(let v):
+                if case .string(let value) = args.ropeScaling?["type"] {
-                ropeScale = 1 / v
+                    return value
-            }
+                } else {
-        } else {
+                    return "default"
-            ropeScale = 1
+                }
-        }
+            }(),
-
+            ropeScaling: args.ropeScaling)
-        self.rope = RoPE(
-            dimensions: headDim, traditional: args.ropeTraditional, base: args.ropeTheta,
-            scale: ropeScale)
     }
 
-    public func callAsFunction(
+    func callAsFunction(
         _ x: MLXArray, mask: MLXArray? = nil, cache: (MLXArray, MLXArray)? = nil
     ) -> (MLXArray, (MLXArray, MLXArray)) {
         let (B, L) = (x.dim(0), x.dim(1))
@@ -62,7 +139,7 @@ private class Attention: Module {
         var keys = wk(x)
         var values = wv(x)
 
-        // prepare the queries, keys and values for the attention computation
+        // Prepare the queries, keys and values for the attention computation
         queries = queries.reshaped(B, L, args.attentionHeads, -1).transposed(0, 2, 1, 3)
         keys = keys.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
         values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
@@ -93,35 +170,35 @@ private class MLP: Module, UnaryLayer {
     @ModuleInfo(key: "down_proj") var down: Linear
     @ModuleInfo(key: "up_proj") var up: Linear
 
-    public init(dimensions: Int, hiddenDimensions: Int) {
+    init(_ args: LlamaConfiguration) {
-        self._gate.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false)
+        self._gate.wrappedValue = Linear(args.hiddenSize, args.intermediateSize, bias: args.mlpBias)
-        self._down.wrappedValue = Linear(hiddenDimensions, dimensions, bias: false)
+        self._down.wrappedValue = Linear(args.intermediateSize, args.hiddenSize, bias: args.mlpBias)
-        self._up.wrappedValue = Linear(dimensions, hiddenDimensions, bias: false)
+        self._up.wrappedValue = Linear(args.hiddenSize, args.intermediateSize, bias: args.mlpBias)
     }
 
-    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+    func callAsFunction(_ x: MLXArray) -> MLXArray {
-        down(silu(gate(x)) * up(x))
+        let activation = silu(gate(x))
+        return down(activation * up(x))
     }
 }
 
 private class TransformerBlock: Module {
-
     @ModuleInfo(key: "self_attn") var attention: Attention
-    let mlp: MLP
+    @ModuleInfo(key: "mlp") var mlp: MLP
 
     @ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm
     @ModuleInfo(key: "post_attention_layernorm") var postAttentionLayerNorm: RMSNorm
 
-    public init(_ args: LlamaConfiguration) {
+    init(_ args: LlamaConfiguration) {
         self._attention.wrappedValue = Attention(args)
-        self.mlp = MLP(dimensions: args.hiddenSize, hiddenDimensions: args.intermediateSize)
+        self._mlp.wrappedValue = MLP(args)
         self._inputLayerNorm.wrappedValue = RMSNorm(
             dimensions: args.hiddenSize, eps: args.rmsNormEps)
         self._postAttentionLayerNorm.wrappedValue = RMSNorm(
             dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
-    public func callAsFunction(
+    func callAsFunction(
         _ x: MLXArray, mask: MLXArray? = nil, cache: (MLXArray, MLXArray)? = nil
     ) -> (MLXArray, (MLXArray, MLXArray)) {
         var (r, cache) = attention(inputLayerNorm(x), mask: mask, cache: cache)
@@ -132,27 +209,24 @@ private class TransformerBlock: Module {
     }
 }
 
-public class LlamaModelInner: Module {
+private class LlamaModelInner: Module {
 
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
 
-    fileprivate let layers: [TransformerBlock]
+    let layers: [TransformerBlock]
     let norm: RMSNorm
 
-    public init(_ args: LlamaConfiguration) {
+    init(_ args: LlamaConfiguration) {
         precondition(args.vocabularySize > 0)
 
         self._embedTokens.wrappedValue = Embedding(
             embeddingCount: args.vocabularySize, dimensions: args.hiddenSize)
 
-        self.layers = (0 ..< args.hiddenLayers)
+        self.layers = (0 ..< args.hiddenLayers).map { _ in TransformerBlock(args) }
-            .map { _ in
-                TransformerBlock(args)
-            }
         self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
-    public func callAsFunction(_ inputs: MLXArray, cache: [(MLXArray, MLXArray)]? = nil) -> (
+    func callAsFunction(_ inputs: MLXArray, cache: [(MLXArray, MLXArray)]? = nil) -> (
         MLXArray, [(MLXArray, MLXArray)]
     ) {
         var h = embedTokens(inputs)
@@ -178,25 +252,31 @@ public class LlamaModelInner: Module {
 public class LlamaModel: Module, LLMModel {
 
     public let vocabularySize: Int
-    let model: LlamaModelInner
+    fileprivate let model: LlamaModelInner
 
-    @ModuleInfo(key: "lm_head") var lmHead: Linear
+    @ModuleInfo(key: "lm_head") var lmHead: Linear?
 
     public init(_ args: LlamaConfiguration) {
         self.vocabularySize = args.vocabularySize
         self.model = LlamaModelInner(args)
-        self._lmHead.wrappedValue = Linear(args.hiddenSize, args.vocabularySize, bias: false)
+        if !args.tieWordEmbeddings {
+            self._lmHead.wrappedValue = Linear(args.hiddenSize, args.vocabularySize, bias: false)
+        }
     }
 
     public func callAsFunction(_ inputs: MLXArray, cache: [(MLXArray, MLXArray)]?) -> (
         MLXArray, [(MLXArray, MLXArray)]
     ) {
         let (out, cache) = model(inputs, cache: cache)
-        return (lmHead(out), cache)
+        if let lmHead {
+            return (lmHead(out), cache)
+        } else {
+            return (model.embedTokens.asLinear(out), cache)
+        }
     }
 
     public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
-        // Remove unused precomputed rotary freqs
+        // Remove unused precomputed rotary frequencies
         weights.filter {
             !$0.key.contains("self_attn.rotary_emb.inv_freq")
         }
@@ -209,55 +289,96 @@ public struct LlamaConfiguration: Codable {
     var hiddenLayers: Int
     var intermediateSize: Int
     var attentionHeads: Int
+    var headDimensions: Int?
     var rmsNormEps: Float
     var vocabularySize: Int
     var kvHeads: Int
+    var maxPositionEmbeddings: Int?
     var ropeTheta: Float = 10_000
     var ropeTraditional: Bool = false
-    var ropeScaling: [String: StringOrNumber]? = nil
+    var ropeScaling: [String: StringOrNumber]?
+    var tieWordEmbeddings: Bool = true
+    var attentionBias: Bool = false
+    var mlpBias: Bool = false
 
     enum CodingKeys: String, CodingKey {
         case hiddenSize = "hidden_size"
         case hiddenLayers = "num_hidden_layers"
         case intermediateSize = "intermediate_size"
         case attentionHeads = "num_attention_heads"
+        case headDimensions = "head_dim"
         case rmsNormEps = "rms_norm_eps"
         case vocabularySize = "vocab_size"
         case kvHeads = "num_key_value_heads"
+        case maxPositionEmbeddings = "max_position_embeddings"
         case ropeTheta = "rope_theta"
         case ropeTraditional = "rope_traditional"
         case ropeScaling = "rope_scaling"
+        case tieWordEmbeddings = "tie_word_embeddings"
+        case attentionBias = "attention_bias"
+        case mlpBias = "mlp_bias"
     }
 
     public init(from decoder: Decoder) throws {
-        // custom implementation to handle optional keys with required values
+        let container = try decoder.container(keyedBy: CodingKeys.self)
-        let container: KeyedDecodingContainer<LlamaConfiguration.CodingKeys> =
-            try decoder.container(
-                keyedBy: LlamaConfiguration.CodingKeys.self)
 
-        self.hiddenSize = try container.decode(
+        hiddenSize = try container.decode(Int.self, forKey: .hiddenSize)
-            Int.self, forKey: LlamaConfiguration.CodingKeys.hiddenSize)
+        hiddenLayers = try container.decode(Int.self, forKey: .hiddenLayers)
-        self.hiddenLayers = try container.decode(
+        intermediateSize = try container.decode(Int.self, forKey: .intermediateSize)
-            Int.self, forKey: LlamaConfiguration.CodingKeys.hiddenLayers)
+        attentionHeads = try container.decode(Int.self, forKey: .attentionHeads)
-        self.intermediateSize = try container.decode(
+        headDimensions = try container.decodeIfPresent(Int.self, forKey: .headDimensions)
-            Int.self, forKey: LlamaConfiguration.CodingKeys.intermediateSize)
+        rmsNormEps = try container.decode(Float.self, forKey: .rmsNormEps)
-        self.attentionHeads = try container.decode(
+        vocabularySize = try container.decode(Int.self, forKey: .vocabularySize)
-            Int.self, forKey: LlamaConfiguration.CodingKeys.attentionHeads)
+        kvHeads = try container.decodeIfPresent(Int.self, forKey: .kvHeads) ?? attentionHeads
-        self.rmsNormEps = try container.decode(
+        maxPositionEmbeddings = try container.decodeIfPresent(
-            Float.self, forKey: LlamaConfiguration.CodingKeys.rmsNormEps)
+            Int.self, forKey: .maxPositionEmbeddings)
-        self.vocabularySize = try container.decode(
+        if let ropeTheta = try container.decodeIfPresent(Float.self, forKey: .ropeTheta) {
-            Int.self, forKey: LlamaConfiguration.CodingKeys.vocabularySize)
+            self.ropeTheta = ropeTheta
-        self.kvHeads = try container.decode(Int.self, forKey: LlamaConfiguration.CodingKeys.kvHeads)
+        }
-        self.ropeTheta =
+        if let ropeTraditional = try container.decodeIfPresent(Bool.self, forKey: .ropeTraditional)
-            try container.decodeIfPresent(
+        {
-                Float.self, forKey: LlamaConfiguration.CodingKeys.ropeTheta)
+            self.ropeTraditional = ropeTraditional
-            ?? 10_000
+        }
-        self.ropeTraditional =
+        ropeScaling = try container.decodeIfPresent(
-            try container.decodeIfPresent(
+            [String: StringOrNumber].self, forKey: .ropeScaling)
-                Bool.self, forKey: LlamaConfiguration.CodingKeys.ropeTraditional) ?? false
+        if let tieWordEmbeddings = try container.decodeIfPresent(
-        self.ropeScaling = try container.decodeIfPresent(
+            Bool.self, forKey: .tieWordEmbeddings)
-            [String: StringOrNumber].self, forKey: LlamaConfiguration.CodingKeys.ropeScaling)
+        {
+            self.tieWordEmbeddings = tieWordEmbeddings
+        }
+        if let attentionBias = try container.decodeIfPresent(Bool.self, forKey: .attentionBias) {
+            self.attentionBias = attentionBias
+        }
+        if let mlpBias = try container.decodeIfPresent(Bool.self, forKey: .mlpBias) {
+            self.mlpBias = mlpBias
+        }
 
+        if let ropeScaling {
+            if ropeScaling["factor"] == nil {
+                throw DecodingError.dataCorruptedError(
+                    forKey: .ropeScaling, in: container,
+                    debugDescription: "rope_scaling must contain 'factor'")
+            }
+            if let ropeType = ropeScaling["type"] ?? ropeScaling["rope_type"] {
+                if case .string = ropeType {
+                    let options = [
+                        StringOrNumber.string("linear"), StringOrNumber.string("dynamic"),
+                        StringOrNumber.string("llama3"),
+                    ]
+                    if !options.contains(ropeType) {
+                        throw DecodingError.dataCorruptedError(
+                            forKey: .ropeScaling, in: container,
+                            debugDescription:
+                                "rope_scaling 'type' currently only supports 'linear', 'dynamic', or 'llama3'"
+                        )
+                    }
+                }
+            } else {
+                throw DecodingError.dataCorruptedError(
+                    forKey: .ropeScaling, in: container,
+                    debugDescription: "rope_scaling must contain either 'type' or 'rope_type'")
+            }
+        }
     }
 }
 

Libraries/LLM/Models.swift

@@ -110,13 +110,27 @@ public struct ModelConfiguration {
 }
 
 extension ModelConfiguration {
+    public static let smolLM_135M_4bit = ModelConfiguration(
+        id: "mlx-community/SmolLM-135M-Instruct-4bit",
+        defaultPrompt: "Tell me about the history of Spain."
+    ) {
+        prompt in
+        "<|im_start|>user\n\(prompt)<|im_end|>\n<|im_start|>assistant\n"
+    }
+
+    public static let mistralNeMo4bit = ModelConfiguration(
+        id: "mlx-community/Mistral-Nemo-Instruct-2407-4bit",
+        defaultPrompt: "Explain quaternions."
+    ) { prompt in
+        "<s>[INST] \(prompt) [/INST] "
+    }
 
     public static let mistral7B4bit = ModelConfiguration(
-        id: "mlx-community/Mistral-7B-v0.1-hf-4bit-mlx",
+        id: "mlx-community/Mistral-7B-Instruct-v0.3-4bit",
-
+        defaultPrompt: "Describe the Swift language."
-        // https://www.promptingguide.ai/models/mistral-7b
+    ) { prompt in
-        defaultPrompt: "describe the swift language"
+        "<s>[INST] \(prompt) [/INST] "
-    )
+    }
 
     public static let codeLlama13b4bit = ModelConfiguration(
         id: "mlx-community/CodeLlama-13b-Instruct-hf-4bit-MLX",
@@ -137,7 +151,7 @@ extension ModelConfiguration {
         defaultPrompt: "Why is the sky blue?"
     )
 
-    public static let phi34bit = ModelConfiguration(
+    public static let phi3_4bit = ModelConfiguration(
         id: "mlx-community/Phi-3-mini-4k-instruct-4bit-no-q-embed",
         defaultPrompt: "what is the gravity on mars and the moon?",
         extraEOSTokens: ["<|end|>"]
@@ -153,6 +167,28 @@ extension ModelConfiguration {
         // https://www.promptingguide.ai/models/gemma
         defaultPrompt: "what is the difference between lettuce and cabbage?"
 
+    ) { prompt in
+        "<start_of_turn>user\n\(prompt)<end_of_turn>\n<start_of_turn>model\n"
+    }
+
+    public static let gemma_2_9b_it_4bit = ModelConfiguration(
+        id: "mlx-community/gemma-2-9b-it-4bit",
+        overrideTokenizer: "PreTrainedTokenizer",
+
+        // https://www.promptingguide.ai/models/gemma
+        defaultPrompt: "What is the difference between lettuce and cabbage?"
+
+    ) { prompt in
+        "<start_of_turn>user\n\(prompt)<end_of_turn>\n<start_of_turn>model\n"
+    }
+
+    public static let gemma_2_2b_it_4bit = ModelConfiguration(
+        id: "mlx-community/gemma-2-2b-it-4bit",
+        overrideTokenizer: "PreTrainedTokenizer",
+
+        // https://www.promptingguide.ai/models/gemma
+        defaultPrompt: "What is the difference between lettuce and cabbage?"
+
     ) { prompt in
         "<start_of_turn>user \(prompt)<end_of_turn><start_of_turn>model"
     }
@@ -174,9 +210,17 @@ extension ModelConfiguration {
         "\(prompt)"
     }
 
-    public static let llama38B4bit = ModelConfiguration(
+    public static let llama3_1_8B_4bit = ModelConfiguration(
+        id: "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
+        defaultPrompt: "What is the difference between a fruit and a vegetable?"
+    ) {
+        prompt in
+        "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\nYou are a helpful assistant<|eot_id|>\n<|start_header_id|>user<|end_header_id|>\n\(prompt)<|eot_id|>\n<|start_header_id|>assistant<|end_header_id|>"
+    }
+
+    public static let llama3_8B_4bit = ModelConfiguration(
         id: "mlx-community/Meta-Llama-3-8B-Instruct-4bit",
-        defaultPrompt: "what is the difference between a fruit and a vegetable?"
+        defaultPrompt: "What is the difference between a fruit and a vegetable?"
     ) {
         prompt in
         "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\nYou are a helpful assistant<|eot_id|>\n<|start_header_id|>user<|end_header_id|>\n\(prompt)<|eot_id|>\n<|start_header_id|>assistant<|end_header_id|>"
@@ -195,11 +239,15 @@ extension ModelConfiguration {
         case .idle:
             bootstrapState = .bootstrapping
             register(configurations: [
+                llama3_1_8B_4bit,
+                mistralNeMo4bit,
+                smolLM_135M_4bit,
                 mistral7B4bit,
                 codeLlama13b4bit,
                 phi4bit,
-                phi34bit,
+                phi3_4bit,
                 gemma2bQuantized,
+                gemma_2_9b_it_4bit,
                 qwen205b4bit,
                 openelm270m4bit,
             ])

Package.swift

@@ -15,9 +15,9 @@ let package = Package(
             targets: ["MLXMNIST"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/ml-explore/mlx-swift", branch: "main"),
+        .package(url: "https://github.com/ml-explore/mlx-swift", from: "0.12.1"),
-        .package(url: "https://github.com/huggingface/swift-transformers", from: "0.1.5"),
+        .package(url: "https://github.com/huggingface/swift-transformers", from: "0.1.8"),
-        .package(url: "https://github.com/1024jp/GzipSwift", from: "6.0.1"),
+        .package(url: "https://github.com/1024jp/GzipSwift", "6.0.1" ... "6.0.1"),
         .package(url: "https://github.com/apple/swift-async-algorithms", from: "1.0.0"),
     ],
     targets: [

mlx-swift-examples.xcodeproj/project.pbxproj

@@ -8,6 +8,7 @@
 
 /* Begin PBXBuildFile section */
 		12305EAF2B9D864400C92FEE /* PredictionView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 12305EAE2B9D864400C92FEE /* PredictionView.swift */; };
+		1C55317A2C5AAB4E00B07ECD /* Gemma2.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1C5531792C5AAB4E00B07ECD /* Gemma2.swift */; };
 		1CD79C702BD80DE100B6C06F /* Phi3.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1CD79C6F2BD80DE100B6C06F /* Phi3.swift */; };
 		525C1E9D2B9A011000B5C356 /* Starcoder2.swift in Sources */ = {isa = PBXBuildFile; fileRef = 525C1E9C2B9A010F00B5C356 /* Starcoder2.swift */; };
 		52A776182B94B5EE00AA6E80 /* Qwen2.swift in Sources */ = {isa = PBXBuildFile; fileRef = 52A776172B94B5EE00AA6E80 /* Qwen2.swift */; };
@@ -218,6 +219,7 @@
 
 /* Begin PBXFileReference section */
 		12305EAE2B9D864400C92FEE /* PredictionView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PredictionView.swift; sourceTree = "<group>"; };
+		1C5531792C5AAB4E00B07ECD /* Gemma2.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Gemma2.swift; sourceTree = "<group>"; };
 		1CD79C6F2BD80DE100B6C06F /* Phi3.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Phi3.swift; sourceTree = "<group>"; };
 		525C1E9C2B9A010F00B5C356 /* Starcoder2.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Starcoder2.swift; sourceTree = "<group>"; };
 		52A776172B94B5EE00AA6E80 /* Qwen2.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Qwen2.swift; sourceTree = "<group>"; };
@@ -480,6 +482,7 @@
 				C3A8B3AB2B9283150002EFB8 /* Models.swift */,
 				C34E48EE2B696E6500FCB841 /* Llama.swift */,
 				C38935E22B86C0FE0037B833 /* Gemma.swift */,
+				1C5531792C5AAB4E00B07ECD /* Gemma2.swift */,
 				C38935C72B869C7A0037B833 /* LLM.h */,
 				C38935E02B869F420037B833 /* LLMModel.swift */,
 				C38935DE2B869DD00037B833 /* Phi.swift */,
@@ -1014,6 +1017,7 @@
 				C38935CE2B869C870037B833 /* Load.swift in Sources */,
 				C3E786AD2B8D4AF50004D037 /* Tokenizer.swift in Sources */,
 				C3A8B3AC2B9283150002EFB8 /* Models.swift in Sources */,
+				1C55317A2C5AAB4E00B07ECD /* Gemma2.swift in Sources */,
 				C3E786AB2B8D1AEC0004D037 /* Evaluate.swift in Sources */,
 				C38935CC2B869C870037B833 /* Llama.swift in Sources */,
 				52A776182B94B5EE00AA6E80 /* Qwen2.swift in Sources */,
@@ -1130,7 +1134,7 @@
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = dwarf;
 				DEVELOPMENT_ASSET_PATHS = "\"Applications/LoRATrainingExample/Preview Content\"";
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_PREVIEWS = YES;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
@@ -1168,7 +1172,7 @@
 				MARKETING_VERSION = 1.0;
 				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.LoRATrainingExample;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.LoRATrainingExample;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SDKROOT = auto;
 				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx xros xrsimulator";
@@ -1222,7 +1226,7 @@
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
 				DEVELOPMENT_ASSET_PATHS = "\"Applications/LoRATrainingExample/Preview Content\"";
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_PREVIEWS = YES;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
@@ -1254,7 +1258,7 @@
 				MARKETING_VERSION = 1.0;
 				MTL_ENABLE_DEBUG_INFO = NO;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.LoRATrainingExample;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.LoRATrainingExample;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SDKROOT = auto;
 				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx xros xrsimulator";
@@ -1302,6 +1306,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = dwarf;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -1367,6 +1372,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -1422,12 +1428,14 @@
 				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
 				CLANG_WARN_UNREACHABLE_CODE = YES;
 				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				CODE_SIGN_IDENTITY = "Apple Development";
 				CODE_SIGN_STYLE = Automatic;
 				COMBINE_HIDPI_IMAGES = YES;
 				COPY_PHASE_STRIP = NO;
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = dwarf;
 				DEFINES_MODULE = YES;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				DYLIB_COMPATIBILITY_VERSION = 1;
 				DYLIB_CURRENT_VERSION = 1;
 				DYLIB_INSTALL_NAME_BASE = "@rpath";
@@ -1465,7 +1473,7 @@
 				MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20";
 				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.MNIST;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.MNIST;
 				PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)";
 				SDKROOT = macosx;
 				SKIP_INSTALL = YES;
@@ -1514,12 +1522,14 @@
 				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
 				CLANG_WARN_UNREACHABLE_CODE = YES;
 				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				CODE_SIGN_IDENTITY = "Apple Development";
 				CODE_SIGN_STYLE = Automatic;
 				COMBINE_HIDPI_IMAGES = YES;
 				COPY_PHASE_STRIP = NO;
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
 				DEFINES_MODULE = YES;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				DYLIB_COMPATIBILITY_VERSION = 1;
 				DYLIB_CURRENT_VERSION = 1;
 				DYLIB_INSTALL_NAME_BASE = "@rpath";
@@ -1551,7 +1561,7 @@
 				MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20";
 				MTL_ENABLE_DEBUG_INFO = NO;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.MNIST;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.MNIST;
 				PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)";
 				SDKROOT = macosx;
 				SKIP_INSTALL = YES;
@@ -1602,6 +1612,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = dwarf;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -1667,6 +1678,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -1722,11 +1734,13 @@
 				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
 				CLANG_WARN_UNREACHABLE_CODE = YES;
 				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				CODE_SIGN_IDENTITY = "Apple Development";
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = dwarf;
 				DEFINES_MODULE = YES;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				DYLIB_COMPATIBILITY_VERSION = 1;
 				DYLIB_CURRENT_VERSION = 1;
 				DYLIB_INSTALL_NAME_BASE = "@rpath";
@@ -1767,7 +1781,7 @@
 				MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20";
 				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.LLM;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.LLM;
 				PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)";
 				SDKROOT = auto;
 				SKIP_INSTALL = YES;
@@ -1815,11 +1829,13 @@
 				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
 				CLANG_WARN_UNREACHABLE_CODE = YES;
 				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				CODE_SIGN_IDENTITY = "Apple Development";
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
 				DEFINES_MODULE = YES;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				DYLIB_COMPATIBILITY_VERSION = 1;
 				DYLIB_CURRENT_VERSION = 1;
 				DYLIB_INSTALL_NAME_BASE = "@rpath";
@@ -1854,7 +1870,7 @@
 				MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20";
 				MTL_ENABLE_DEBUG_INFO = NO;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.LLM;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.LLM;
 				PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)";
 				SDKROOT = auto;
 				SKIP_INSTALL = YES;
@@ -1922,6 +1938,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = dwarf;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -1988,6 +2005,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -2046,6 +2064,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = dwarf;
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -2111,6 +2130,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
@@ -2174,6 +2194,7 @@
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = dwarf;
 				DEVELOPMENT_ASSET_PATHS = "\"Applications/MNISTTrainer/Preview Content\"";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_PREVIEWS = YES;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
@@ -2212,7 +2233,7 @@
 				MARKETING_VERSION = 1.0;
 				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.MNISTTrainer;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.MNISTTrainer;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SDKROOT = auto;
 				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
@@ -2265,6 +2286,7 @@
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
 				DEVELOPMENT_ASSET_PATHS = "\"Applications/MNISTTrainer/Preview Content\"";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_PREVIEWS = YES;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
@@ -2297,7 +2319,7 @@
 				MARKETING_VERSION = 1.0;
 				MTL_ENABLE_DEBUG_INFO = NO;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.MNISTTrainer;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.mlx.MNISTTrainer;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SDKROOT = auto;
 				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
@@ -2349,7 +2371,7 @@
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = dwarf;
 				DEVELOPMENT_ASSET_PATHS = "\"Applications/LLMEval/Preview Content\"";
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_PREVIEWS = YES;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
@@ -2379,7 +2401,7 @@
 				"INFOPLIST_KEY_UIStatusBarStyle[sdk=iphonesimulator*]" = UIStatusBarStyleDefault;
 				INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
 				INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
-				IPHONEOS_DEPLOYMENT_TARGET = 17.2;
+				IPHONEOS_DEPLOYMENT_TARGET = 17.0;
 				LD_RUNPATH_SEARCH_PATHS = "@executable_path/Frameworks";
 				"LD_RUNPATH_SEARCH_PATHS[sdk=macosx*]" = "@executable_path/../Frameworks";
 				LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
@@ -2387,7 +2409,7 @@
 				MARKETING_VERSION = 1.0;
 				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.LLMEval;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.LLMEval;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SDKROOT = auto;
 				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx xros xrsimulator";
@@ -2440,7 +2462,7 @@
 				CURRENT_PROJECT_VERSION = 1;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
 				DEVELOPMENT_ASSET_PATHS = "\"Applications/LLMEval/Preview Content\"";
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = GZPQ4W3XF5;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_PREVIEWS = YES;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
@@ -2464,7 +2486,7 @@
 				"INFOPLIST_KEY_UIStatusBarStyle[sdk=iphonesimulator*]" = UIStatusBarStyleDefault;
 				INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
 				INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
-				IPHONEOS_DEPLOYMENT_TARGET = 17.2;
+				IPHONEOS_DEPLOYMENT_TARGET = 17.0;
 				LD_RUNPATH_SEARCH_PATHS = "@executable_path/Frameworks";
 				"LD_RUNPATH_SEARCH_PATHS[sdk=macosx*]" = "@executable_path/../Frameworks";
 				LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
@@ -2472,7 +2494,7 @@
 				MARKETING_VERSION = 1.0;
 				MTL_ENABLE_DEBUG_INFO = NO;
 				MTL_FAST_MATH = YES;
-				PRODUCT_BUNDLE_IDENTIFIER = mlx.LLMEval;
+				PRODUCT_BUNDLE_IDENTIFIER = cn.bewords.LLMEval;
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SDKROOT = auto;
 				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx xros xrsimulator";
@@ -2591,8 +2613,8 @@
 			isa = XCRemoteSwiftPackageReference;
 			repositoryURL = "https://github.com/1024jp/GzipSwift";
 			requirement = {
-				kind = upToNextMajorVersion;
+				kind = exactVersion;
-				minimumVersion = 6.0.1;
+				version = 6.0.1;
 			};
 		};
 		C382DE882B630889000F8F03 /* XCRemoteSwiftPackageReference "swift-async-algorithms" */ = {
@@ -2607,8 +2629,8 @@
 			isa = XCRemoteSwiftPackageReference;
 			repositoryURL = "https://github.com/huggingface/swift-transformers";
 			requirement = {
-				branch = main;
+				kind = upToNextMajorVersion;
-				kind = branch;
+				minimumVersion = 0.1.8;
 			};
 		};
 		C392736E2B60699100368D5D /* XCRemoteSwiftPackageReference "swift-argument-parser" */ = {
@@ -2623,8 +2645,8 @@
 			isa = XCRemoteSwiftPackageReference;
 			repositoryURL = "https://github.com/ml-explore/mlx-swift";
 			requirement = {
-				branch = main;
+				kind = upToNextMajorVersion;
-				kind = branch;
+				minimumVersion = 0.12.1;
 			};
 		};
 /* End XCRemoteSwiftPackageReference section */

mlx-swift-examples.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -15,8 +15,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/ml-explore/mlx-swift",
       "state" : {
-        "branch" : "main",
+        "revision" : "36d63a1fc386a551df14f5b67df1756dc17d2ebc",
-        "revision" : "d6d9472da5bf7ec2654e8914bd1d15622f45b6a9"
+        "version" : "0.12.1"
       }
     },
     {
@@ -78,8 +78,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/huggingface/swift-transformers",
       "state" : {
-        "branch" : "main",
+        "revision" : "fc6543263e4caed9bf6107466d625cfae9357f08",
-        "revision" : "fc6543263e4caed9bf6107466d625cfae9357f08"
+        "version" : "0.1.8"
       }
     }
   ],

mlx-swift-examples.xcodeproj/xcshareddata/xcschemes/LLMEval.xcscheme

@@ -31,8 +31,8 @@
    </TestAction>
    <LaunchAction
       buildConfiguration = "Release"
-      selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
+      selectedDebuggerIdentifier = ""
-      selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
+      selectedLauncherIdentifier = "Xcode.IDEFoundation.Launcher.PosixSpawn"
       launchStyle = "0"
       useCustomWorkingDirectory = "NO"
       ignoresPersistentStateOnLaunch = "NO"

mlx-swift-examples.xcodeproj/xcshareddata/xcschemes/llm-tool.xcscheme

@@ -31,8 +31,8 @@
    </TestAction>
    <LaunchAction
       buildConfiguration = "Release"
-      selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
+      selectedDebuggerIdentifier = ""
-      selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
+      selectedLauncherIdentifier = "Xcode.IDEFoundation.Launcher.PosixSpawn"
       launchStyle = "0"
       useCustomWorkingDirectory = "NO"
       ignoresPersistentStateOnLaunch = "NO"