llama.go 21.2 KB
Newer Older
1
2
package llama

3
4
//go:generate make -j 8

5
/*
6
7
#cgo CFLAGS: -O3 -std=c17 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE -DGGML_USE_CPU -DGGML_USE_CPU_AARCH64
#cgo CXXFLAGS: -O3 -std=c++17 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE -DGGML_USE_CPU -DGGML_USE_CPU_AARCH64
8
9
#cgo amd64,avx CFLAGS: -mavx
#cgo amd64,avx CXXFLAGS: -mavx
10
11
#cgo amd64,avx2 CFLAGS: -mavx2 -mfma -mf16c
#cgo amd64,avx2 CXXFLAGS: -mavx2 -mfma -mf16c
12
13
14
15
16
17
18
19
#cgo amd64,avx512 CFLAGS: -mavx512f -mavx512dq -mavx512bw
#cgo amd64,avx512 CXXFLAGS: -mavx512f -mavx512dq -mavx512bw
#cgo amd64,avx512bf16 CFLAGS: -mavx512bf16 -D__AVX512BF16__
#cgo amd64,avx512bf16 CXXFLAGS: -mavx512bf16 -D__AVX512BF16__
#cgo amd64,avx512vbmi CFLAGS: -mavx512vbmi -D__AVX512VBMI__
#cgo amd64,avx512vbmi CXXFLAGS: -mavx512vbmi -D__AVX512VBMI__
#cgo amd64,avx512vnni CFLAGS: -mavx512vnni -D__AVX512VNNI__
#cgo amd64,avx512vnni CXXFLAGS: -mavx512vnni -D__AVX512VNNI__
20
21
22
23
24
25
#cgo amd64,f16c CFLAGS: -mf16c
#cgo amd64,f16c CXXFLAGS: -mf16c
#cgo amd64,fma CFLAGS: -mfma
#cgo amd64,fma CXXFLAGS: -mfma
#cgo cuda CFLAGS: -fPIE -DGGML_USE_CUDA -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
#cgo cuda CXXFLAGS: -DGGML_USE_CUDA -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
26
27
28
29
#cgo cuda_jetpack5 LDFLAGS: -lggml_cuda_jetpack5
#cgo cuda_jetpack6 LDFLAGS: -lggml_cuda_jetpack6
#cgo cuda_v11 LDFLAGS: -lggml_cuda_v11
#cgo cuda_v12 LDFLAGS: -lggml_cuda_v12
30
31
32
#cgo darwin,amd64 CFLAGS: -Wno-incompatible-pointer-types-discards-qualifiers
#cgo darwin,amd64 CXXFLAGS: -Wno-incompatible-pointer-types-discards-qualifiers
#cgo darwin,amd64 LDFLAGS: -framework Foundation
33
34
35
#cgo darwin,amd64,avx2 CFLAGS: -DGGML_USE_ACCELERATE -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
#cgo darwin,amd64,avx2 CXXFLAGS: -DGGML_USE_ACCELERATE -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
#cgo darwin,amd64,avx2 LDFLAGS: -framework Accelerate
36
37
#cgo darwin,arm64 CFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS -DGGML_BLAS_USE_ACCELERATE
#cgo darwin,arm64 CXXFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS -DGGML_BLAS_USE_ACCELERATE
38
#cgo darwin,arm64 LDFLAGS: -framework Foundation -framework Metal -framework MetalKit -framework Accelerate
39
40
#cgo linux CFLAGS: -D_GNU_SOURCE
#cgo linux CXXFLAGS: -D_GNU_SOURCE
41
#cgo linux LDFLAGS: -ldl
42
#cgo linux,amd64 LDFLAGS: -L${SRCDIR}/build/linux-amd64
43
44
#cgo linux,arm64 CFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
#cgo linux,arm64 CXXFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
45
#cgo linux,arm64 LDFLAGS: -L${SRCDIR}/build/linux-arm64
46
47
#cgo linux,arm64,sve CFLAGS: -march=armv8.6-a+sve
#cgo linux,arm64,sve CXXFLAGS: -march=armv8.6-a+sve
48
49
50
51
#cgo linux,cuda LDFLAGS: -lcuda -lcudart -lcublas -lcublasLt -lpthread -lrt -lresolv
#cgo linux,rocm LDFLAGS: -lpthread -lrt -lresolv
#cgo rocm CFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIP -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
#cgo rocm CXXFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIP -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
52
#cgo rocm LDFLAGS: -L${SRCDIR} -lggml_rocm -lhipblas -lamdhip64 -lrocblas
53
54
#cgo windows CFLAGS: -Wno-discarded-qualifiers -D_WIN32_WINNT=0x602
#cgo windows CXXFLAGS: -D_WIN32_WINNT=0x602
55
#cgo windows LDFLAGS: -lmsvcrt -static-libstdc++ -static-libgcc -static
56
#cgo windows,amd64 LDFLAGS: -L${SRCDIR}/build/windows-amd64
57
58
#cgo windows,arm64 CFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
#cgo windows,arm64 CXXFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
59
#cgo windows,arm64 LDFLAGS: -L${SRCDIR}/build/windows-arm64
60
61
62
63
64
65
#cgo windows,cuda LDFLAGS: -lcuda -lcudart -lcublas -lcublasLt
#cgo windows,rocm LDFLAGS: -lggml_rocm -lhipblas -lamdhip64 -lrocblas

#include <stdlib.h>
#include "llama.h"
#include "clip.h"
66
#include "ggml.h"
67
#include "llava.h"
68
#include "mllama.h"
69
70
#include "sampling_ext.h"

71
72
extern bool llamaProgressCallback(float progress, void *user_data);
extern void llamaLog(int level, char* text, void* user_data);
73
74
75
76
77
78
79
80
81
82
83

typedef enum {COMP_UNKNOWN,COMP_GCC,COMP_CLANG} COMPILER;
COMPILER inline get_compiler() {
#if defined(__clang__)
	return COMP_CLANG;
#elif defined(__GNUC__)
	return COMP_GCC;
#else
	return UNKNOWN_COMPILER;
#endif
}
84

85
86
87
88
89
90
91
92
93
*/
import "C"

