llama.go 17.8 KB
Newer Older
1
2
3
package llama

/*
Michael Yang's avatar
Michael Yang committed
4
5
#cgo CFLAGS: -std=c11
#cgo CXXFLAGS: -std=c++17
Michael Yang's avatar
Michael Yang committed
6
#cgo CPPFLAGS: -O3
Michael Yang's avatar
Michael Yang committed
7
8
9
10
11
#cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/include
#cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/common
#cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/examples/llava
#cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/src
#cgo CPPFLAGS: -I${SRCDIR}/../ml/backend/ggml/ggml/include
12
13

#include <stdlib.h>
Michael Yang's avatar
Michael Yang committed
14
#include "ggml.h"
15
16
17
#include "llama.h"
#include "clip.h"
#include "llava.h"
Michael Yang's avatar
Michael Yang committed
18

19
#include "mllama.h"
20
21
#include "sampling_ext.h"

22
23
extern bool llamaProgressCallback(float progress, void *user_data);
extern void llamaLog(int level, char* text, void* user_data);
24
25
26
27
28
29
30
31
32
33
34

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
}
35

36
37
38
39
40
41
42
*/
import "C"

import (
	_ "embed"
	"errors"
	"fmt"
43
	"os"
44
45
	"runtime"
	"runtime/cgo"
Jesse Gross's avatar
Jesse Gross committed
46
	"slices"
47
	"strings"
48
	"sync/atomic"
49
	"unsafe"
Michael Yang's avatar
Michael Yang committed
50
51
52
53
54

	_ "github.com/ollama/ollama/llama/llama.cpp/common"
	_ "github.com/ollama/ollama/llama/llama.cpp/examples/llava"
	_ "github.com/ollama/ollama/llama/llama.cpp/src"
	"github.com/ollama/ollama/ml/backend/ggml/ggml/src"
55
56
57
)

func BackendInit() {
Michael Yang's avatar
Michael Yang committed
58
	ggml.OnceLoad()
59
60
61
62
	C.llama_backend_init()
}

func PrintSystemInfo() string {
63
64
65
66
67
68
69
70
71
72
	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
73
74
}

75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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
	}

92
	fmt.Fprint(os.Stderr, C.GoString(text))
93
94
}

95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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
}

117
118
119
120
type ContextParams struct {
	c C.struct_llama_context_params
}

121
func NewContextParams(numCtx int, batchSize int, numSeqMax int, threads int, flashAttention bool, kvCacheType string) ContextParams {
122
123
124
125
126
127
128
129
	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)
130
131
132
	params.type_k = kvCacheTypeFromStr(strings.ToLower(kvCacheType))
	params.type_v = kvCacheTypeFromStr(strings.ToLower(kvCacheType))

133
134
135
	return ContextParams{c: params}
}

136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
// 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
	}
}

152
153
154
155
156
type Context struct {
	c          *C.struct_llama_context
	numThreads int
}

157
var ErrKvCacheFull = errors.New("could not find a kv cache slot")
158
159
160
161
162
163
164
165
166
167
168
169
170

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 {
171
		return ErrKvCacheFull
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
	}

	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))
}

193
194
195
196
197
198
199
200
func (c *Context) KvCacheClear() {
	C.llama_kv_cache_clear(c.c)
}

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

201
202
203
204
205
206
207
208
209
210
211
// 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 {
212
213
214
215
216
217
	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())
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
}

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
}

238
func LoadModelFromFile(modelPath string, params ModelParams) (*Model, error) {
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
	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)
	}

268
	m := Model{c: C.llama_load_model_from_file(C.CString(modelPath), cparams)}
Jesse Gross's avatar
Jesse Gross committed
269
	if m.c == nil {
270
271
272
273
		return nil, fmt.Errorf("unable to load model: %s", modelPath)
	}

	return &m, nil
274
275
276
277
278
279
}

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

280
281
func NewContextWithModel(model *Model, params ContextParams) (*Context, error) {
	c := Context{
282
283
284
		c:          C.llama_new_context_with_model(model.c, params.c),
		numThreads: int(params.c.n_threads),
	}
Jesse Gross's avatar
Jesse Gross committed
285
	if c.c == nil {
286
287
288
289
		return nil, errors.New("unable to create llama context")
	}

	return &c, nil
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
}

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
309
310
311
	if loraAdapter == nil {
		return errors.New("unable to load lora")
	}
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326

	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
327
	maxSeq    int
328
329
330
	embedSize int
}

331
332
333
// 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
334
335
func NewBatch(batchSize int, maxSeq int, embedSize int) (*Batch, error) {
	b := Batch{
336
337
338
339
		c:         C.llama_batch_init(C.int(batchSize*maxSeq), C.int(embedSize), C.int(maxSeq)),
		batchSize: batchSize,
		maxSeq:    maxSeq,
		embedSize: embedSize,
340
	}
Jesse Gross's avatar
Jesse Gross committed
341
342
343
344
345
346
347
348
349
350
351
352

	// 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
353
354
}

355
356
357
358
359
360
361
362
func (b *Batch) Size() int {
	return b.batchSize
}

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

