From 67d2bd24668efddcedcdf582cfca645eb06432a9 Mon Sep 17 00:00:00 2001 From: ensan Date: Fri, 1 Dec 2023 14:43:00 +0900 Subject: [PATCH] Fix token_to_piece implementation in Swift --- examples/batched.swift/Sources/main.swift | 10 ++++++++++ .../llama.swiftui/llama.cpp.swift/LibLlama.swift | 13 ++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/examples/batched.swift/Sources/main.swift b/examples/batched.swift/Sources/main.swift index ba15197ae..aa6f7beba 100644 --- a/examples/batched.swift/Sources/main.swift +++ b/examples/batched.swift/Sources/main.swift @@ -242,6 +242,16 @@ private func token_to_piece(token: llama_token, buffer: inout [CChar]) -> String Int32(result.count) ) assert(check == nTokens) + } else if nTokens > 8 { + result.removeAll() + result = [CChar](repeating: 0, count: nTokens) + let check = llama_token_to_piece( + model, + token, + &result, + Int32(result.count) + ) + assert(check == nTokens) } else { result.removeLast(result.count - Int(nTokens)) } diff --git a/examples/llama.swiftui/llama.cpp.swift/LibLlama.swift b/examples/llama.swiftui/llama.cpp.swift/LibLlama.swift index aaef09611..387d3f9c3 100644 --- a/examples/llama.swiftui/llama.cpp.swift/LibLlama.swift +++ b/examples/llama.swiftui/llama.cpp.swift/LibLlama.swift @@ -165,7 +165,18 @@ actor LlamaContext { let result = UnsafeMutablePointer.allocate(capacity: 8) result.initialize(repeating: Int8(0), count: 8) - let _ = llama_token_to_piece(model, token, result, 8) + let nTokens = llama_token_to_piece(model, token, result, 8) + + if nTokens > 8 { + result.removeAll() + result = [CChar](repeating: 0, count: nTokens) + _ = llama_token_to_piece( + model, + token, + &result, + Int32(result.count) + ) + } let resultStr = String(cString: result)