llama.go 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630
  1. package llama
  2. //go:generate make -j 8
  3. /*
  4. #cgo CFLAGS: -O2 -std=c11 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE
  5. #cgo CXXFLAGS: -O2 -std=c++11 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE
  6. #cgo amd64,avx CFLAGS: -mavx
  7. #cgo amd64,avx CXXFLAGS: -mavx
  8. #cgo amd64,avx2 CFLAGS: -mavx2 -mfma
  9. #cgo amd64,avx2 CXXFLAGS: -mavx2 -mfma
  10. #cgo amd64,f16c CFLAGS: -mf16c
  11. #cgo amd64,f16c CXXFLAGS: -mf16c
  12. #cgo amd64,fma CFLAGS: -mfma
  13. #cgo amd64,fma CXXFLAGS: -mfma
  14. #cgo avx CFLAGS: -mavx
  15. #cgo avx CXXFLAGS: -mavx
  16. #cgo avx2 CFLAGS: -mavx2 -mfma -mf16c
  17. #cgo avx2 CXXFLAGS: -mavx2 -mfma -mf16c
  18. #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
  19. #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
  20. #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
  21. #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
  22. #cgo cuda_v11 LDFLAGS: -lggml_cuda_v11 -L/usr/local/cuda-11/lib64
  23. #cgo cuda_v12 LDFLAGS: -lggml_cuda_v12 -L/usr/local/cuda-12/lib64
  24. #cgo darwin,amd64 CFLAGS: -Wno-incompatible-pointer-types-discards-qualifiers
  25. #cgo darwin,amd64 CXXFLAGS: -Wno-incompatible-pointer-types-discards-qualifiers
  26. #cgo darwin,amd64 LDFLAGS: -framework Foundation
  27. #cgo darwin,amd64,avx2 CFLAGS: -DGGML_USE_ACCELERATE -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
  28. #cgo darwin,amd64,avx2 CXXFLAGS: -DGGML_USE_ACCELERATE -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
  29. #cgo darwin,amd64,avx2 LDFLAGS: -framework Accelerate
  30. #cgo darwin,arm64 CFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS
  31. #cgo darwin,arm64 CXXFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS
  32. #cgo darwin,arm64 LDFLAGS: -framework Foundation -framework Metal -framework MetalKit -framework Accelerate
  33. #cgo linux CFLAGS: -D_GNU_SOURCE
  34. #cgo linux CXXFLAGS: -D_GNU_SOURCE
  35. #cgo linux,amd64 LDFLAGS: -L${SRCDIR}/build/Linux/amd64
  36. #cgo linux,amd64 LDFLAGS: -L${SRCDIR}/build/Linux/amd64
  37. #cgo linux,arm64 CFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA -D__ARM_FEATURE_MATMUL_INT8
  38. #cgo linux,arm64 CXXFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA -D__ARM_FEATURE_MATMUL_INT8
  39. #cgo linux,arm64 LDFLAGS: -L${SRCDIR}/build/Linux/arm64
  40. #cgo linux,arm64,sve CFLAGS: -march=armv8.6-a+sve
  41. #cgo linux,arm64,sve CXXFLAGS: -march=armv8.6-a+sve
  42. #cgo linux,cuda LDFLAGS: -lcuda -lcudart -lcublas -lcublasLt -lpthread -ldl -lrt -lresolv
  43. #cgo linux,rocm LDFLAGS: -L/opt/rocm/lib -lpthread -ldl -lrt -lresolv
  44. #cgo rocm CFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIPBLAS -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
  45. #cgo rocm CXXFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIPBLAS -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
  46. #cgo rocm LDFLAGS: -L${SRCDIR} -lggml_rocm -lhipblas -lamdhip64 -lrocblas
  47. #cgo windows CFLAGS: -Wno-discarded-qualifiers -D_WIN32_WINNT=0x602
  48. #cgo windows CXXFLAGS: -D_WIN32_WINNT=0x602
  49. #cgo windows LDFLAGS: -lmsvcrt
  50. #cgo windows LDFLAGS: -lmsvcrt -static-libstdc++ -static-libgcc -static
  51. #cgo windows,amd64 LDFLAGS: -L${SRCDIR}/build/Windows/amd64
  52. #cgo windows,amd64 LDFLAGS: -L${SRCDIR}/build/Windows/amd64
  53. #cgo windows,arm64 CFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
  54. #cgo windows,arm64 CXXFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
  55. #cgo windows,arm64 LDFLAGS: -L${SRCDIR}/build/Windows/arm64
  56. #cgo windows,arm64 LDFLAGS: -L${SRCDIR}/build/Windows/arm64
  57. #cgo windows,cuda LDFLAGS: -lcuda -lcudart -lcublas -lcublasLt
  58. #cgo windows,rocm LDFLAGS: -lggml_rocm -lhipblas -lamdhip64 -lrocblas
  59. #include <stdlib.h>
  60. #include "llama.h"
  61. #include "clip.h"
  62. #include "ggml.h"
  63. #include "llava.h"
  64. #include "mllama.h"
  65. #include "sampling_ext.h"
  66. bool llamaProgressCallback(float progress, void *user_data);
  67. typedef enum {COMP_UNKNOWN,COMP_GCC,COMP_CLANG} COMPILER;
  68. COMPILER inline get_compiler() {
  69. #if defined(__clang__)
  70. return COMP_CLANG;
  71. #elif defined(__GNUC__)
  72. return COMP_GCC;
  73. #else
  74. return UNKNOWN_COMPILER;
  75. #endif
  76. }
  77. */
  78. import "C"
  79. import (
  80. _ "embed"
  81. "errors"
  82. "fmt"
  83. "runtime"
  84. "runtime/cgo"
  85. "strings"
  86. "unsafe"
  87. )
  88. var CpuFeatures = ""
  89. func BackendInit() {
  90. C.llama_backend_init()
  91. }
  92. func PrintSystemInfo() string {
  93. var compiler string
  94. switch C.get_compiler() {
  95. case C.COMP_UNKNOWN:
  96. compiler = "cgo(unknown_compiler)"
  97. case C.COMP_GCC:
  98. compiler = "cgo(gcc)"
  99. case C.COMP_CLANG:
  100. compiler = "cgo(clang)"
  101. }
  102. return C.GoString(C.llama_print_system_info()) + compiler
  103. }
  104. type ContextParams struct {
  105. c C.struct_llama_context_params
  106. }
  107. func NewContextParams(numCtx int, batchSize int, numSeqMax int, threads int, flashAttention bool) ContextParams {
  108. params := C.llama_context_default_params()
  109. params.n_ctx = C.uint(numCtx)
  110. params.n_batch = C.uint(batchSize)
  111. params.n_seq_max = C.uint(numSeqMax)
  112. params.n_threads = C.int(threads)
  113. params.n_threads_batch = params.n_threads
  114. params.embeddings = C.bool(true)
  115. params.flash_attn = C.bool(flashAttention)
  116. return ContextParams{c: params}
  117. }
  118. type Context struct {
  119. c *C.struct_llama_context
  120. numThreads int
  121. }
  122. func (c *Context) KvCacheClear() {
  123. C.llama_kv_cache_clear(c.c)
  124. }
  125. func (c *Context) Decode(batch *Batch) error {
  126. // Positive return values does not mean a fatal error, but rather a warning.
  127. // 0 - success
  128. // 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context)
  129. // < 0 - error
  130. code := int(C.llama_decode(c.c, batch.c))
  131. if code < 0 {
  132. return fmt.Errorf("llama_decode failed with code %d", code)
  133. }
  134. if code > 0 {
  135. return fmt.Errorf("could not find a KV slot for the batch - try reducing the size of the batch or increase the context. code: %d", code)
  136. }
  137. return nil
  138. }
  139. func (c *Context) Model() *Model {
  140. return &Model{c: C.llama_get_model(c.c)}
  141. }
  142. func (c *Context) KvCacheSeqAdd(seqId int, p0 int, p1 int, delta int) {
  143. C.llama_kv_cache_seq_add(c.c, C.int(seqId), C.int(p0), C.int(p1), C.int(delta))
  144. }
  145. func (c *Context) KvCacheSeqRm(seqId int, p0 int, p1 int) bool {
  146. return bool(C.llama_kv_cache_seq_rm(c.c, C.int(seqId), C.int(p0), C.int(p1)))
  147. }
  148. func (c *Context) KvCacheSeqCp(srcSeqId int, dstSeqId int, p0 int, p1 int) {
  149. C.llama_kv_cache_seq_cp(c.c, C.int(srcSeqId), C.int(dstSeqId), C.int(p0), C.int(p1))
  150. }
  151. // Get the embeddings for a sequence id
  152. func (c *Context) GetEmbeddingsSeq(seqId int) []float32 {
  153. embeddings := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId)))
  154. if embeddings == nil {
  155. return nil
  156. }
  157. return unsafe.Slice((*float32)(embeddings), c.Model().NEmbd())
  158. }
  159. func (c *Context) GetEmbeddingsIth(i int) []float32 {
  160. embeddings := unsafe.Pointer(C.llama_get_embeddings_ith(c.c, C.int32_t(i)))
  161. if embeddings == nil {
  162. return nil
  163. }
  164. return unsafe.Slice((*float32)(embeddings), c.Model().NEmbd())
  165. }
  166. type ModelParams struct {
  167. NumGpuLayers int
  168. MainGpu int
  169. UseMmap bool
  170. UseMlock bool
  171. TensorSplit []float32
  172. Progress func(float32)
  173. VocabOnly bool
  174. }
  175. //export llamaProgressCallback
  176. func llamaProgressCallback(progress C.float, userData unsafe.Pointer) C.bool {
  177. handle := *(*cgo.Handle)(userData)
  178. callback := handle.Value().(func(float32))
  179. callback(float32(progress))
  180. return true
  181. }
  182. func LoadModelFromFile(modelPath string, params ModelParams) (*Model, error) {
  183. cparams := C.llama_model_default_params()
  184. cparams.n_gpu_layers = C.int(params.NumGpuLayers)
  185. cparams.main_gpu = C.int32_t(params.MainGpu)
  186. cparams.use_mmap = C.bool(params.UseMmap)
  187. cparams.use_mlock = C.bool(params.UseMlock)
  188. cparams.vocab_only = C.bool(params.VocabOnly)
  189. if len(params.TensorSplit) > 0 {
  190. tensorSplitData := &params.TensorSplit[0]
  191. var tensorSplitPin runtime.Pinner
  192. tensorSplitPin.Pin(tensorSplitData)
  193. defer tensorSplitPin.Unpin()
  194. cparams.tensor_split = (*C.float)(unsafe.Pointer(tensorSplitData))
  195. }
  196. if params.Progress != nil {
  197. handle := cgo.NewHandle(params.Progress)
  198. defer handle.Delete()
  199. var handlePin runtime.Pinner
  200. handlePin.Pin(&handle)
  201. defer handlePin.Unpin()
  202. cparams.progress_callback = C.llama_progress_callback(C.llamaProgressCallback)
  203. cparams.progress_callback_user_data = unsafe.Pointer(&handle)
  204. }
  205. m := Model{c: C.llama_load_model_from_file(C.CString(modelPath), cparams)}
  206. if m.c == (*C.struct_llama_model)(C.NULL) {
  207. return nil, fmt.Errorf("unable to load model: %s", modelPath)
  208. }
  209. return &m, nil
  210. }
  211. func FreeModel(model *Model) {
  212. C.llama_free_model(model.c)
  213. }
  214. func NewContextWithModel(model *Model, params ContextParams) (*Context, error) {
  215. c := Context{
  216. c: C.llama_new_context_with_model(model.c, params.c),
  217. numThreads: int(params.c.n_threads),
  218. }
  219. if c.c == (*C.struct_llama_context)(C.NULL) {
  220. return nil, errors.New("unable to create llama context")
  221. }
  222. return &c, nil
  223. }
  224. func (m *Model) NumVocab() int {
  225. return int(C.llama_n_vocab(m.c))
  226. }
  227. func (m *Model) TokenIsEog(token int) bool {
  228. return bool(C.llama_token_is_eog(m.c, C.llama_token(token)))
  229. }
  230. func (m *Model) AddBOSToken() bool {
  231. return bool(C.llama_add_bos_token(m.c))
  232. }
  233. func (m *Model) ApplyLoraFromFile(context *Context, loraPath string, scale float32, threads int) error {
  234. cLoraPath := C.CString(loraPath)
  235. defer C.free(unsafe.Pointer(cLoraPath))
  236. loraAdapter := C.llama_lora_adapter_init(m.c, cLoraPath)
  237. err := -1
  238. if loraAdapter != nil {
  239. err = int(C.llama_lora_adapter_set(context.c, loraAdapter, C.float(scale)))
  240. }
  241. if err != 0 {
  242. return errors.New("error applying lora from file")
  243. }
  244. return nil
  245. }
  246. type Batch struct {
  247. c C.struct_llama_batch
  248. batchSize int
  249. embedSize int
  250. }
  251. // Creates a new batch for either word tokens if embed is 0 or
  252. // image embeddings if embed is specified. Batches cannot contain
  253. // both types at the same time
  254. func NewBatch(nTokens int, embed int, maxSeq int) *Batch {
  255. return &Batch{
  256. c: C.llama_batch_init(C.int(nTokens), C.int(embed), C.int(maxSeq)),
  257. batchSize: nTokens,
  258. embedSize: embed,
  259. }
  260. }
  261. func (b *Batch) NumTokens() int {
  262. return int(b.c.n_tokens)
  263. }
  264. func (b *Batch) IsEmbedding() bool {
  265. return b.embedSize != 0
  266. }
  267. // Add adds either a token or an image embedding to the batch depending on the type
  268. // when the batch was initialized. The other argument will be ignored. Adds to the
  269. // batch with the given position for the given sequence ids, and optionally instructs
  270. // to include logits.
  271. func (b *Batch) Add(token int, embed []float32, pos int, seqIds []int, logits bool) {
  272. if !b.IsEmbedding() {
  273. unsafe.Slice(b.c.token, b.batchSize)[b.c.n_tokens] = C.llama_token(token)
  274. } else {
  275. copy(unsafe.Slice((*float32)(b.c.embd), b.batchSize*b.embedSize)[int(b.c.n_tokens)*b.embedSize:], embed)
  276. }
  277. unsafe.Slice(b.c.pos, b.batchSize)[b.c.n_tokens] = C.llama_pos(pos)
  278. unsafe.Slice(b.c.n_seq_id, b.batchSize)[b.c.n_tokens] = C.int(len(seqIds))
  279. for i, s := range seqIds {
  280. unsafe.Slice((unsafe.Slice(b.c.seq_id, b.batchSize)[b.c.n_tokens]), C.int(len(seqIds)))[i] = C.int32_t(s)
  281. }
  282. if logits {
  283. unsafe.Slice(b.c.logits, b.batchSize)[b.c.n_tokens] = 1
  284. }
  285. b.c.n_tokens += 1
  286. }
  287. func (b *Batch) Clear() {
  288. b.c.n_tokens = 0
  289. }
  290. func (b *Batch) Free() {
  291. b.batchSize = 0
  292. C.llama_batch_free(b.c)
  293. }
  294. type Model struct {
  295. c *C.struct_llama_model
  296. }
  297. func (m *Model) TokenToPiece(token int) string {
  298. tokenLen := 12
  299. buf := make([]byte, tokenLen)
  300. tokenLen = int(C.llama_token_to_piece(
  301. m.c,
  302. C.int32_t(token),
  303. (*C.char)(unsafe.Pointer(&buf[0])),
  304. C.int32_t(tokenLen),
  305. C.int32_t(0),
  306. C.bool(true),
  307. ))
  308. if tokenLen < 0 {
  309. tokenLen = -tokenLen
  310. buf = make([]byte, tokenLen)
  311. C.llama_token_to_piece(
  312. m.c,
  313. C.int32_t(token),
  314. (*C.char)(unsafe.Pointer(&buf[0])),
  315. C.int32_t(tokenLen),
  316. C.int32_t(0),
  317. C.bool(true),
  318. )
  319. }
  320. return strings.TrimRight(string(buf), "\x00")
  321. }
  322. func (m *Model) Tokenize(text string, addSpecial bool, parseSpecial bool) ([]int, error) {
  323. maxTokens := len(text) + 2
  324. cTokens := make([]C.llama_token, maxTokens)
  325. cText := C.CString(text)
  326. defer C.free(unsafe.Pointer(cText))
  327. result := C.llama_tokenize(
  328. m.c,
  329. cText,
  330. C.int32_t(len(text)),
  331. &cTokens[0],
  332. C.int32_t(maxTokens),
  333. C.bool(addSpecial),
  334. C.bool(parseSpecial),
  335. )
  336. // if the result is negative, reallocate and retry with the correct buffer size
  337. if result < 0 {
  338. maxTokens = int(-result)
  339. cTokens = make([]C.llama_token, maxTokens)
  340. result = C.llama_tokenize(
  341. m.c,
  342. cText,
  343. C.int32_t(len(text)),
  344. &cTokens[0],
  345. C.int32_t(maxTokens),
  346. C.bool(addSpecial),
  347. C.bool(parseSpecial),
  348. )
  349. if result < 0 {
  350. return nil, fmt.Errorf("tokenization failed, required %d tokens", -result)
  351. }
  352. }
  353. tokens := make([]int, result)
  354. for i := range result {
  355. tokens[i] = int(cTokens[i])
  356. }
  357. return tokens, nil
  358. }
  359. func (m *Model) NEmbd() int {
  360. return int(C.llama_n_embd(m.c))
  361. }
  362. func Quantize(infile, outfile string, ftype uint32) error {
  363. cinfile := C.CString(infile)
  364. defer C.free(unsafe.Pointer(cinfile))
  365. coutfile := C.CString(outfile)
  366. defer C.free(unsafe.Pointer(coutfile))
  367. params := C.llama_model_quantize_default_params()
  368. params.nthread = -1
  369. params.ftype = ftype
  370. if rc := C.llama_model_quantize(cinfile, coutfile, &params); rc != 0 {
  371. return fmt.Errorf("llama_model_quantize: %d", rc)
  372. }
  373. return nil
  374. }
  375. // llava
  376. type ClipContext struct {
  377. c *C.struct_clip_ctx
  378. m *C.struct_mllama_ctx
  379. IsMllama bool
  380. embedPin runtime.Pinner
  381. pinned bool
  382. }
  383. func getVisionArch(mp *C.char) (string, error) {
  384. gguf_ctx := C.gguf_init_from_file(mp, C.struct_gguf_init_params{no_alloc: true, ctx: (**C.struct_ggml_context)(C.NULL)})
  385. if gguf_ctx == nil {
  386. return "", errors.New("unable to load vision projector")
  387. }
  388. defer C.gguf_free(gguf_ctx)
  389. arch_index := C.gguf_find_key(gguf_ctx, C.CString("general.architecture"))
  390. if int(arch_index) < 0 {
  391. return "", errors.New("unknown vision model architecture")
  392. }
  393. arch := C.gguf_get_val_str(gguf_ctx, arch_index)
  394. return C.GoString(arch), nil
  395. }
  396. func NewClipContext(modelPath string) (*ClipContext, error) {
  397. mp := C.CString(modelPath)
  398. defer C.free(unsafe.Pointer(mp))
  399. arch, err := getVisionArch(mp)
  400. if err != nil {
  401. return nil, err
  402. }
  403. var cc ClipContext
  404. if arch == "clip" {
  405. cc.c = C.clip_model_load(mp, 1)
  406. } else if arch == "mllama" {
  407. cc.m = C.mllama_model_load(mp, 1)
  408. cc.IsMllama = true
  409. } else {
  410. return nil, fmt.Errorf("unknown vision model architecture: %s", arch)
  411. }
  412. // XXX: check embedding size?
  413. return &cc, nil
  414. }
  415. func (c *ClipContext) Free() {
  416. if c.c != nil {
  417. C.clip_free(c.c)
  418. }
  419. if c.m != nil {
  420. C.mllama_free(c.m)
  421. }
  422. }
  423. func NewLlavaImageEmbed(llamaContext *Context, clipContext *ClipContext, data []byte) [][]float32 {
  424. c := C.llava_image_embed_make_with_bytes(clipContext.c, C.int(llamaContext.numThreads), (*C.uchar)(unsafe.Pointer(&data[0])), C.int(len(data)))
  425. numTokens := int(c.n_image_pos)
  426. numEmbed := llamaContext.Model().NEmbd()
  427. s := unsafe.Slice((*float32)(c.embed), numEmbed*numTokens)
  428. embed := make([][]float32, numTokens)
  429. rows := make([]float32, len(s))
  430. copy(rows, s)
  431. for i := range embed {
  432. embed[i] = rows[i*numEmbed : (i+1)*numEmbed]
  433. }
  434. C.llava_image_embed_free(c)
  435. return embed
  436. }
  437. func NewMllamaImageEmbed(llamaContext *Context, clipContext *ClipContext, data []byte, aspectRatioId int) [][]float32 {
  438. img := C.mllama_image_init()
  439. defer C.mllama_image_free(img)
  440. C.mllama_image_load_from_data(unsafe.Pointer(&data[0]), C.int(len(data)), 560, 560, 3, 4, C.int(aspectRatioId), img)
  441. numTokens := int(C.mllama_n_positions(clipContext.m) * C.mllama_n_tiles(clipContext.m))
  442. numEmbed := llamaContext.Model().NEmbd()
  443. rows := make([]float32, numEmbed*numTokens)
  444. C.mllama_image_encode(clipContext.m, C.int(llamaContext.numThreads), img, (*C.float)(unsafe.Pointer(&rows[0])))
  445. embed := make([][]float32, numTokens)
  446. for i := range embed {
  447. embed[i] = rows[i*numEmbed : (i+1)*numEmbed]
  448. }
  449. return embed
  450. }
  451. // This really needs to be set on a batch instead
  452. func MllamaSetCrossAttn(llamaContext *Context, clipContext *ClipContext, embed [][]float32) {
  453. if embed != nil {
  454. if clipContext.pinned {
  455. panic("Cross attention state already pinned")
  456. }
  457. embedData := &embed[0][0]
  458. clipContext.embedPin.Pin(embedData)
  459. clipContext.pinned = true
  460. C.llama_set_cross_attn_state(llamaContext.c, (*C.float)(unsafe.Pointer(embedData)))
  461. } else {
  462. C.llama_set_cross_attn_state(llamaContext.c, (*C.float)(C.NULL))
  463. if clipContext.pinned {
  464. clipContext.embedPin.Unpin()
  465. clipContext.pinned = false
  466. }
  467. }
  468. }
  469. // sampling
  470. // TODO: this is a temporary wrapper to allow calling C++ code from CGo
  471. type SamplingContext struct {
  472. c *C.struct_gpt_sampler
  473. }
  474. type SamplingParams struct {
  475. TopK int
  476. TopP float32
  477. MinP float32
  478. TfsZ float32
  479. TypicalP float32
  480. Temp float32
  481. RepeatLastN int
  482. PenaltyRepeat float32
  483. PenaltyFreq float32
  484. PenaltyPresent float32
  485. Mirostat int
  486. MirostatTau float32
  487. MirostatEta float32
  488. PenalizeNl bool
  489. Seed uint32
  490. Grammar string
  491. }
  492. func NewSamplingContext(model *Model, params SamplingParams) *SamplingContext {
  493. var cparams C.struct_gpt_sampler_cparams
  494. cparams.top_k = C.int32_t(params.TopK)
  495. cparams.top_p = C.float(params.TopP)
  496. cparams.min_p = C.float(params.MinP)
  497. cparams.tfs_z = C.float(params.TfsZ)
  498. cparams.typical_p = C.float(params.TypicalP)
  499. cparams.temp = C.float(params.Temp)
  500. cparams.penalty_last_n = C.int32_t(params.RepeatLastN)
  501. cparams.penalty_repeat = C.float(params.PenaltyRepeat)
  502. cparams.penalty_freq = C.float(params.PenaltyFreq)
  503. cparams.penalty_present = C.float(params.PenaltyFreq)
  504. cparams.mirostat = C.int32_t(params.Mirostat)
  505. cparams.mirostat_tau = C.float(params.MirostatTau)
  506. cparams.mirostat_eta = C.float(params.MirostatEta)
  507. cparams.penalize_nl = C.bool(params.PenalizeNl)
  508. cparams.seed = C.uint32_t(params.Seed)
  509. grammar := C.CString(params.Grammar)
  510. defer C.free(unsafe.Pointer(grammar))
  511. cparams.grammar = grammar
  512. context := &SamplingContext{c: C.gpt_sampler_cinit(model.c, &cparams)}
  513. runtime.SetFinalizer(context, func(s *SamplingContext) { C.gpt_sampler_cfree(s.c) })
  514. return context
  515. }
  516. func (s *SamplingContext) Reset() {
  517. C.gpt_sampler_creset(s.c)
  518. }
  519. func (s *SamplingContext) Sample(llamaContext *Context, idx int) int {
  520. return int(C.gpt_sampler_csample(s.c, llamaContext.c, C.int(idx)))
  521. }
  522. func (s *SamplingContext) Accept(id int, applyGrammar bool) {
  523. C.gpt_sampler_caccept(s.c, C.llama_token(id), C.bool(applyGrammar))
  524. }