import (
	_ "embed"
	"errors"
	"fmt"
	"runtime"
	"runtime/cgo"
Jesse Gross's avatar
Jesse Gross committed
94
	"slices"
95
	"strings"
96
	"sync/atomic"
97
98
99
100
101
102
103
104
	"unsafe"
)

func BackendInit() {
	C.llama_backend_init()
}

func PrintSystemInfo() string {
105
106
107
108
109
110
111
112
113
114
	var compiler string
	switch C.get_compiler() {
	case C.COMP_UNKNOWN:
		compiler = "cgo(unknown_compiler)"
	case C.COMP_GCC:
		compiler = "cgo(gcc)"
	case C.COMP_CLANG:
		compiler = "cgo(clang)"
	}
	return C.GoString(C.llama_print_system_info()) + compiler
115
116
}

117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
var logLevel atomic.Int32

func init() {
	logLevel.Store(int32(C.GGML_LOG_LEVEL_INFO))
	C.llama_log_set((C.ggml_log_callback)(C.llamaLog), nil)
}

func EnableDebug() {
	logLevel.Store(int32(C.GGML_LOG_LEVEL_DEBUG))
}

//export llamaLog
func llamaLog(level int32, text *C.char, _ unsafe.Pointer) {
	if level < logLevel.Load() {
		return
	}

	fmt.Print(C.GoString(text))
}

137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
func GetModelArch(modelPath string) (string, error) {
	mp := C.CString(modelPath)
	defer C.free(unsafe.Pointer(mp))

	gguf_ctx := C.gguf_init_from_file(mp, C.struct_gguf_init_params{no_alloc: true, ctx: (**C.struct_ggml_context)(C.NULL)})
	if gguf_ctx == nil {
		return "", errors.New("unable to load model file")
	}
	defer C.gguf_free(gguf_ctx)

	key := C.CString("general.architecture")
	defer C.free(unsafe.Pointer(key))
	arch_index := C.gguf_find_key(gguf_ctx, key)
	if int(arch_index) < 0 {
		return "", errors.New("unknown model architecture")
	}

	arch := C.gguf_get_val_str(gguf_ctx, arch_index)

	return C.GoString(arch), nil
}

159
160
161
162
type ContextParams struct {
	c C.struct_llama_context_params
}

