1import Foundation
2import llama
3
4enum LlamaError: Error {
5 case couldNotInitializeContext
6}
7
8func llama_batch_clear(_ batch: inout llama_batch) {
9 batch.n_tokens = 0
10}
11
12func llama_batch_add(_ batch: inout llama_batch, _ id: llama_token, _ pos: llama_pos, _ seq_ids: [llama_seq_id], _ logits: Bool) {
13 batch.token [Int(batch.n_tokens)] = id
14 batch.pos [Int(batch.n_tokens)] = pos
15 batch.n_seq_id[Int(batch.n_tokens)] = Int32(seq_ids.count)
16 for i in 0..<seq_ids.count {
17 batch.seq_id[Int(batch.n_tokens)]![Int(i)] = seq_ids[i]
18 }
19 batch.logits [Int(batch.n_tokens)] = logits ? 1 : 0
20
21 batch.n_tokens += 1
22}
23
24actor LlamaContext {
25 private var model: OpaquePointer
26 private var context: OpaquePointer
27 private var vocab: OpaquePointer
28 private var sampling: UnsafeMutablePointer<llama_sampler>
29 private var batch: llama_batch
30 private var tokens_list: [llama_token]
31 var is_done: Bool = false
32
33 /// This variable is used to store temporarily invalid cchars
34 private var temporary_invalid_cchars: [CChar]
35
36 var n_len: Int32 = 1024
37 var n_cur: Int32 = 0
38
39 var n_decode: Int32 = 0
40
41 init(model: OpaquePointer, context: OpaquePointer) {
42 self.model = model
43 self.context = context
44 self.tokens_list = []
45 self.batch = llama_batch_init(512, 0, 1)
46 self.temporary_invalid_cchars = []
47 let sparams = llama_sampler_chain_default_params()
48 self.sampling = llama_sampler_chain_init(sparams)
49 llama_sampler_chain_add(self.sampling, llama_sampler_init_temp(0.4))
50 llama_sampler_chain_add(self.sampling, llama_sampler_init_dist(1234))
51 vocab = llama_model_get_vocab(model)
52 }
53
54 deinit {
55 llama_sampler_free(sampling)
56 llama_batch_free(batch)
57 llama_model_free(model)
58 llama_free(context)
59 llama_backend_free()
60 }
61
62 static func create_context(path: String) throws -> LlamaContext {
63 llama_backend_init()
64 var model_params = llama_model_default_params()
65
66#if targetEnvironment(simulator)
67 model_params.n_gpu_layers = 0
68 print("Running on simulator, force use n_gpu_layers = 0")
69#endif
70 let model = llama_model_load_from_file(path, model_params)
71 guard let model else {
72 print("Could not load model at \(path)")
73 throw LlamaError.couldNotInitializeContext
74 }
75
76 let n_threads = max(1, min(8, ProcessInfo.processInfo.processorCount - 2))
77 print("Using \(n_threads) threads")
78
79 var ctx_params = llama_context_default_params()
80 ctx_params.n_ctx = 2048
81 ctx_params.n_threads = Int32(n_threads)
82 ctx_params.n_threads_batch = Int32(n_threads)
83
84 let context = llama_init_from_model(model, ctx_params)
85 guard let context else {
86 print("Could not load context!")
87 throw LlamaError.couldNotInitializeContext
88 }
89
90 return LlamaContext(model: model, context: context)
91 }
92
93 func model_info() -> String {
94 let result = UnsafeMutablePointer<Int8>.allocate(capacity: 256)
95 result.initialize(repeating: Int8(0), count: 256)
96 defer {
97 result.deallocate()
98 }
99
100 // TODO: this is probably very stupid way to get the string from C
101
102 let nChars = llama_model_desc(model, result, 256)
103 let bufferPointer = UnsafeBufferPointer(start: result, count: Int(nChars))
104
105 var SwiftString = ""
106 for char in bufferPointer {
107 SwiftString.append(Character(UnicodeScalar(UInt8(char))))
108 }
109
110 return SwiftString
111 }
112
113 func get_n_tokens() -> Int32 {
114 return batch.n_tokens;
115 }
116
117 func completion_init(text: String) {
118 print("attempting to complete \"\(text)\"")
119
120 tokens_list = tokenize(text: text, add_bos: true)
121 temporary_invalid_cchars = []
122
123 let n_ctx = llama_n_ctx(context)
124 let n_kv_req = tokens_list.count + (Int(n_len) - tokens_list.count)
125
126 print("\n n_len = \(n_len), n_ctx = \(n_ctx), n_kv_req = \(n_kv_req)")
127
128 if n_kv_req > n_ctx {
129 print("error: n_kv_req > n_ctx, the required KV cache size is not big enough")
130 }
131
132 for id in tokens_list {
133 print(String(cString: token_to_piece(token: id) + [0]))
134 }
135
136 llama_batch_clear(&batch)
137
138 for i1 in 0..<tokens_list.count {
139 let i = Int(i1)
140 llama_batch_add(&batch, tokens_list[i], Int32(i), [0], false)
141 }
142 batch.logits[Int(batch.n_tokens) - 1] = 1 // true
143
144 if llama_decode(context, batch) != 0 {
145 print("llama_decode() failed")
146 }
147
148 n_cur = batch.n_tokens
149 }
150
151 func completion_loop() -> String {
152 var new_token_id: llama_token = 0
153
154 new_token_id = llama_sampler_sample(sampling, context, batch.n_tokens - 1)
155
156 if llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len {
157 print("\n")
158 is_done = true
159 let new_token_str = String(cString: temporary_invalid_cchars + [0])
160 temporary_invalid_cchars.removeAll()
161 return new_token_str
162 }
163
164 let new_token_cchars = token_to_piece(token: new_token_id)
165 temporary_invalid_cchars.append(contentsOf: new_token_cchars)
166 let new_token_str: String
167 if let string = String(validatingUTF8: temporary_invalid_cchars + [0]) {
168 temporary_invalid_cchars.removeAll()
169 new_token_str = string
170 } else if (0 ..< temporary_invalid_cchars.count).contains(where: {$0 != 0 && String(validatingUTF8: Array(temporary_invalid_cchars.suffix($0)) + [0]) != nil}) {
171 // in this case, at least the suffix of the temporary_invalid_cchars can be interpreted as UTF8 string
172 let string = String(cString: temporary_invalid_cchars + [0])
173 temporary_invalid_cchars.removeAll()
174 new_token_str = string
175 } else {
176 new_token_str = ""
177 }
178 print(new_token_str)
179 // tokens_list.append(new_token_id)
180
181 llama_batch_clear(&batch)
182 llama_batch_add(&batch, new_token_id, n_cur, [0], true)
183
184 n_decode += 1
185 n_cur += 1
186
187 if llama_decode(context, batch) != 0 {
188 print("failed to evaluate llama!")
189 }
190
191 return new_token_str
192 }
193
194 func bench(pp: Int, tg: Int, pl: Int, nr: Int = 1) -> String {
195 var pp_avg: Double = 0
196 var tg_avg: Double = 0
197
198 var pp_std: Double = 0
199 var tg_std: Double = 0
200
201 for _ in 0..<nr {
202 // bench prompt processing
203
204 llama_batch_clear(&batch)
205
206 let n_tokens = pp
207
208 for i in 0..<n_tokens {
209 llama_batch_add(&batch, 0, Int32(i), [0], false)
210 }
211 batch.logits[Int(batch.n_tokens) - 1] = 1 // true
212
213 llama_memory_clear(llama_get_memory(context), false)
214
215 let t_pp_start = DispatchTime.now().uptimeNanoseconds / 1000;
216
217 if llama_decode(context, batch) != 0 {
218 print("llama_decode() failed during prompt")
219 }
220 llama_synchronize(context)
221
222 let t_pp_end = DispatchTime.now().uptimeNanoseconds / 1000;
223
224 // bench text generation
225
226 llama_memory_clear(llama_get_memory(context), false)
227
228 let t_tg_start = DispatchTime.now().uptimeNanoseconds / 1000;
229
230 for i in 0..<tg {
231 llama_batch_clear(&batch)
232
233 for j in 0..<pl {
234 llama_batch_add(&batch, 0, Int32(i), [Int32(j)], true)
235 }
236
237 if llama_decode(context, batch) != 0 {
238 print("llama_decode() failed during text generation")
239 }
240 llama_synchronize(context)
241 }
242
243 let t_tg_end = DispatchTime.now().uptimeNanoseconds / 1000;
244
245 llama_memory_clear(llama_get_memory(context), false)
246
247 let t_pp = Double(t_pp_end - t_pp_start) / 1000000.0
248 let t_tg = Double(t_tg_end - t_tg_start) / 1000000.0
249
250 let speed_pp = Double(pp) / t_pp
251 let speed_tg = Double(pl*tg) / t_tg
252
253 pp_avg += speed_pp
254 tg_avg += speed_tg
255
256 pp_std += speed_pp * speed_pp
257 tg_std += speed_tg * speed_tg
258
259 print("pp \(speed_pp) t/s, tg \(speed_tg) t/s")
260 }
261
262 pp_avg /= Double(nr)
263 tg_avg /= Double(nr)
264
265 if nr > 1 {
266 pp_std = sqrt(pp_std / Double(nr - 1) - pp_avg * pp_avg * Double(nr) / Double(nr - 1))
267 tg_std = sqrt(tg_std / Double(nr - 1) - tg_avg * tg_avg * Double(nr) / Double(nr - 1))
268 } else {
269 pp_std = 0
270 tg_std = 0
271 }
272
273 let model_desc = model_info();
274 let model_size = String(format: "%.2f GiB", Double(llama_model_size(model)) / 1024.0 / 1024.0 / 1024.0);
275 let model_n_params = String(format: "%.2f B", Double(llama_model_n_params(model)) / 1e9);
276 let backend = "Metal";
277 let pp_avg_str = String(format: "%.2f", pp_avg);
278 let tg_avg_str = String(format: "%.2f", tg_avg);
279 let pp_std_str = String(format: "%.2f", pp_std);
280 let tg_std_str = String(format: "%.2f", tg_std);
281
282 var result = ""
283
284 result += String("| model | size | params | backend | test | t/s |\n")
285 result += String("| --- | --- | --- | --- | --- | --- |\n")
286 result += String("| \(model_desc) | \(model_size) | \(model_n_params) | \(backend) | pp \(pp) | \(pp_avg_str) ± \(pp_std_str) |\n")
287 result += String("| \(model_desc) | \(model_size) | \(model_n_params) | \(backend) | tg \(tg) | \(tg_avg_str) ± \(tg_std_str) |\n")
288
289 return result;
290 }
291
292 func clear() {
293 tokens_list.removeAll()
294 temporary_invalid_cchars.removeAll()
295 llama_memory_clear(llama_get_memory(context), true)
296 }
297
298 private func tokenize(text: String, add_bos: Bool) -> [llama_token] {
299 let utf8Count = text.utf8.count
300 let n_tokens = utf8Count + (add_bos ? 1 : 0) + 1
301 let tokens = UnsafeMutablePointer<llama_token>.allocate(capacity: n_tokens)
302 let tokenCount = llama_tokenize(vocab, text, Int32(utf8Count), tokens, Int32(n_tokens), add_bos, false)
303
304 var swiftTokens: [llama_token] = []
305 for i in 0..<tokenCount {
306 swiftTokens.append(tokens[Int(i)])
307 }
308
309 tokens.deallocate()
310
311 return swiftTokens
312 }
313
314 /// - note: The result does not contain null-terminator
315 private func token_to_piece(token: llama_token) -> [CChar] {
316 let result = UnsafeMutablePointer<Int8>.allocate(capacity: 8)
317 result.initialize(repeating: Int8(0), count: 8)
318 defer {
319 result.deallocate()
320 }
321 let nTokens = llama_token_to_piece(vocab, token, result, 8, 0, false)
322
323 if nTokens < 0 {
324 let newResult = UnsafeMutablePointer<Int8>.allocate(capacity: Int(-nTokens))
325 newResult.initialize(repeating: Int8(0), count: Int(-nTokens))
326 defer {
327 newResult.deallocate()
328 }
329 let nNewTokens = llama_token_to_piece(vocab, token, newResult, -nTokens, 0, false)
330 let bufferPointer = UnsafeBufferPointer(start: newResult, count: Int(nNewTokens))
331 return Array(bufferPointer)
332 } else {
333 let bufferPointer = UnsafeBufferPointer(start: result, count: Int(nTokens))
334 return Array(bufferPointer)
335 }
336 }
337}