Skip to content

Commit 50ff67a

Browse files
jolonfjolonf
committed
swift-format
1 parent 6f48543 commit 50ff67a

File tree

3 files changed

+31
-25
lines changed

3 files changed

+31
-25
lines changed

‎Applications/MLXChatExample/Models/PromptCache.swift

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -44,30 +44,30 @@ public class PromptCache: @unchecked Sendable {
4444
/// - If the cache is not trimmable return nil for the caller
4545
/// to create a new cache.
4646
public func getUncachedSuffix(prompt: MLXArray) -> MLXArray? {
47-
47+
4848
print("[getUncachedSuffix] self.tokens.size = \(self.tokens.size)")
49-
49+
5050
print("cache[\(self.tokens.size)]: \(self.tokens)")
5151
print("prompt[\(prompt.size)]: \(prompt)")
52-
52+
5353
let comPrefixLength = commonPrefixLength(newPromptTokens: prompt)
5454
print("[getUncachedSuffix] comPrefixLength: \(comPrefixLength)")
55-
55+
5656
if comPrefixLength == self.tokens.size {
57-
let suffix = prompt[comPrefixLength..<prompt.size]
57+
let suffix = prompt[comPrefixLength ..< prompt.size]
5858
print("Concating...")
5959
self.tokens = concatenated([self.tokens, suffix], axis: 0)
6060
return suffix
61-
} else if (comPrefixLength < self.tokens.size) {
61+
} else if comPrefixLength < self.tokens.size {
6262
if isTrimmable() {
6363
print("trimming: \(self.tokens.size - comPrefixLength)")
6464
let trimmedLen = self.trim(self.tokens.size - comPrefixLength)
6565
print("trimmed: \(trimmedLen)")
6666
if trimmedLen != self.tokens.size - comPrefixLength {
6767
print("Warning: request trimmed amount and actual trimmed amount are different")
6868
}
69-
self.tokens = self.tokens[0..<comPrefixLength]
70-
let suffix = prompt[comPrefixLength..<prompt.size]
69+
self.tokens = self.tokens[0 ..< comPrefixLength]
70+
let suffix = prompt[comPrefixLength ..< prompt.size]
7171
self.tokens = concatenated([self.tokens, suffix], axis: 0)
7272
return suffix
7373
} else {
@@ -81,15 +81,15 @@ public class PromptCache: @unchecked Sendable {
8181

8282
/// - Returns: true if all KV caches are trimmable
8383
public func isTrimmable() -> Bool {
84-
return cache.allSatisfy { $0.isTrimmable()}
84+
return cache.allSatisfy { $0.isTrimmable() }
8585
}
8686

8787
/// Trims all KV caches.
8888
/// - Parameters:
8989
/// - n: Amount to trim.
9090
/// - Returns: Amount KV Caches were trimmed (may be less than ``n``).
9191
public func trim(_ n: Int) -> Int {
92-
if !self.isTrimmable(){
92+
if !self.isTrimmable() {
9393
return 0
9494
}
9595
return cache.map { $0.trim(n: n) }.max() ?? 0
@@ -103,7 +103,7 @@ public class PromptCache: @unchecked Sendable {
103103
public func commonPrefixLength(newPromptTokens: MLXArray) -> Int {
104104
return commonPrefixLength(self.tokens, newPromptTokens)
105105
}
106-
106+
107107
/// Finds the common prefix between ``MLXArray``s.
108108
/// - Parameters:
109109
/// - array1: First array
@@ -113,7 +113,7 @@ public class PromptCache: @unchecked Sendable {
113113
// TODO: Add test cases
114114
print("Calculating common prefix: array1[\(array1.size)] array2[\(array2.size)]")
115115
let minLength = min(array1.size, array2.size)
116-
for i in 0..<minLength {
116+
for i in 0 ..< minLength {
117117
if all(array1[i] .!= array2[i]).item(Bool.self) {
118118
return i
119119
}

‎Applications/MLXChatExample/Services/MLXService.swift

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ class MLXService {
1919
/// Includes both language models (LLM) and vision-language models (VLM).
2020
static let availableModels: [LMModel] = [
2121
LMModel(name: "llama3.2:1b", configuration: LLMRegistry.llama3_2_1B_4bit, type: .llm),
22+
LMModel(name: "llama3.2:3b", configuration: LLMRegistry.llama3_2_3B_4bit, type: .llm),
2223
LMModel(name: "qwen2.5:1.5b", configuration: LLMRegistry.qwen2_5_1_5b, type: .llm),
2324
LMModel(name: "smolLM:135m", configuration: LLMRegistry.smolLM_135M_4bit, type: .llm),
2425
LMModel(name: "qwen3:0.6b", configuration: LLMRegistry.qwen3_0_6b_4bit, type: .llm),
@@ -72,13 +73,13 @@ class MLXService {
7273
self.modelDownloadProgress = progress
7374
}
7475
}
75-
76+
7677
// Clear out the promptCache
7778
promptCache.removeObject(forKey: model.name as NSString)
78-
79+
7980
// Cache the loaded model for future use
8081
modelCache.setObject(container, forKey: model.name as NSString)
81-
82+
8283
return container
8384
}
8485
}
@@ -127,15 +128,20 @@ class MLXService {
127128
// Get the prompt cache and adjust new prompt to remove
128129
// prefix already in cache, trim cache if cache is
129130
// inconsistent with new prompt.
130-
let (cache, lmInput) = getPromptCache(fullPrompt: fullPrompt, parameters: parameters, context: context, modelName: model.name)
131-
131+
let (cache, lmInput) = getPromptCache(
132+
fullPrompt: fullPrompt, parameters: parameters, context: context,
133+
modelName: model.name)
134+
132135
// TODO: The generated tokens should be added to the prompt cache but not possible with AsyncStream
133136
return try MLXLMCommon.generate(
134137
input: lmInput, parameters: parameters, context: context, cache: cache.cache)
135138
}
136139
}
137-
138-
func getPromptCache(fullPrompt: LMInput, parameters: GenerateParameters, context: ModelContext, modelName: String) -> (PromptCache, LMInput) {
140+
141+
func getPromptCache(
142+
fullPrompt: LMInput, parameters: GenerateParameters, context: ModelContext,
143+
modelName: String
144+
) -> (PromptCache, LMInput) {
139145
let cache: PromptCache
140146
if let existingCache = promptCache.object(forKey: modelName as NSString) {
141147
cache = existingCache
@@ -146,7 +152,7 @@ class MLXService {
146152
}
147153

148154
let lmInput: LMInput
149-
155+
150156
/// Remove prefix from prompt that is already in cache
151157
if let suffix = cache.getUncachedSuffix(prompt: fullPrompt.text.tokens) {
152158
lmInput = LMInput(text: LMInput.Text(tokens: suffix))
@@ -157,7 +163,7 @@ class MLXService {
157163
self.promptCache.setObject(newCache, forKey: modelName as NSString)
158164
lmInput = fullPrompt
159165
}
160-
166+
161167
return (cache, lmInput)
162168
}
163169
}

‎Libraries/MLXLMCommon/KVCache.swift

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ public protocol KVCache: Evaluatable {
1212
var offset: Int { get }
1313

1414
func update(keys: MLXArray, values: MLXArray) -> (MLXArray, MLXArray)
15-
15+
1616
func isTrimmable() -> Bool
17-
17+
1818
func trim(n: Int) -> Int
1919
}
2020

@@ -100,11 +100,11 @@ public class KVCacheSimple: KVCache, Evaluatable {
100100
self.values![.ellipsis, ..<self.offset, 0...]
101101
)
102102
}
103-
103+
104104
public func isTrimmable() -> Bool {
105105
return true
106106
}
107-
107+
108108
public func trim(n: Int) -> Int {
109109
let toTrim = min(self.offset, n)
110110
self.offset -= toTrim

0 commit comments

Comments
 (0)