163
func NewContextParams(numCtx int, batchSize int, numSeqMax int, threads int, flashAttention bool, kvCacheType string) ContextParams {
164
165
166
167
168
169
170
171
	params := C.llama_context_default_params()
	params.n_ctx = C.uint(numCtx)
	params.n_batch = C.uint(batchSize)
	params.n_seq_max = C.uint(numSeqMax)
	params.n_threads = C.int(threads)
	params.n_threads_batch = params.n_threads
	params.embeddings = C.bool(true)
	params.flash_attn = C.bool(flashAttention)
172
173
174
	params.type_k = kvCacheTypeFromStr(strings.ToLower(kvCacheType))
	params.type_v = kvCacheTypeFromStr(strings.ToLower(kvCacheType))

175
176
177
	return ContextParams{c: params}
}

178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
// kvCacheTypeFromStr converts a string cache type to the corresponding GGML type value
func kvCacheTypeFromStr(s string) C.enum_ggml_type {
	if s == "" {
		return C.GGML_TYPE_F16
	}

	switch s {
	case "q8_0":
		return C.GGML_TYPE_Q8_0
	case "q4_0":
		return C.GGML_TYPE_Q4_0
	default:
		return C.GGML_TYPE_F16
	}
}

194
195
196
197
198
type Context struct {
	c          *C.struct_llama_context
	numThreads int
}

199
var ErrKvCacheFull = errors.New("could not find a kv cache slot")
200
201
202
203
204
205
206
207
208
209
210
211
212

func (c *Context) Decode(batch *Batch) error {
	// Positive return values does not mean a fatal error, but rather a warning.
	//   0 - success
	//   1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context)
	// < 0 - error
	code := int(C.llama_decode(c.c, batch.c))

	if code < 0 {
		return fmt.Errorf("llama_decode failed with code %d", code)
	}

	if code > 0 {
213
		return ErrKvCacheFull
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
	}

	return nil
}

func (c *Context) Model() *Model {
	return &Model{c: C.llama_get_model(c.c)}
}

func (c *Context) KvCacheSeqAdd(seqId int, p0 int, p1 int, delta int) {
	C.llama_kv_cache_seq_add(c.c, C.int(seqId), C.int(p0), C.int(p1), C.int(delta))
}

func (c *Context) KvCacheSeqRm(seqId int, p0 int, p1 int) bool {
	return bool(C.llama_kv_cache_seq_rm(c.c, C.int(seqId), C.int(p0), C.int(p1)))
}

func (c *Context) KvCacheSeqCp(srcSeqId int, dstSeqId int, p0 int, p1 int) {
	C.llama_kv_cache_seq_cp(c.c, C.int(srcSeqId), C.int(dstSeqId), C.int(p0), C.int(p1))
}

235
236
237
238
239
240
241
242
func (c *Context) KvCacheClear() {
	C.llama_kv_cache_clear(c.c)
}

func (c *Context) KvCacheDefrag() {
	C.llama_kv_cache_defrag(c.c)
}

243
244
245
246
247
248
249
250
251
252
253
// Get the embeddings for a sequence id
func (c *Context) GetEmbeddingsSeq(seqId int) []float32 {
	embeddings := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId)))
	if embeddings == nil {
		return nil
	}

	return unsafe.Slice((*float32)(embeddings), c.Model().NEmbd())
}

func (c *Context) GetEmbeddingsIth(i int) []float32 {
254
255
256
257
258
259
	embeddings := unsafe.Pointer(C.llama_get_embeddings_ith(c.c, C.int32_t(i)))
	if embeddings == nil {
		return nil
	}

	return unsafe.Slice((*float32)(embeddings), c.Model().NEmbd())
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
}

type ModelParams struct {
	NumGpuLayers int
	MainGpu      int
	UseMmap      bool
	UseMlock     bool
	TensorSplit  []float32
	Progress     func(float32)
	VocabOnly    bool
}

//export llamaProgressCallback
func llamaProgressCallback(progress C.float, userData unsafe.Pointer) C.bool {
	handle := *(*cgo.Handle)(userData)
	callback := handle.Value().(func(float32))
	callback(float32(progress))
	return true
}

