Skip to content

Commit b220222

Browse files
authored
swift : fix token_to_piece implementation (#4278)
* Fix token_to_piece implementation in Swift * Fix errors
1 parent 511f52c commit b220222

File tree

2 files changed

+19
-15
lines changed

2 files changed

+19
-15
lines changed

examples/batched.swift/Sources/main.swift

+3-7
Original file line numberDiff line numberDiff line change
@@ -230,18 +230,15 @@ private func token_to_piece(token: llama_token, buffer: inout [CChar]) -> String
230230
var result = [CChar](repeating: 0, count: 8)
231231
let nTokens = llama_token_to_piece(model, token, &result, Int32(result.count))
232232
if nTokens < 0 {
233-
if result.count >= -Int(nTokens) {
234-
result.removeLast(-Int(nTokens))
235-
} else {
236-
result.removeAll()
237-
}
233+
let actualTokensCount = -Int(nTokens)
234+
result = .init(repeating: 0, count: actualTokensCount)
238235
let check = llama_token_to_piece(
239236
model,
240237
token,
241238
&result,
242239
Int32(result.count)
243240
)
244-
assert(check == nTokens)
241+
assert(check == actualTokensCount)
245242
} else {
246243
result.removeLast(result.count - Int(nTokens))
247244
}
@@ -259,5 +256,4 @@ private func token_to_piece(token: llama_token, buffer: inout [CChar]) -> String
259256
buffer = []
260257
return bufferString
261258
}
262-
return nil
263259
}

examples/llama.swiftui/llama.cpp.swift/LibLlama.swift

+16-8
Original file line numberDiff line numberDiff line change
@@ -164,13 +164,21 @@ actor LlamaContext {
164164
private func token_to_piece(token: llama_token) -> String {
165165
let result = UnsafeMutablePointer<Int8>.allocate(capacity: 8)
166166
result.initialize(repeating: Int8(0), count: 8)
167-
168-
let _ = llama_token_to_piece(model, token, result, 8)
169-
170-
let resultStr = String(cString: result)
171-
172-
result.deallocate()
173-
174-
return resultStr
167+
defer {
168+
result.deallocate()
169+
}
170+
let nTokens = llama_token_to_piece(model, token, result, 8)
171+
172+
if nTokens < 0 {
173+
let newResult = UnsafeMutablePointer<Int8>.allocate(capacity: Int(-nTokens))
174+
newResult.initialize(repeating: Int8(0), count: Int(-nTokens))
175+
defer {
176+
newResult.deallocate()
177+
}
178+
_ = llama_token_to_piece(model, token, newResult, -nTokens)
179+
return String(cString: newResult)
180+
} else {
181+
return String(cString: result)
182+
}
175183
}
176184
}

0 commit comments

Comments
 (0)