jeans/ollama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

runner: avoid buffer overwrite when generating multiple embeddings (#8714)

Browse Source 
Shield the code processing the embedding result
from subsequent calls that may overwrite the same
buffer to process a second input when retrieving
model embeddings.
This commit is contained in:
Diego Pereira 2025-02-05 21:53:33 -03:00 committed by GitHub
parent 5b446cc815
commit 928911bc68
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
llama/llama.go
View File

@ -199,21 +199,25 @@ func (c *Context) KvCacheDefrag() {
// Get the embeddings for a sequence id // Get the embeddings for a sequence id
func (c *Context) GetEmbeddingsSeq(seqId int) []float32 { func (c *Context) GetEmbeddingsSeq(seqId int) []float32 {
embeddings := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId))) e := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId)))
if embeddings == nil { if e == nil {
return nil return nil
} }
return unsafe.Slice((*float32)(embeddings), c.Model().NEmbd()) embeddings := make([]float32, c.Model().NEmbd())
_ = copy(embeddings, unsafe.Slice((*float32)(e), c.Model().NEmbd()))
return embeddings
} }
func (c *Context) GetEmbeddingsIth(i int) []float32 { func (c *Context) GetEmbeddingsIth(i int) []float32 {
embeddings := unsafe.Pointer(C.llama_get_embeddings_ith(c.c, C.int32_t(i))) e := unsafe.Pointer(C.llama_get_embeddings_ith(c.c, C.int32_t(i)))
if embeddings == nil { if e == nil {
return nil return nil
} }
return unsafe.Slice((*float32)(embeddings), c.Model().NEmbd()) embeddings := make([]float32, c.Model().NEmbd())
_ = copy(embeddings, unsafe.Slice((*float32)(e), c.Model().NEmbd()))
return embeddings
} }
type ModelParams struct { type ModelParams struct {