280
func LoadModelFromFile(modelPath string, params ModelParams) (*Model, error) {
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
	cparams := C.llama_model_default_params()
	cparams.n_gpu_layers = C.int(params.NumGpuLayers)
	cparams.main_gpu = C.int32_t(params.MainGpu)
	cparams.use_mmap = C.bool(params.UseMmap)
	cparams.use_mlock = C.bool(params.UseMlock)
	cparams.vocab_only = C.bool(params.VocabOnly)

	if len(params.TensorSplit) > 0 {
		tensorSplitData := &params.TensorSplit[0]

		var tensorSplitPin runtime.Pinner
		tensorSplitPin.Pin(tensorSplitData)
		defer tensorSplitPin.Unpin()

		cparams.tensor_split = (*C.float)(unsafe.Pointer(tensorSplitData))
	}

	if params.Progress != nil {
		handle := cgo.NewHandle(params.Progress)
		defer handle.Delete()

		var handlePin runtime.Pinner
		handlePin.Pin(&handle)
		defer handlePin.Unpin()

		cparams.progress_callback = C.llama_progress_callback(C.llamaProgressCallback)
		cparams.progress_callback_user_data = unsafe.Pointer(&handle)
	}

310
	m := Model{c: C.llama_load_model_from_file(C.CString(modelPath), cparams)}
Jesse Gross's avatar
Jesse Gross committed
311
	if m.c == nil {
312
313
314
315
		return nil, fmt.Errorf("unable to load model: %s", modelPath)
	}

	return &m, nil
316
317
318
319
320
321
}

func FreeModel(model *Model) {
	C.llama_free_model(model.c)
}

322
323
func NewContextWithModel(model *Model, params ContextParams) (*Context, error) {
	c := Context{
324
325
326
		c:          C.llama_new_context_with_model(model.c, params.c),
		numThreads: int(params.c.n_threads),
	}
Jesse Gross's avatar
Jesse Gross committed
327
	if c.c == nil {
328
329
330
331
		return nil, errors.New("unable to create llama context")
	}

	return &c, nil
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
}

func (m *Model) NumVocab() int {
	return int(C.llama_n_vocab(m.c))
}

func (m *Model) TokenIsEog(token int) bool {
	return bool(C.llama_token_is_eog(m.c, C.llama_token(token)))
}

func (m *Model) AddBOSToken() bool {
	return bool(C.llama_add_bos_token(m.c))
}

func (m *Model) ApplyLoraFromFile(context *Context, loraPath string, scale float32, threads int) error {
	cLoraPath := C.CString(loraPath)
	defer C.free(unsafe.Pointer(cLoraPath))

	loraAdapter := C.llama_lora_adapter_init(m.c, cLoraPath)
Jesse Gross's avatar
Jesse Gross committed
351
352
353
	if loraAdapter == nil {
		return errors.New("unable to load lora")
	}
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368

	err := -1
	if loraAdapter != nil {
		err = int(C.llama_lora_adapter_set(context.c, loraAdapter, C.float(scale)))
	}
	if err != 0 {
		return errors.New("error applying lora from file")
	}

	return nil
}

type Batch struct {
	c         C.struct_llama_batch
	batchSize int
369
	maxSeq    int
370
371
372
	embedSize int
}

373
374
375
// Creates a new batch for either word tokens or image embeddings (if embedSize is non-zero).
// Batches cannot contain both types at the same time. batchSize is the maximum number of entries
// that can be added per sequence
Jesse Gross's avatar
Jesse Gross committed
376
377
func NewBatch(batchSize int, maxSeq int, embedSize int) (*Batch, error) {
	b := Batch{
378
379
380
381
		c:         C.llama_batch_init(C.int(batchSize*maxSeq), C.int(embedSize), C.int(maxSeq)),
		batchSize: batchSize,
		maxSeq:    maxSeq,
		embedSize: embedSize,
382
	}
Jesse Gross's avatar
Jesse Gross committed
383
384
385
386
387
388
389
390
391
392
393
394

	// Check to see if any of the allocations in llama_batch_init() failed
	nilPointer := (embedSize == 0 && b.c.token == nil) || (embedSize != 0 && b.c.embd == nil) ||
		b.c.pos == nil || b.c.n_seq_id == nil || b.c.seq_id == nil || b.c.logits == nil ||
		slices.Contains(unsafe.Slice(b.c.seq_id, b.allocSize()), nil)

	if nilPointer {
		C.llama_batch_free(b.c)
		return nil, fmt.Errorf("unable to allocate batch (batchSize=%v maxSeq=%v embedSize=%v)", batchSize, maxSeq, embedSize)
	}

	return &b, nil
395
396
}

397
398
399
400
401
402
403
404
func (b *Batch) Size() int {
	return b.batchSize
}

