forked from third-party-mirrors/ollama
embeddings
This commit is contained in:
jmorganca
parent
ce15ed6d69
commit
b2ef3bf490
@ -55,6 +55,7 @@ func NewContextParams() ContextParams {
|
||||
params.n_ctx = C.uint(2048)
|
||||
params.n_threads = C.uint(runtime.NumCPU())
|
||||
params.n_threads_batch = params.n_threads
|
||||
params.embeddings = C.bool(true)
|
||||
return ContextParams{c: params}
|
||||
}
|
||||
|
||||
@ -124,6 +125,20 @@ 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)))
|
||||
}
|
||||
|
||||
// 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 {
|
||||
return unsafe.Slice((*float32)(unsafe.Pointer(C.llama_get_embeddings_ith(c.c, C.int32_t(i)))), c.Model().NEmbd())
|
||||
}
|
||||
|
||||
func LoadModelFromFile(modelPath string, params ModelParams) *Model {
|
||||
return &Model{c: C.llama_load_model_from_file(C.CString(modelPath), params.c)}
|
||||
}
|
||||
@ -225,6 +240,10 @@ func (m *Model) Tokenize(text string, maxTokens int, addSpecial bool, parseSpeci
|
||||
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))
|
||||
|
||||
@ -23,9 +23,16 @@ type Sequence struct {
|
||||
// tokens left to evaluate
|
||||
tokens []int
|
||||
|
||||
// channel to send responses over
|
||||
responses chan string
|
||||
|
||||
samplingCtx *llama.SamplingContext
|
||||
|
||||
// channel to send back the embedding if embedding only
|
||||
embedding chan []float32
|
||||
|
||||
// true if an embedding are to be returned instead of text generation
|
||||
embeddingOnly bool
|
||||
}
|
||||
|
||||
// prompt returns true if the prompt is still being processed
|
||||
@ -33,38 +40,26 @@ func (s *Sequence) prompt() bool {
|
||||
return s.nPast < len(s.tokens)-1
|
||||
}
|
||||
|
||||
func (s *Server) NewSequence(r Request, w http.ResponseWriter) *Sequence {
|
||||
var samplingParams llama.SamplingParams
|
||||
samplingParams.TopK = r.TopK
|
||||
samplingParams.TopP = r.TopP
|
||||
samplingParams.TfsZ = r.TFSZ
|
||||
samplingParams.TypicalP = r.TypicalP
|
||||
samplingParams.Temp = r.Temperature
|
||||
samplingParams.PenaltyRepeat = r.RepeatPenalty
|
||||
samplingParams.PenaltyFreq = r.FrequencyPenalty
|
||||
samplingParams.PenaltyPresent = r.PresencePenalty
|
||||
samplingParams.Mirostat = r.Mirostat
|
||||
samplingParams.MirostatTau = r.MirostatTau
|
||||
samplingParams.MirostatEta = r.MirostatEta
|
||||
samplingParams.PenalizeNl = r.PenalizeNewline
|
||||
samplingParams.Seed = uint32(r.Seed)
|
||||
samplingParams.Grammar = r.Grammar
|
||||
|
||||
tokens, err := s.lc.Model().Tokenize(r.Prompt, 2048, false, true)
|
||||
func (s *Server) NewSequence(prompt string, params *llama.SamplingParams, embedding bool) *Sequence {
|
||||
tokens, err := s.lc.Model().Tokenize(prompt, 2048, false, true)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
sc := llama.NewSamplingContext(samplingParams)
|
||||
|
||||
for _, t := range tokens {
|
||||
sc.Accept(s.lc, t, false)
|
||||
var sc *llama.SamplingContext
|
||||
if params != nil {
|
||||
sc = llama.NewSamplingContext(*params)
|
||||
for _, t := range tokens {
|
||||
sc.Accept(s.lc, t, false)
|
||||
}
|
||||
}
|
||||
|
||||
return &Sequence{
|
||||
tokens: tokens,
|
||||
responses: make(chan string, 1),
|
||||
samplingCtx: sc,
|
||||
tokens: tokens,
|
||||
responses: make(chan string, 1),
|
||||
embedding: make(chan []float32, 1),
|
||||
samplingCtx: sc,
|
||||
embeddingOnly: embedding,
|
||||
}
|
||||
}
|
||||
|
||||
@ -152,6 +147,20 @@ func (s *Server) run(ctx context.Context) {
|
||||
continue
|
||||
}
|
||||
|
||||
// if done processing the prompt, generating an embedding and return
|
||||
if seq.embeddingOnly {
|
||||
embd := s.lc.GetEmbeddingsSeq(i)
|
||||
if embd == nil {
|
||||
embd = s.lc.GetEmbeddingsIth(ibatch[i])
|
||||
}
|
||||
|
||||
seq.embedding <- embd
|
||||
close(seq.embedding)
|
||||
s.lc.KvCacheSeqRm(i, 0, -1)
|
||||
s.seqs[i] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
// sample a token
|
||||
// logits := s.lc.GetLogitsIth(ibatch[i])
|
||||
// token := s.lc.SampleTokenGreedy(logits)
|
||||
@ -178,7 +187,7 @@ func (s *Server) run(ctx context.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
type Request struct {
|
||||
type CompletionRequest struct {
|
||||
Prompt string `json:"prompt"`
|
||||
Images []string `json:"images"`
|
||||
Grammar string `json:"grammar"`
|
||||
@ -186,14 +195,14 @@ type Request struct {
|
||||
api.Options
|
||||
}
|
||||
|
||||
type Response struct {
|
||||
type CompletionResponse struct {
|
||||
Token string `json:"token"`
|
||||
}
|
||||
|
||||
func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
|
||||
var request Request
|
||||
request.Options = api.DefaultOptions()
|
||||
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
||||
func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
|
||||
var req CompletionRequest
|
||||
req.Options = api.DefaultOptions()
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "Bad request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
@ -203,8 +212,26 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Transfer-Encoding", "chunked")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
seq := s.NewSequence(request, w)
|
||||
var samplingParams llama.SamplingParams
|
||||
samplingParams.TopK = req.TopK
|
||||
samplingParams.TopP = req.TopP
|
||||
samplingParams.TfsZ = req.TFSZ
|
||||
samplingParams.TypicalP = req.TypicalP
|
||||
samplingParams.Temp = req.Temperature
|
||||
samplingParams.PenaltyRepeat = req.RepeatPenalty
|
||||
samplingParams.PenaltyFreq = req.FrequencyPenalty
|
||||
samplingParams.PenaltyPresent = req.PresencePenalty
|
||||
samplingParams.Mirostat = req.Mirostat
|
||||
samplingParams.MirostatTau = req.MirostatTau
|
||||
samplingParams.MirostatEta = req.MirostatEta
|
||||
samplingParams.PenalizeNl = req.PenalizeNewline
|
||||
samplingParams.Seed = uint32(req.Seed)
|
||||
samplingParams.Grammar = req.Grammar
|
||||
|
||||
seq := s.NewSequence(req.Prompt, &samplingParams, false)
|
||||
|
||||
// TODO (jmorganca): add to sequence queue instead of
|
||||
// failing if a slot isn't available
|
||||
s.mu.Lock()
|
||||
for i, sq := range s.seqs {
|
||||
if sq == nil {
|
||||
@ -215,8 +242,9 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// stream the response
|
||||
for token := range seq.responses {
|
||||
if err := json.NewEncoder(w).Encode(&Response{
|
||||
if err := json.NewEncoder(w).Encode(&CompletionResponse{
|
||||
Token: token,
|
||||
}); err != nil {
|
||||
log.Println("Failed to encode result:", err)
|
||||
@ -233,6 +261,46 @@ func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
type EmbeddingRequest struct {
|
||||
Prompt string `json:"prompt"`
|
||||
}
|
||||
|
||||
type EmbeddingResponse struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
}
|
||||
|
||||
// TODO (jmorganca): is it safe to do this concurrently with decoding?
|
||||
func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
|
||||
var req EmbeddingRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "Bad request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
|
||||
seq := s.NewSequence(req.Prompt, nil, true)
|
||||
|
||||
s.mu.Lock()
|
||||
for i, sq := range s.seqs {
|
||||
if sq == nil {
|
||||
s.seqs[i] = seq
|
||||
s.cond.Signal()
|
||||
break
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
embedding := <-seq.embedding
|
||||
|
||||
if err := json.NewEncoder(w).Encode(&EmbeddingResponse{
|
||||
Embedding: embedding,
|
||||
}); err != nil {
|
||||
log.Println("Failed to encode result:", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func main() {
|
||||
mpath := flag.String("model", "", "Path to model binary file")
|
||||
ppath := flag.String("projector", "", "Path to projector binary file")
|
||||
@ -279,8 +347,12 @@ func main() {
|
||||
}
|
||||
defer listener.Close()
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embeddings", server.embeddings)
|
||||
mux.HandleFunc("/completion", server.completion)
|
||||
|
||||
httpServer := http.Server{
|
||||
Handler: http.HandlerFunc(server.handler),
|
||||
Handler: mux,
|
||||
}
|
||||
|
||||
log.Println("Server listening on", addr)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user