Skip to content

Commit 9a6d9a3

Browse files
davidkoskidavidkoski
committed
use ReerCodable macro to allow for default values
- make all configuration Sendable and public var as well
1 parent 51f51cf commit 9a6d9a3

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

55 files changed

+1140
-2744
lines changed

‎.circleci/config.yml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,17 +38,17 @@ jobs:
3838
xcrun --show-sdk-build-version
3939
swift --version
4040
find . -name Package.resolved -exec rm {} \;
41-
xcodebuild test -scheme mlx-libraries-Package -destination 'platform=OS X'
41+
xcodebuild test -scheme mlx-libraries-Package -destination 'platform=OS X' -skipMacroValidation
4242
- run:
4343
name: Build Examples
4444
command: |
4545
xcodebuild -version
4646
xcrun --show-sdk-build-version
4747
swift --version
4848
find . -name Package.resolved -exec rm {} \;
49-
xcodebuild -scheme llm-tool
50-
xcodebuild -scheme image-tool
51-
xcodebuild -scheme mnist-tool
49+
xcodebuild -scheme llm-tool -skipMacroValidation
50+
xcodebuild -scheme image-tool -skipMacroValidation
51+
xcodebuild -scheme mnist-tool -skipMacroValidation
5252
5353
workflows:
5454
build_and_test:

‎Libraries/Embedders/Pooling.swift‎

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,15 @@ import Foundation
44
import MLX
55
import MLXLinalg
66
import MLXNN
7+
import ReerCodable
78