func (b *Batch) allocSize() int {
	return b.batchSize * b.maxSeq
}

405
406
407
408
409
410
411
412
413
414
415
416
func (b *Batch) NumTokens() int {
	return int(b.c.n_tokens)
}

func (b *Batch) IsEmbedding() bool {
	return b.embedSize != 0
}

// Add adds either a token or an image embedding to the batch depending on the type
// when the batch was initialized. The other argument will be ignored. Adds to the
// batch with the given position for the given sequence ids, and optionally instructs
// to include logits.
417
func (b *Batch) Add(token int, embed []float32, pos int, logits bool, seqIds ...int) {
418
	if !b.IsEmbedding() {
419
		unsafe.Slice(b.c.token, b.allocSize())[b.c.n_tokens] = C.llama_token(token)
420
	} else {
421
		copy(unsafe.Slice((*float32)(b.c.embd), b.allocSize()*b.embedSize)[int(b.c.n_tokens)*b.embedSize:], embed)
422
	}
423
424
	unsafe.Slice(b.c.pos, b.allocSize())[b.c.n_tokens] = C.llama_pos(pos)
	unsafe.Slice(b.c.n_seq_id, b.allocSize())[b.c.n_tokens] = C.int(len(seqIds))
425
426

	for i, s := range seqIds {
427
		unsafe.Slice((unsafe.Slice(b.c.seq_id, b.allocSize())[b.c.n_tokens]), C.int(len(seqIds)))[i] = C.int32_t(s)
428
429
430
	}

	if logits {
431
		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 1
432
433
	} else {
		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 0
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
	}

	b.c.n_tokens += 1
}

func (b *Batch) Clear() {
	b.c.n_tokens = 0
}

func (b *Batch) Free() {
	b.batchSize = 0
	C.llama_batch_free(b.c)
}

type Model struct {
	c *C.struct_llama_model
}

func (m *Model) TokenToPiece(token int) string {
	tokenLen := 12
	buf := make([]byte, tokenLen)
	tokenLen = int(C.llama_token_to_piece(
		m.c,
		C.int32_t(token),
		(*C.char)(unsafe.Pointer(&buf[0])),
		C.int32_t(tokenLen),
		C.int32_t(0),
		C.bool(true),
	))
	if tokenLen < 0 {
		tokenLen = -tokenLen

		buf = make([]byte, tokenLen)
		C.llama_token_to_piece(
			m.c,
			C.int32_t(token),
			(*C.char)(unsafe.Pointer(&buf[0])),
			C.int32_t(tokenLen),
			C.int32_t(0),
			C.bool(true),
		)
	}
	return strings.TrimRight(string(buf), "\x00")
}

func (m *Model) Tokenize(text string, addSpecial bool, parseSpecial bool) ([]int, error) {
	maxTokens := len(text) + 2
	cTokens := make([]C.llama_token, maxTokens)
	cText := C.CString(text)
	defer C.free(unsafe.Pointer(cText))

	result := C.llama_tokenize(
		m.c,
		cText,
		C.int32_t(len(text)),
		&cTokens[0],
		C.int32_t(maxTokens),
		C.bool(addSpecial),
		C.bool(parseSpecial),
	)

	// if the result is negative, reallocate and retry with the correct buffer size
	if result < 0 {
		maxTokens = int(-result)
		cTokens = make([]C.llama_token, maxTokens)
		result = C.llama_tokenize(
			m.c,
			cText,
			C.int32_t(len(text)),
			&cTokens[0],
			C.int32_t(maxTokens),
			C.bool(addSpecial),
			C.bool(parseSpecial),
		)
		if result < 0 {
			return nil, fmt.Errorf("tokenization failed, required %d tokens", -result)
		}
	}

	tokens := make([]int, result)
	for i := range result {
		tokens[i] = int(cTokens[i])
	}

	return tokens, nil
}

func (m *Model) NEmbd() int {
	return int(C.llama_n_embd(m.c))
}

func Quantize(infile, outfile string, ftype uint32) error {
	cinfile := C.CString(infile)
	defer C.free(unsafe.Pointer(cinfile))

	coutfile := C.CString(outfile)
	defer C.free(unsafe.Pointer(coutfile))

	params := C.llama_model_quantize_default_params()
	params.nthread = -1
	params.ftype = ftype

	if rc := C.llama_model_quantize(cinfile, coutfile, &params); rc != 0 {
		return fmt.Errorf("llama_model_quantize: %d", rc)
	}

	return nil
}