363
364
365
366
367
368
369
370
371
372
373
374
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.
375
func (b *Batch) Add(token int, embed []float32, pos int, logits bool, seqIds ...int) {
376
	if !b.IsEmbedding() {
377
		unsafe.Slice(b.c.token, b.allocSize())[b.c.n_tokens] = C.llama_token(token)
378
	} else {
379
		copy(unsafe.Slice((*float32)(b.c.embd), b.allocSize()*b.embedSize)[int(b.c.n_tokens)*b.embedSize:], embed)
380
	}
381
382
	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))
383
384

	for i, s := range seqIds {
385
		unsafe.Slice((unsafe.Slice(b.c.seq_id, b.allocSize())[b.c.n_tokens]), C.int(len(seqIds)))[i] = C.int32_t(s)
386
387
388
	}

	if logits {
389
		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 1
390
391
	} else {
		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 0
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
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
	}

	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
}

501
// vision processing
502
type ClipContext struct {
503
	c *C.struct_clip_ctx
504
505
}

506
func NewClipContext(llamaContext *Context, modelPath string) (*ClipContext, error) {
507
508
	mp := C.CString(modelPath)
	defer C.free(unsafe.Pointer(mp))
509
	c := C.clip_model_load(mp, 1)
Jesse Gross's avatar
Jesse Gross committed
510
511
512
	if c == nil {
		return nil, fmt.Errorf("unable to load clip model: %v", modelPath)
	}
513

514
515
516
517
	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)
518
519
	}

520
	return &ClipContext{c: c}, nil
521
522
523
}

func (c *ClipContext) Free() {
524
	C.clip_free(c.c)
525
526
}

Jesse Gross's avatar
Jesse Gross committed
527
func (c *ClipContext) NewEmbed(llamaContext *Context, data []byte) ([][]float32, error) {
528
	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
529
530
531
	if l == nil {
		return nil, errors.New("unable to make llava embedding from image")
	}
532

533
	numTokens := int(l.n_image_pos)
534
535
	numEmbed := llamaContext.Model().NEmbd()

536
	s := unsafe.Slice((*float32)(l.embed), numEmbed*numTokens)
537
538
539
540
541
542
543
544
545

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

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

546
	C.llava_image_embed_free(l)
547

Jesse Gross's avatar
Jesse Gross committed
548
	return embed, nil
549
550
}

551
552
553
554
555
556
557
558
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
559
560
561
	if c == nil {
		return nil, fmt.Errorf("unable to load mllama model: %v", modelPath)
	}
562
563
564
565
566
567
568
569
570
571
572
573
574
575

	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
576
func (m *MllamaContext) NewEmbed(llamaContext *Context, data []byte, aspectRatioId int) ([][]float32, error) {
577
578
579
	img := C.mllama_image_init()
	defer C.mllama_image_free(img)

Jesse Gross's avatar
Jesse Gross committed
580
581
582
583
	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")
	}
584

585
	rows := make([]float32, m.EmbedSize(llamaContext))
Jesse Gross's avatar
Jesse Gross committed
586
587
588
589
	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")
	}
590

591
592
	embed := make([][]float32, 1)
	embed[0] = rows
593

Jesse Gross's avatar
Jesse Gross committed
594
	return embed, nil
595
596
}

597
598
599
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()
600

601
602
	return numTokens * numEmbed
}
603

604
605
func (c *Context) SetCrossAttention(state bool) {
	C.llama_set_cross_attention(c.c, C.bool(state))
606
607
}

608
609
610
611
func (c *Context) Synchronize() {
	C.llama_synchronize(c.c)
}

612
613
614
// sampling
// TODO: this is a temporary wrapper to allow calling C++ code from CGo
type SamplingContext struct {
615
	c *C.struct_common_sampler
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
}

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
636
func NewSamplingContext(model *Model, params SamplingParams) (*SamplingContext, error) {
637
	var cparams C.struct_common_sampler_cparams
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
	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.seed = C.uint32_t(params.Seed)

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

	cparams.grammar = grammar
656
	context := &SamplingContext{c: C.common_sampler_cinit(model.c, &cparams)}
Jesse Gross's avatar
Jesse Gross committed
657
658
659
660
	if context.c == nil {
		return nil, errors.New("unable to create sampling context")
	}

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

Jesse Gross's avatar
Jesse Gross committed
663
	return context, nil
664
665
666
}

func (s *SamplingContext) Reset() {
667
	C.common_sampler_creset(s.c)
668
669
}

670
func (s *SamplingContext) Sample(llamaContext *Context, idx int) int {
671
	return int(C.common_sampler_csample(s.c, llamaContext.c, C.int(idx)))
672
673
}

674
func (s *SamplingContext) Accept(id int, applyGrammar bool) {
675
	C.common_sampler_caccept(s.c, C.llama_token(id), C.bool(applyGrammar))
676
}
677

678
679
680
681
// 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))
682
683
684
685
686
687
688
	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
689
690
691
692
	n := C.schema_to_grammar(cStr, (*C.char)(unsafe.Pointer(&buf[0])), C.size_t(maxLen))
	if n == 0 {
		// preserve nil
		return nil
693
	}
694
	return buf[:n]
695
}