llama.go 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703
  1. package llama
  2. /*
  3. #cgo CFLAGS: -std=c11
  4. #cgo CXXFLAGS: -std=c++17
  5. #cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/include
  6. #cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/common
  7. #cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/examples/llava
  8. #cgo CPPFLAGS: -I${SRCDIR}/llama.cpp/src
  9. #cgo CPPFLAGS: -I${SRCDIR}/../ml/backend/ggml/ggml/include
  10. #include <stdlib.h>
  11. #include "ggml.h"
  12. #include "llama.h"
  13. #include "clip.h"
  14. #include "llava.h"
  15. #include "gguf.h"
  16. #include "mllama.h"
  17. #include "sampling_ext.h"
  18. extern bool llamaProgressCallback(float progress, void *user_data);
  19. extern void llamaLog(int level, char* text, void* user_data);
  20. typedef enum {COMP_UNKNOWN,COMP_GCC,COMP_CLANG} COMPILER;
  21. COMPILER inline get_compiler() {
  22. #if defined(__clang__)
  23. return COMP_CLANG;
  24. #elif defined(__GNUC__)
  25. return COMP_GCC;
  26. #else
  27. return UNKNOWN_COMPILER;
  28. #endif
  29. }
  30. */
  31. import "C"
  32. import (
  33. _ "embed"
  34. "errors"
  35. "fmt"
  36. "os"
  37. "runtime"
  38. "runtime/cgo"
  39. "slices"
  40. "strings"
  41. "sync/atomic"
  42. "unsafe"
  43. _ "github.com/ollama/ollama/llama/llama.cpp/common"
  44. _ "github.com/ollama/ollama/llama/llama.cpp/examples/llava"
  45. _ "github.com/ollama/ollama/llama/llama.cpp/src"
  46. "github.com/ollama/ollama/ml/backend/ggml/ggml/src"
  47. )
  48. func BackendInit() {
  49. ggml.OnceLoad()
  50. C.llama_backend_init()
  51. }
  52. func PrintSystemInfo() string {
  53. var compiler string
  54. switch C.get_compiler() {
  55. case C.COMP_UNKNOWN:
  56. compiler = "cgo(unknown_compiler)"
  57. case C.COMP_GCC:
  58. compiler = "cgo(gcc)"
  59. case C.COMP_CLANG:
  60. compiler = "cgo(clang)"
  61. }
  62. return C.GoString(C.llama_print_system_info()) + compiler
  63. }
  64. var logLevel atomic.Int32
  65. func init() {
  66. logLevel.Store(int32(C.GGML_LOG_LEVEL_INFO))
  67. C.llama_log_set((C.ggml_log_callback)(C.llamaLog), nil)
  68. }
  69. func EnableDebug() {
  70. logLevel.Store(int32(C.GGML_LOG_LEVEL_DEBUG))
  71. }
  72. //export llamaLog
  73. func llamaLog(level int32, text *C.char, _ unsafe.Pointer) {
  74. if level < logLevel.Load() {
  75. return
  76. }
  77. fmt.Fprint(os.Stderr, C.GoString(text))
  78. }
  79. func GetModelArch(modelPath string) (string, error) {
  80. mp := C.CString(modelPath)
  81. defer C.free(unsafe.Pointer(mp))
  82. gguf_ctx := C.gguf_init_from_file(mp, C.struct_gguf_init_params{no_alloc: true, ctx: (**C.struct_ggml_context)(C.NULL)})
  83. if gguf_ctx == nil {
  84. return "", errors.New("unable to load model file")
  85. }
  86. defer C.gguf_free(gguf_ctx)
  87. key := C.CString("general.architecture")
  88. defer C.free(unsafe.Pointer(key))
  89. arch_index := C.gguf_find_key(gguf_ctx, key)
  90. if int(arch_index) < 0 {
  91. return "", errors.New("unknown model architecture")
  92. }
  93. arch := C.gguf_get_val_str(gguf_ctx, arch_index)
  94. return C.GoString(arch), nil
  95. }
  96. type ContextParams struct {
  97. c C.struct_llama_context_params
  98. }
  99. func NewContextParams(numCtx int, batchSize int, numSeqMax int, threads int, flashAttention bool, kvCacheType string) ContextParams {
  100. params := C.llama_context_default_params()
  101. params.n_ctx = C.uint(numCtx)
  102. params.n_batch = C.uint(batchSize)
  103. params.n_seq_max = C.uint(numSeqMax)
  104. params.n_threads = C.int(threads)
  105. params.n_threads_batch = params.n_threads
  106. params.embeddings = C.bool(true)
  107. params.flash_attn = C.bool(flashAttention)
  108. params.type_k = kvCacheTypeFromStr(strings.ToLower(kvCacheType))
  109. params.type_v = kvCacheTypeFromStr(strings.ToLower(kvCacheType))
  110. return ContextParams{c: params}
  111. }
  112. // kvCacheTypeFromStr converts a string cache type to the corresponding GGML type value
  113. func kvCacheTypeFromStr(s string) C.enum_ggml_type {
  114. if s == "" {
  115. return C.GGML_TYPE_F16
  116. }
  117. switch s {
  118. case "q8_0":
  119. return C.GGML_TYPE_Q8_0
  120. case "q4_0":
  121. return C.GGML_TYPE_Q4_0
  122. default:
  123. return C.GGML_TYPE_F16
  124. }
  125. }
  126. type Context struct {
  127. c *C.struct_llama_context
  128. numThreads int
  129. }
  130. var ErrKvCacheFull = errors.New("could not find a kv cache slot")
  131. func (c *Context) Decode(batch *Batch) error {
  132. // Positive return values does not mean a fatal error, but rather a warning.
  133. // 0 - success
  134. // 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context)
  135. // < 0 - error
  136. code := int(C.llama_decode(c.c, batch.c))
  137. if code < 0 {
  138. return fmt.Errorf("llama_decode failed with code %d", code)
  139. }
  140. if code > 0 {
  141. return ErrKvCacheFull
  142. }
  143. return nil
  144. }
  145. func (c *Context) Model() *Model {
  146. return &Model{c: C.llama_get_model(c.c)}
  147. }
  148. func (c *Context) KvCacheSeqAdd(seqId int, p0 int, p1 int, delta int) {
  149. C.llama_kv_cache_seq_add(c.c, C.int(seqId), C.int(p0), C.int(p1), C.int(delta))
  150. }
  151. func (c *Context) KvCacheSeqRm(seqId int, p0 int, p1 int) bool {
  152. return bool(C.llama_kv_cache_seq_rm(c.c, C.int(seqId), C.int(p0), C.int(p1)))
  153. }
  154. func (c *Context) KvCacheSeqCp(srcSeqId int, dstSeqId int, p0 int, p1 int) {
  155. C.llama_kv_cache_seq_cp(c.c, C.int(srcSeqId), C.int(dstSeqId), C.int(p0), C.int(p1))
  156. }
  157. func (c *Context) KvCacheClear() {
  158. C.llama_kv_cache_clear(c.c)
  159. }
  160. func (c *Context) KvCacheDefrag() {
  161. C.llama_kv_cache_defrag(c.c)
  162. }
  163. // Get the embeddings for a sequence id
  164. func (c *Context) GetEmbeddingsSeq(seqId int) []float32 {
  165. e := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId)))
  166. if e == nil {
  167. return nil
  168. }
  169. embeddings := make([]float32, c.Model().NEmbd())
  170. _ = copy(embeddings, unsafe.Slice((*float32)(e), c.Model().NEmbd()))
  171. return embeddings
  172. }
  173. func (c *Context) GetEmbeddingsIth(i int) []float32 {
  174. e := unsafe.Pointer(C.llama_get_embeddings_ith(c.c, C.int32_t(i)))
  175. if e == nil {
  176. return nil
  177. }
  178. embeddings := make([]float32, c.Model().NEmbd())
  179. _ = copy(embeddings, unsafe.Slice((*float32)(e), c.Model().NEmbd()))
  180. return embeddings
  181. }
  182. type ModelParams struct {
  183. NumGpuLayers int
  184. MainGpu int
  185. UseMmap bool
  186. UseMlock bool
  187. TensorSplit []float32
  188. Progress func(float32)
  189. VocabOnly bool
  190. }
  191. //export llamaProgressCallback
  192. func llamaProgressCallback(progress C.float, userData unsafe.Pointer) C.bool {
  193. handle := *(*cgo.Handle)(userData)
  194. callback := handle.Value().(func(float32))
  195. callback(float32(progress))
  196. return true
  197. }
  198. func LoadModelFromFile(modelPath string, params ModelParams) (*Model, error) {
  199. cparams := C.llama_model_default_params()
  200. cparams.n_gpu_layers = C.int(params.NumGpuLayers)
  201. cparams.main_gpu = C.int32_t(params.MainGpu)
  202. cparams.use_mmap = C.bool(params.UseMmap)
  203. cparams.use_mlock = C.bool(params.UseMlock)
  204. cparams.vocab_only = C.bool(params.VocabOnly)
  205. if len(params.TensorSplit) > 0 {
  206. tensorSplitData := &params.TensorSplit[0]
  207. var tensorSplitPin runtime.Pinner
  208. tensorSplitPin.Pin(tensorSplitData)
  209. defer tensorSplitPin.Unpin()
  210. cparams.tensor_split = (*C.float)(unsafe.Pointer(tensorSplitData))
  211. }
  212. if params.Progress != nil {
  213. handle := cgo.NewHandle(params.Progress)
  214. defer handle.Delete()
  215. var handlePin runtime.Pinner
  216. handlePin.Pin(&handle)
  217. defer handlePin.Unpin()
  218. cparams.progress_callback = C.llama_progress_callback(C.llamaProgressCallback)
  219. cparams.progress_callback_user_data = unsafe.Pointer(&handle)
  220. }
  221. m := Model{c: C.llama_load_model_from_file(C.CString(modelPath), cparams)}
  222. if m.c == nil {
  223. return nil, fmt.Errorf("unable to load model: %s", modelPath)
  224. }
  225. return &m, nil
  226. }
  227. func FreeModel(model *Model) {
  228. C.llama_free_model(model.c)
  229. }
  230. func NewContextWithModel(model *Model, params ContextParams) (*Context, error) {
  231. c := Context{
  232. c: C.llama_new_context_with_model(model.c, params.c),
  233. numThreads: int(params.c.n_threads),
  234. }
  235. if c.c == nil {
  236. return nil, errors.New("unable to create llama context")
  237. }
  238. return &c, nil
  239. }
  240. func (m *Model) NumVocab() int {
  241. return int(C.llama_n_vocab(m.Vocab()))
  242. }
  243. func (m *Model) TokenIsEog(token int) bool {
  244. return bool(C.llama_token_is_eog(m.Vocab(), C.llama_token(token)))
  245. }
  246. func (m *Model) AddBOSToken() bool {
  247. return bool(C.llama_add_bos_token(m.Vocab()))
  248. }
  249. func (m *Model) ApplyLoraFromFile(context *Context, loraPath string, scale float32, threads int) error {
  250. cLoraPath := C.CString(loraPath)
  251. defer C.free(unsafe.Pointer(cLoraPath))
  252. loraAdapter := C.llama_adapter_lora_init(m.c, cLoraPath)
  253. if loraAdapter == nil {
  254. return errors.New("unable to load lora")
  255. }
  256. err := -1
  257. if loraAdapter != nil {
  258. err = int(C.llama_set_adapter_lora(context.c, loraAdapter, C.float(scale)))
  259. }
  260. if err != 0 {
  261. return errors.New("error applying lora from file")
  262. }
  263. return nil
  264. }
  265. func (m *Model) Vocab() *C.struct_llama_vocab {
  266. return C.llama_model_get_vocab(m.c)
  267. }
  268. type Batch struct {
  269. c C.struct_llama_batch
  270. batchSize int
  271. maxSeq int
  272. embedSize int
  273. }
  274. // Creates a new batch for either word tokens or image embeddings (if embedSize is non-zero).
  275. // Batches cannot contain both types at the same time. batchSize is the maximum number of entries
  276. // that can be added per sequence
  277. func NewBatch(batchSize int, maxSeq int, embedSize int) (*Batch, error) {
  278. b := Batch{
  279. c: C.llama_batch_init(C.int(batchSize*maxSeq), C.int(embedSize), C.int(maxSeq)),
  280. batchSize: batchSize,
  281. maxSeq: maxSeq,
  282. embedSize: embedSize,
  283. }
  284. // Check to see if any of the allocations in llama_batch_init() failed
  285. nilPointer := (embedSize == 0 && b.c.token == nil) || (embedSize != 0 && b.c.embd == nil) ||
  286. b.c.pos == nil || b.c.n_seq_id == nil || b.c.seq_id == nil || b.c.logits == nil ||
  287. slices.Contains(unsafe.Slice(b.c.seq_id, b.allocSize()), nil)
  288. if nilPointer {
  289. C.llama_batch_free(b.c)
  290. return nil, fmt.Errorf("unable to allocate batch (batchSize=%v maxSeq=%v embedSize=%v)", batchSize, maxSeq, embedSize)
  291. }
  292. return &b, nil
  293. }
  294. func (b *Batch) Size() int {
  295. return b.batchSize
  296. }
  297. func (b *Batch) allocSize() int {
  298. return b.batchSize * b.maxSeq
  299. }
  300. func (b *Batch) NumTokens() int {
  301. return int(b.c.n_tokens)
  302. }
  303. func (b *Batch) IsEmbedding() bool {
  304. return b.embedSize != 0
  305. }
  306. // Add adds either a token or an image embedding to the batch depending on the type
  307. // when the batch was initialized. The other argument will be ignored. Adds to the
  308. // batch with the given position for the given sequence ids, and optionally instructs
  309. // to include logits.
  310. func (b *Batch) Add(token int, embed []float32, pos int, logits bool, seqIds ...int) {
  311. if !b.IsEmbedding() {
  312. unsafe.Slice(b.c.token, b.allocSize())[b.c.n_tokens] = C.llama_token(token)
  313. } else {
  314. copy(unsafe.Slice((*float32)(b.c.embd), b.allocSize()*b.embedSize)[int(b.c.n_tokens)*b.embedSize:], embed)
  315. }
  316. unsafe.Slice(b.c.pos, b.allocSize())[b.c.n_tokens] = C.llama_pos(pos)
  317. unsafe.Slice(b.c.n_seq_id, b.allocSize())[b.c.n_tokens] = C.int(len(seqIds))
  318. for i, s := range seqIds {
  319. unsafe.Slice((unsafe.Slice(b.c.seq_id, b.allocSize())[b.c.n_tokens]), C.int(len(seqIds)))[i] = C.int32_t(s)
  320. }
  321. if logits {
  322. unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 1
  323. } else {
  324. unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 0
  325. }
  326. b.c.n_tokens += 1
  327. }
  328. func (b *Batch) Clear() {
  329. b.c.n_tokens = 0
  330. }
  331. func (b *Batch) Free() {
  332. b.batchSize = 0
  333. C.llama_batch_free(b.c)
  334. }
  335. type Model struct {
  336. c *C.struct_llama_model
  337. }
  338. func (m *Model) TokenToPiece(token int) string {
  339. tokenLen := 12
  340. buf := make([]byte, tokenLen)
  341. tokenLen = int(C.llama_token_to_piece(
  342. m.Vocab(),
  343. C.int32_t(token),
  344. (*C.char)(unsafe.Pointer(&buf[0])),
  345. C.int32_t(tokenLen),
  346. C.int32_t(0),
  347. C.bool(true),
  348. ))
  349. if tokenLen < 0 {
  350. tokenLen = -tokenLen
  351. buf = make([]byte, tokenLen)
  352. C.llama_token_to_piece(
  353. m.Vocab(),
  354. C.int32_t(token),
  355. (*C.char)(unsafe.Pointer(&buf[0])),
  356. C.int32_t(tokenLen),
  357. C.int32_t(0),
  358. C.bool(true),
  359. )
  360. }
  361. return strings.TrimRight(string(buf), "\x00")
  362. }
  363. func (m *Model) Tokenize(text string, addSpecial bool, parseSpecial bool) ([]int, error) {
  364. maxTokens := len(text) + 2
  365. cTokens := make([]C.llama_token, maxTokens)
  366. cText := C.CString(text)
  367. defer C.free(unsafe.Pointer(cText))
  368. result := C.llama_tokenize(
  369. m.Vocab(),
  370. cText,
  371. C.int32_t(len(text)),
  372. &cTokens[0],
  373. C.int32_t(maxTokens),
  374. C.bool(addSpecial),
  375. C.bool(parseSpecial),
  376. )
  377. // if the result is negative, reallocate and retry with the correct buffer size
  378. if result < 0 {
  379. maxTokens = int(-result)
  380. cTokens = make([]C.llama_token, maxTokens)
  381. result = C.llama_tokenize(
  382. m.Vocab(),
  383. cText,
  384. C.int32_t(len(text)),
  385. &cTokens[0],
  386. C.int32_t(maxTokens),
  387. C.bool(addSpecial),
  388. C.bool(parseSpecial),
  389. )
  390. if result < 0 {
  391. return nil, fmt.Errorf("tokenization failed, required %d tokens", -result)
  392. }
  393. }
  394. tokens := make([]int, result)
  395. for i := range result {
  396. tokens[i] = int(cTokens[i])
  397. }
  398. return tokens, nil
  399. }
  400. func (m *Model) NEmbd() int {
  401. return int(C.llama_n_embd(m.c))
  402. }
  403. func Quantize(infile, outfile string, ftype uint32) error {
  404. cinfile := C.CString(infile)
  405. defer C.free(unsafe.Pointer(cinfile))
  406. coutfile := C.CString(outfile)
  407. defer C.free(unsafe.Pointer(coutfile))
  408. params := C.llama_model_quantize_default_params()
  409. params.nthread = -1
  410. params.ftype = ftype
  411. if rc := C.llama_model_quantize(cinfile, coutfile, &params); rc != 0 {
  412. return fmt.Errorf("llama_model_quantize: %d", rc)
  413. }
  414. return nil
  415. }
  416. // vision processing
  417. type ClipContext struct {
  418. c *C.struct_clip_ctx
  419. }
  420. func NewClipContext(llamaContext *Context, modelPath string) (*ClipContext, error) {
  421. mp := C.CString(modelPath)
  422. defer C.free(unsafe.Pointer(mp))
  423. c := C.clip_model_load(mp, 1)
  424. if c == nil {
  425. return nil, fmt.Errorf("unable to load clip model: %v", modelPath)
  426. }
  427. projEmbedSize := int(C.clip_n_mmproj_embd(c))
  428. modelEmbedSize := llamaContext.Model().NEmbd()
  429. if projEmbedSize != modelEmbedSize {
  430. return nil, fmt.Errorf("projector embedding size (%d) does not match model (%d)", projEmbedSize, modelEmbedSize)
  431. }
  432. return &ClipContext{c: c}, nil
  433. }
  434. func (c *ClipContext) Free() {
  435. C.clip_free(c.c)
  436. }
  437. func (c *ClipContext) NewEmbed(llamaContext *Context, data []byte) ([][]float32, error) {
  438. l := C.llava_image_embed_make_with_bytes(c.c, C.int(llamaContext.numThreads), (*C.uchar)(unsafe.Pointer(&data[0])), C.int(len(data)))
  439. if l == nil {
  440. return nil, errors.New("unable to make llava embedding from image")
  441. }
  442. numTokens := int(l.n_image_pos)
  443. numEmbed := llamaContext.Model().NEmbd()
  444. s := unsafe.Slice((*float32)(l.embed), numEmbed*numTokens)
  445. embed := make([][]float32, numTokens)
  446. rows := make([]float32, len(s))
  447. copy(rows, s)
  448. for i := range embed {
  449. embed[i] = rows[i*numEmbed : (i+1)*numEmbed]
  450. }
  451. C.llava_image_embed_free(l)
  452. return embed, nil
  453. }
  454. type MllamaContext struct {
  455. c *C.struct_mllama_ctx
  456. }
  457. func NewMllamaContext(llamaContext *Context, modelPath string) (*MllamaContext, error) {
  458. mp := C.CString(modelPath)
  459. defer C.free(unsafe.Pointer(mp))
  460. c := C.mllama_model_load(mp, 1)
  461. if c == nil {
  462. return nil, fmt.Errorf("unable to load mllama model: %v", modelPath)
  463. }
  464. projEmbedSize := int(C.mllama_n_embd(c))
  465. modelEmbedSize := llamaContext.Model().NEmbd()
  466. if projEmbedSize != modelEmbedSize {
  467. return nil, fmt.Errorf("projector embedding size (%d) does not match model (%d)", projEmbedSize, modelEmbedSize)
  468. }
  469. return &MllamaContext{c: c}, nil
  470. }
  471. func (m *MllamaContext) Free() {
  472. C.mllama_free(m.c)
  473. }
  474. func (m *MllamaContext) NewEmbed(llamaContext *Context, data []byte, aspectRatioId int) ([][]float32, error) {
  475. img := C.mllama_image_init()
  476. defer C.mllama_image_free(img)
  477. ok := bool(C.mllama_image_load_from_data(unsafe.Pointer(&data[0]), C.int(len(data)), 560, 560, 3, 4, C.int(aspectRatioId), img))
  478. if !ok {
  479. return nil, errors.New("unable to load mllama image data")
  480. }
  481. rows := make([]float32, m.EmbedSize(llamaContext))
  482. ok = bool(C.mllama_image_encode(m.c, C.int(llamaContext.numThreads), img, (*C.float)(unsafe.Pointer(&rows[0]))))
  483. if !ok {
  484. return nil, errors.New("unable to make mllama embedding from image")
  485. }
  486. embed := make([][]float32, 1)
  487. embed[0] = rows
  488. return embed, nil
  489. }
  490. func (m *MllamaContext) EmbedSize(llamaContext *Context) int {
  491. numTokens := int(C.mllama_n_positions(m.c) * C.mllama_n_tiles(m.c))
  492. numEmbed := llamaContext.Model().NEmbd()
  493. return numTokens * numEmbed
  494. }
  495. func (c *Context) SetCrossAttention(state bool) {
  496. C.llama_set_cross_attention(c.c, C.bool(state))
  497. }
  498. func (c *Context) Synchronize() {
  499. C.llama_synchronize(c.c)
  500. }
  501. // sampling
  502. // TODO: this is a temporary wrapper to allow calling C++ code from CGo
  503. type SamplingContext struct {
  504. c *C.struct_common_sampler
  505. }
  506. type SamplingParams struct {
  507. TopK int
  508. TopP float32
  509. MinP float32
  510. TypicalP float32
  511. Temp float32
  512. RepeatLastN int
  513. PenaltyRepeat float32
  514. PenaltyFreq float32
  515. PenaltyPresent float32
  516. Mirostat int
  517. MirostatTau float32
  518. MirostatEta float32
  519. PenalizeNl bool
  520. Seed uint32
  521. Grammar string
  522. }
  523. func NewSamplingContext(model *Model, params SamplingParams) (*SamplingContext, error) {
  524. var cparams C.struct_common_sampler_cparams
  525. cparams.top_k = C.int32_t(params.TopK)
  526. cparams.top_p = C.float(params.TopP)
  527. cparams.min_p = C.float(params.MinP)
  528. cparams.typical_p = C.float(params.TypicalP)
  529. cparams.temp = C.float(params.Temp)
  530. cparams.penalty_last_n = C.int32_t(params.RepeatLastN)
  531. cparams.penalty_repeat = C.float(params.PenaltyRepeat)
  532. cparams.penalty_freq = C.float(params.PenaltyFreq)
  533. cparams.penalty_present = C.float(params.PenaltyFreq)
  534. cparams.mirostat = C.int32_t(params.Mirostat)
  535. cparams.mirostat_tau = C.float(params.MirostatTau)
  536. cparams.mirostat_eta = C.float(params.MirostatEta)
  537. cparams.seed = C.uint32_t(params.Seed)
  538. grammar := C.CString(params.Grammar)
  539. defer C.free(unsafe.Pointer(grammar))
  540. cparams.grammar = grammar
  541. context := &SamplingContext{c: C.common_sampler_cinit(model.c, &cparams)}
  542. if context.c == nil {
  543. return nil, errors.New("unable to create sampling context")
  544. }
  545. runtime.SetFinalizer(context, func(s *SamplingContext) { C.common_sampler_cfree(s.c) })
  546. return context, nil
  547. }
  548. func (s *SamplingContext) Reset() {
  549. C.common_sampler_creset(s.c)
  550. }
  551. func (s *SamplingContext) Sample(llamaContext *Context, idx int) int {
  552. return int(C.common_sampler_csample(s.c, llamaContext.c, C.int(idx)))
  553. }
  554. func (s *SamplingContext) Accept(id int, applyGrammar bool) {
  555. C.common_sampler_caccept(s.c, C.llama_token(id), C.bool(applyGrammar))
  556. }
  557. // SchemaToGrammar converts the provided JSON schema to a grammar. It returns
  558. // nil if the provided schema is invalid JSON or an invalid JSON schema.
  559. func SchemaToGrammar(schema []byte) []byte {
  560. cStr := C.CString(string(schema))
  561. defer C.free(unsafe.Pointer(cStr))
  562. // Allocate buffer for grammar output with reasonable size
  563. const maxLen = 32768 // 32KB
  564. buf := make([]byte, maxLen)
  565. // Call C function to convert schema to grammar
  566. n := C.schema_to_grammar(cStr, (*C.char)(unsafe.Pointer(&buf[0])), C.size_t(maxLen))
  567. if n == 0 {
  568. // preserve nil
  569. return nil
  570. }
  571. return buf[:n]
  572. }