543
// vision processing
544
type ClipContext struct {
545
	c *C.struct_clip_ctx
546
547
}

548
func NewClipContext(llamaContext *Context, modelPath string) (*ClipContext, error) {
549
550
	mp := C.CString(modelPath)
	defer C.free(unsafe.Pointer(mp))
551
	c := C.clip_model_load(mp, 1)
Jesse Gross's avatar
Jesse Gross committed
552
553
554
	if c == nil {
		return nil, fmt.Errorf("unable to load clip model: %v", modelPath)
	}
555

556
557
558
559
	projEmbedSize := int(C.clip_n_mmproj_embd(c))
	modelEmbedSize := llamaContext.Model().NEmbd()
	if projEmbedSize != modelEmbedSize {
		return nil, fmt.Errorf("projector embedding size (%d) does not match model (%d)", projEmbedSize, modelEmbedSize)
560
561
	}

562
	return &ClipContext{c: c}, nil
563
564
565
}

func (c *ClipContext) Free() {
566
	C.clip_free(c.c)
567
568
}

Jesse Gross's avatar
Jesse Gross committed
569
func (c *ClipContext) NewEmbed(llamaContext *Context, data []byte) ([][]float32, error) {
570
	l := C.llava_image_embed_make_with_bytes(c.c, C.int(llamaContext.numThreads), (*C.uchar)(unsafe.Pointer(&data[0])), C.int(len(data)))
Jesse Gross's avatar
Jesse Gross committed
571
572
573
	if l == nil {
		return nil, errors.New("unable to make llava embedding from image")
	}
574

575
	numTokens := int(l.n_image_pos)
576
577
	numEmbed := llamaContext.Model().NEmbd()

578
	s := unsafe.Slice((*float32)(l.embed), numEmbed*numTokens)
579
580
581
582
583
584
585
586
587

	embed := make([][]float32, numTokens)
	rows := make([]float32, len(s))
	copy(rows, s)

	for i := range embed {
		embed[i] = rows[i*numEmbed : (i+1)*numEmbed]
	}

588
	C.llava_image_embed_free(l)
589

Jesse Gross's avatar
Jesse Gross committed
590
	return embed, nil
591
592
}

593
594
595
596
597
598
599
600
type MllamaContext struct {
	c *C.struct_mllama_ctx
}

func NewMllamaContext(llamaContext *Context, modelPath string) (*MllamaContext, error) {
	mp := C.CString(modelPath)
	defer C.free(unsafe.Pointer(mp))
	c := C.mllama_model_load(mp, 1)
Jesse Gross's avatar
Jesse Gross committed
601
602
603
	if c == nil {
		return nil, fmt.Errorf("unable to load mllama model: %v", modelPath)
	}
604
605
606
607
608
609
610
611
612
613
614
615
616
617

	projEmbedSize := int(C.mllama_n_embd(c))
	modelEmbedSize := llamaContext.Model().NEmbd()
	if projEmbedSize != modelEmbedSize {
		return nil, fmt.Errorf("projector embedding size (%d) does not match model (%d)", projEmbedSize, modelEmbedSize)
	}

	return &MllamaContext{c: c}, nil
}

func (m *MllamaContext) Free() {
	C.mllama_free(m.c)
}

Jesse Gross's avatar
Jesse Gross committed
618
func (m *MllamaContext) NewEmbed(llamaContext *Context, data []byte, aspectRatioId int) ([][]float32, error) {
619
620
621
	img := C.mllama_image_init()
	defer C.mllama_image_free(img)

Jesse Gross's avatar
Jesse Gross committed
622
623
624
625
	ok := bool(C.mllama_image_load_from_data(unsafe.Pointer(&data[0]), C.int(len(data)), 560, 560, 3, 4, C.int(aspectRatioId), img))
	if !ok {
		return nil, errors.New("unable to load mllama image data")
	}
626

627
	rows := make([]float32, m.EmbedSize(llamaContext))
Jesse Gross's avatar
Jesse Gross committed
628
629
630
631
	ok = bool(C.mllama_image_encode(m.c, C.int(llamaContext.numThreads), img, (*C.float)(unsafe.Pointer(&rows[0]))))
	if !ok {
		return nil, errors.New("unable to make mllama embedding from image")
	}
632

633
634
	embed := make([][]float32, 1)
	embed[0] = rows
635

Jesse Gross's avatar
Jesse Gross committed
636
	return embed, nil
637
638
}