8-
public struct PoolingConfiguration: Codable {
9-
public let dimension: Int
10-
public let poolingModeClsToken: Bool
11-
public let poolingModeMeanTokens: Bool
12-
public let poolingModeMaxTokens: Bool
13-
public let poolingModeLastToken: Bool
14-
15-
enum CodingKeys: String, CodingKey {
16-
case dimension = "word_embedding_dimension"
17-
case poolingModeClsToken = "pooling_mode_cls_token"
18-
case poolingModeMeanTokens = "pooling_mode_mean_tokens"
19-
case poolingModeMaxTokens = "pooling_mode_max_tokens"
20-
case poolingModeLastToken = "pooling_mode_lasttoken"
21-
}
9+
@Codable
10+
public struct PoolingConfiguration: Sendable {
11+
@CodingKey("word_embedding_dimension") public let dimension: Int
12+
@CodingKey("pooling_mode_cls_token") public let poolingModeClsToken: Bool
13+
@CodingKey("pooling_mode_mean_tokens") public let poolingModeMeanTokens: Bool
14+
@CodingKey("pooling_mode_max_tokens") public let poolingModeMaxTokens: Bool
15+
@CodingKey("pooling_mode_lasttoken") public let poolingModeLastToken: Bool
2216
}
2317

2418
func loadPooling(modelDirectory: URL) -> Pooling {

‎Libraries/MLXLLM/Codable+Support.swift‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
import Foundation
2+
3+
/// `swift-transformers` also declares a public `Decoder` and it conflicts with the `Codable`
4+
/// implementations.
5+
public typealias Decoder = Swift.Decoder

‎Libraries/MLXLLM/Documentation.docc/adding-model.md‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,17 +14,12 @@ and create a `.swift` file for your new model:
1414
Create a configuration struct to match the `config.json` (any parameters needed).
1515

1616
```swift
17-
public struct YourModelConfiguration: Codable, Sendable {
18-
public let hiddenSize: Int
19-
20-
// use this pattern for values that need defaults
21-
public let _layerNormEps: Float?
22-
public var layerNormEps: Float { _layerNormEps ?? 1e-6 }
23-
24-
enum CodingKeys: String, CodingKey {
25-
case hiddenSize = "hidden_size"
26-
case _layerNormEps = "layer_norm_eps"
27-
}
17+
import ReerCodable
18+
19+
@Codable
20+
public struct YourModelConfiguration: Sendable {
21+
@CodingKey("hidden_size") public var hiddenSize: Int
22+
@CodingKey("layer_norm_eps") public var layerNormEps: Float = 1e-6
2823
}
2924
```
3025

‎Libraries/MLXLLM/LLMModelFactory.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,9 @@ public class LLMTypeRegistry: ModelTypeRegistry, @unchecked Sendable {
3535
"phimoe": create(PhiMoEConfiguration.self, PhiMoEModel.init),
3636
"gemma": create(GemmaConfiguration.self, GemmaModel.init),
3737
"gemma2": create(Gemma2Configuration.self, Gemma2Model.init),
38-
"gemma3": create(Gemma3TextConfiguration.self, Gemma3TextModel.init),
39-
"gemma3_text": create(Gemma3TextConfiguration.self, Gemma3TextModel.init),
40-
"gemma3n": create(Gemma3nTextConfiguration.self, Gemma3nTextModel.init),
38+
"gemma3": create(Gemma3TextConfigurationContainer.self, Gemma3TextModel.init),
39+
"gemma3_text": create(Gemma3TextConfigurationContainer.self, Gemma3TextModel.init),
40+
"gemma3n": create(Gemma3nTextConfigurationContainer.self, Gemma3nTextModel.init),
4141
"qwen2": create(Qwen2Configuration.self, Qwen2Model.init),
4242
"qwen3": create(Qwen3Configuration.self, Qwen3Model.init),
4343
"qwen3_moe": create(Qwen3MoEConfiguration.self, Qwen3MoEModel.init),

‎Libraries/MLXLLM/Lora+Data.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public func loadLoRAData(url: URL) throws -> [String] {
4848

4949
func loadJSONL(url: URL) throws -> [String] {
5050

51-
struct Line: Codable {
51+
struct Line: Codable, Sendable {
5252
let text: String?
5353
}
5454

‎Libraries/MLXLLM/Models/BaichuanM1.swift‎

Lines changed: 18 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -11,39 +11,24 @@ import MLXFast
1111
import MLXLMCommon
1212
import MLXNN
1313
import MLXRandom
14-
15-
public struct BaichuanM1Configuration: Codable, Sendable {
16-
var vocabularySize: Int
17-
var hiddenSize: Int
18-
var intermediateSize: Int
19-
var hiddenLayers: Int
20-
var attentionHeads: Int
21-
var kvHeads: Int
22-
var ropeTheta: Float
23-
var slidingWindow: Int
24-
var slidingWindowLayers: [Int]
25-
var convWindow: Int
26-
var rmsNormEps: Float
27-
var swaAttentionHeads: Int?
28-
var swaKvHeads: Int?
29-
var tieWordEmbeddings: Bool = false
30-
31-
enum CodingKeys: String, CodingKey {
32-
case vocabularySize = "vocab_size"
33-
case hiddenSize = "hidden_size"
34-
case intermediateSize = "intermediate_size"
35-
case hiddenLayers = "num_hidden_layers"
36-
case attentionHeads = "num_attention_heads"
37-
case kvHeads = "num_key_value_heads"
38-
case ropeTheta = "rope_theta"
39-
case slidingWindow = "sliding_window"
40-
case slidingWindowLayers = "sliding_window_layers"
41-
case convWindow = "conv_window"
42-
case rmsNormEps = "rms_norm_eps"
43-
case swaAttentionHeads = "num_swa_attention_heads"
44-
case swaKvHeads = "num_swa_key_value_heads"
45-
case tieWordEmbeddings = "tie_word_embeddings"
46-
}
14+
import ReerCodable
15+
16+
@Codable
17+
public struct BaichuanM1Configuration: Sendable {
18+
@CodingKey("vocab_size") public var vocabularySize: Int
19+
@CodingKey("hidden_size") public var hiddenSize: Int
20+
@CodingKey("intermediate_size") public var intermediateSize: Int
21+
@CodingKey("num_hidden_layers") public var hiddenLayers: Int
22+
@CodingKey("num_attention_heads") public var attentionHeads: Int
23+
@CodingKey("num_key_value_heads") public var kvHeads: Int
24+
@CodingKey("rope_theta") public var ropeTheta: Float
25+
@CodingKey("sliding_window") public var slidingWindow: Int
26+
@CodingKey("sliding_window_layers") public var slidingWindowLayers: [Int]
27+
@CodingKey("conv_window") public var convWindow: Int
28+
@CodingKey("rms_norm_eps") public var rmsNormEps: Float
29+
@CodingKey("num_swa_attention_heads") public var swaAttentionHeads: Int?
30+
@CodingKey("num_swa_key_value_heads") public var swaKvHeads: Int?
31+
@CodingKey("tie_word_embeddings") public var tieWordEmbeddings: Bool = false
4732
}
4833

4934
private class Attention: Module {

‎Libraries/MLXLLM/Models/BailingMoe.swift‎

Lines changed: 32 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -10,69 +10,40 @@ import Foundation
1010
import MLX
1111
import MLXLMCommon
1212
import MLXNN
13-
14-
public struct BailingMoeConfiguration: Codable, Sendable {
15-
var modelType: String
16-
var hiddenSize: Int
17-
var intermediateSize: Int
18-
var maxPositionEmbeddings: Int?
19-
var moeIntermediateSize: Int
20-
var numExperts: Int
21-
var numSharedExperts: Int
22-
var normTopkProb: Bool
23-
var attentionHeads: Int
24-
var numExpertsPerToken: Int
25-
var hiddenLayers: Int
26-
var kvHeads: Int
27-
var rmsNormEps: Float
28-
var ropeTheta: Float
29-
var vocabularySize: Int
30-
var firstKDenseReplace: Int
13+
import ReerCodable
14+
15+
@Codable
16+
public struct BailingMoeConfiguration: Sendable {
17+
@CodingKey("model_type") public var modelType: String
18+
@CodingKey("hidden_size") public var hiddenSize: Int
19+
@CodingKey("intermediate_size") public var intermediateSize: Int
20+
@CodingKey("max_position_embeddings") public var maxPositionEmbeddings: Int?
21+
@CodingKey("moe_intermediate_size") public var moeIntermediateSize: Int
22+
@CodingKey("num_experts") public var numExperts: Int
23+
@CodingKey("num_shared_experts") public var numSharedExperts: Int
24+
@CodingKey("norm_topk_prob") public var normTopkProb: Bool
25+
@CodingKey("num_attention_heads") public var attentionHeads: Int
26+
@CodingKey("num_experts_per_tok") public var numExpertsPerToken: Int
27+
@CodingKey("num_hidden_layers") public var hiddenLayers: Int
28+
@CodingKey("num_key_value_heads") public var kvHeads: Int
29+
@CodingKey("rms_norm_eps") public var rmsNormEps: Float
30+
@CodingKey("rope_theta") public var ropeTheta: Float
31+
@CodingKey("vocab_size") public var vocabularySize: Int
32+
@CodingKey("first_k_dense_replace") public var firstKDenseReplace: Int
3133

3234
// Optional features
33-
var ropeScaling: [String: StringOrNumber]? = nil
34-
var useBias: Bool = false
35-
var useQKVBias: Bool = false
36-
var useQKNorm: Bool = false
37-
var tieWordEmbeddings: Bool = false
38-
var partialRotaryFactor: Float = 1.0
39-
var moeRouterEnableExpertBias: Bool = false
40-
var routedScalingFactor: Float = 1.0
41-
var scoreFunction: String = "softmax"
42-
var nGroup: Int = 1
43-
var topkGroup: Int = 4
44-
var moeSharedExpertIntermediateSize: Int? = nil
45-
46-
enum CodingKeys: String, CodingKey {
47-
case modelType = "model_type"
48-
case hiddenSize = "hidden_size"
49-
case intermediateSize = "intermediate_size"
50-
case maxPositionEmbeddings = "max_position_embeddings"
51-
case moeIntermediateSize = "moe_intermediate_size"
52-
case numExperts = "num_experts"
53-
case numSharedExperts = "num_shared_experts"
54-
case normTopkProb = "norm_topk_prob"
55-
case attentionHeads = "num_attention_heads"
56-
case numExpertsPerToken = "num_experts_per_tok"
57-
case hiddenLayers = "num_hidden_layers"
58-
case kvHeads = "num_key_value_heads"
59-
case rmsNormEps = "rms_norm_eps"
60-
case ropeTheta = "rope_theta"
61-
case vocabularySize = "vocab_size"
62-
case firstKDenseReplace = "first_k_dense_replace"
63-
case ropeScaling = "rope_scaling"
64-
case useBias = "use_bias"
65-
case useQKVBias = "use_qkv_bias"
66-
case useQKNorm = "use_qk_norm"
67-
case tieWordEmbeddings = "tie_word_embeddings"
68-
case partialRotaryFactor = "partial_rotary_factor"
69-
case moeRouterEnableExpertBias = "moe_router_enable_expert_bias"
70-
case routedScalingFactor = "routed_scaling_factor"
71-
case scoreFunction = "score_function"
72-
case nGroup = "n_group"
73-
case topkGroup = "topk_group"
74-
case moeSharedExpertIntermediateSize = "moe_shared_expert_intermediate_size"
75-
}
35+
@CodingKey("rope_scaling") public var ropeScaling: [String: StringOrNumber]? = nil
36+
@CodingKey("use_bias") public var useBias: Bool = false
37+
@CodingKey("use_qkv_bias") public var useQKVBias: Bool = false
38+
@CodingKey("use_qk_norm") public var useQKNorm: Bool = false
39+
@CodingKey("tie_word_embeddings") public var tieWordEmbeddings: Bool = false
40+
@CodingKey("partial_rotary_factor") public var partialRotaryFactor: Float = 1.0
41+
@CodingKey("moe_router_enable_expert_bias") public var moeRouterEnableExpertBias: Bool = false
42+
@CodingKey("routed_scaling_factor") public var routedScalingFactor: Float = 1.0
43+
@CodingKey("score_function") public var scoreFunction: String = "softmax"
44+
@CodingKey("n_group") public var nGroup: Int = 1
45+
@CodingKey("topk_group") public var topkGroup: Int = 4
46+
@CodingKey("moe_shared_expert_intermediate_size") public var moeSharedExpertIntermediateSize: Int? = nil
7647
}
7748

7849
private class Attention: Module {

0 commit comments

Comments
 (0)