639
640
641
func (m *MllamaContext) EmbedSize(llamaContext *Context) int {
	numTokens := int(C.mllama_n_positions(m.c) * C.mllama_n_tiles(m.c))
	numEmbed := llamaContext.Model().NEmbd()
642

643
644
	return numTokens * numEmbed
}
645

646
647
func (c *Context) SetCrossAttention(state bool) {
	C.llama_set_cross_attention(c.c, C.bool(state))
648
649
}

650
651
652
653
func (c *Context) Synchronize() {
	C.llama_synchronize(c.c)
}

654
655
656
// sampling
// TODO: this is a temporary wrapper to allow calling C++ code from CGo
type SamplingContext struct {
657
	c *C.struct_common_sampler
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
}

type SamplingParams struct {
	TopK           int
	TopP           float32
	MinP           float32
	TypicalP       float32
	Temp           float32
	RepeatLastN    int
	PenaltyRepeat  float32
	PenaltyFreq    float32
	PenaltyPresent float32
	Mirostat       int
	MirostatTau    float32
	MirostatEta    float32
	PenalizeNl     bool
	Seed           uint32
	Grammar        string
}

Jesse Gross's avatar
Jesse Gross committed
678
func NewSamplingContext(model *Model, params SamplingParams) (*SamplingContext, error) {
679
	var cparams C.struct_common_sampler_cparams
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
	cparams.top_k = C.int32_t(params.TopK)
	cparams.top_p = C.float(params.TopP)
	cparams.min_p = C.float(params.MinP)
	cparams.typical_p = C.float(params.TypicalP)
	cparams.temp = C.float(params.Temp)
	cparams.penalty_last_n = C.int32_t(params.RepeatLastN)
	cparams.penalty_repeat = C.float(params.PenaltyRepeat)
	cparams.penalty_freq = C.float(params.PenaltyFreq)
	cparams.penalty_present = C.float(params.PenaltyFreq)
	cparams.mirostat = C.int32_t(params.Mirostat)
	cparams.mirostat_tau = C.float(params.MirostatTau)
	cparams.mirostat_eta = C.float(params.MirostatEta)
	cparams.penalize_nl = C.bool(params.PenalizeNl)
	cparams.seed = C.uint32_t(params.Seed)

	grammar := C.CString(params.Grammar)
	defer C.free(unsafe.Pointer(grammar))

	cparams.grammar = grammar
699
	context := &SamplingContext{c: C.common_sampler_cinit(model.c, &cparams)}
Jesse Gross's avatar
Jesse Gross committed
700
701
702
703
	if context.c == nil {
		return nil, errors.New("unable to create sampling context")
	}

704
	runtime.SetFinalizer(context, func(s *SamplingContext) { C.common_sampler_cfree(s.c) })
705

Jesse Gross's avatar
Jesse Gross committed
706
	return context, nil
707
708
709
}

func (s *SamplingContext) Reset() {
710
	C.common_sampler_creset(s.c)
711
712
}

713
func (s *SamplingContext) Sample(llamaContext *Context, idx int) int {
714
	return int(C.common_sampler_csample(s.c, llamaContext.c, C.int(idx)))
715
716
}

717
func (s *SamplingContext) Accept(id int, applyGrammar bool) {
718
	C.common_sampler_caccept(s.c, C.llama_token(id), C.bool(applyGrammar))
719
}
720

721
722
723
724
// SchemaToGrammar converts the provided JSON schema to a grammar. It returns
// nil if the provided schema is invalid JSON or an invalid JSON schema.
func SchemaToGrammar(schema []byte) []byte {
	cStr := C.CString(string(schema))
725
726
727
728
729
730
731
	defer C.free(unsafe.Pointer(cStr))

	// Allocate buffer for grammar output with reasonable size
	const maxLen = 32768 // 32KB
	buf := make([]byte, maxLen)

	// Call C function to convert schema to grammar
732
733
734
735
	n := C.schema_to_grammar(cStr, (*C.char)(unsafe.Pointer(&buf[0])), C.size_t(maxLen))
	if n == 0 {
		// preserve nil
		return nil
736
	}
737
	return buf[:n]
738
}