Repository: ollama/ollama
Pull request: Embeddings fixes


Changed files:
filename: llama/llama.go
patch: @@ -429,7 +429,7 @@ func NewSamplingContext(params SamplingParams) *SamplingContext {
 }
 
 func (s *SamplingContext) Free() {
-	if s.c != nil {
+	if s != nil {
 		C.llama_sampling_cfree(s.c)
 	}
 }

filename: llama/runner/runner.go
patch: @@ -65,10 +65,10 @@ type Sequence struct {
 	doneReason string
 
 	// Metrics
-	t_start_process_prompt time.Time
-	t_start_genereration   time.Time
-	n_decoded              int
-	n_prompt_tokens        int
+	startProcessingTime time.Time
+	startGenerationTime time.Time
+	numDecoded          int
+	numPromptTokens     int
 }
 
 type NewSequenceParams struct {
@@ -88,9 +88,15 @@ func (s *Server) NewSequence(prompt string, params NewSequenceParams) *Sequence
 	if params.numKeep < 0 {
 		params.numKeep = len(tokens)
 	}
-	// Subtracting 4 ensures that at least 1 token can be discarded during shift
-	params.numKeep = min(params.numKeep, s.numCtx-4)
-	params.numKeep += s.bosToken
+
+	if !params.embedding {
+		// Subtracting 4 ensures that at least 1 token can be discarded during shift
+		params.numKeep = min(params.numKeep, s.numCtx-4)
+		params.numKeep += s.bosToken
+	} else {
+		// Embeddings are 1 shot - just truncate to the context window, without ever shifting
+		params.numKeep = min(params.numKeep, s.numCtx)
+	}
 
 	// truncate to fit in context window
 	if len(tokens) > s.numCtx {
@@ -110,7 +116,7 @@ func (s *Server) NewSequence(prompt string, params NewSequenceParams) *Sequence
 
 	return &Sequence{
 		tokens:           tokens,
-		n_prompt_tokens:  len(tokens),
+		numPromptTokens:  len(tokens),
 		numPredict:       params.numPredict,
 		pendingResponses: make([]string, 0),
 		responses:        make(chan string, 1),
@@ -256,8 +262,8 @@ func (s *Server) processBatch() {
 			s.shiftContext(i)
 		}
 
-		if seq.t_start_process_prompt.IsZero() {
-			seq.t_start_process_prompt = time.Now()
+		if seq.startProcessingTime.IsZero() {
+			seq.startProcessingTime = time.Now()
 		}
 
 		var numTokensProcessed int
@@ -294,7 +300,7 @@ func (s *Server) processBatch() {
 			continue
 		}
 
-		// if done processing the prompt, generating an embedding and return
+		// if done processing the prompt, generate an embedding and return
 		if seq.embeddingOnly {
 			embd := s.lc.GetEmbeddingsSeq(i)
 			if embd == nil {
@@ -307,14 +313,12 @@ func (s *Server) processBatch() {
 		}
 
 		// sample a token
-		// logits := s.lc.GetLogitsIth(ibatch[i])
-		// token := s.lc.SampleTokenGreedy(logits)
 		token := seq.samplingCtx.Sample(s.lc, nil, seq.iBatch)
 
 		seq.samplingCtx.Accept(s.lc, token, true)
-		seq.n_decoded += 1
-		if seq.n_decoded == 1 {
-			seq.t_start_genereration = time.Now()
+		seq.numDecoded += 1
+		if seq.numDecoded == 1 {
+			seq.startGenerationTime = time.Now()
 		}
 		piece := s.model.TokenToPiece(token)
 
@@ -505,10 +509,10 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	if err := json.NewEncoder(w).Encode(&CompletionResponse{
 		Stop: true,
 		Timings: Timings{
-			PromptN:     seq.n_prompt_tokens,
-			PromptMS:    float64(seq.t_start_genereration.Sub(seq.t_start_process_prompt).Milliseconds()),
-			PredictedN:  seq.n_decoded,
-			PredictedMS: float64(time.Since(seq.t_start_genereration).Milliseconds()),
+			PromptN:     seq.numPromptTokens,
+			PromptMS:    float64(seq.startGenerationTime.Sub(seq.startProcessingTime).Milliseconds()),
+			PredictedN:  seq.numDecoded,
+			PredictedMS: float64(time.Since(seq.startGenerationTime).Milliseconds()),
 		},
 	}); err != nil {
 		log.Println("Failed to encode result:", err)
@@ -525,14 +529,13 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 }
 
 type EmbeddingRequest struct {
-	Content []string `json:"content"`
+	Content string `json:"content"`
 }
 
 type EmbeddingResponse struct {
-	Embedding [][]float32 `json:"embedding"`
+	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 {
@@ -543,36 +546,24 @@ func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
 
 	slog.Debug("embedding request", "content", req.Content)
-	seqs := make([]*Sequence, len(req.Content))
-	embeddings := make([][]float32, len(req.Content))
-	var processed int
-	for i, content := range req.Content {
-		seqs[i] = s.NewSequence(content, NewSequenceParams{embedding: true})
-	}
-
-	// TODO - refactor to go routines to add seq's and drain the responses
-	// so we don't stall until each set is iterated through
-	for processed < len(seqs) {
-		s.mu.Lock()
-		for i, sq := range s.seqs {
-			if processed >= len(seqs) {
-				break
-			}
-			if sq == nil {
-				s.seqs[i] = seqs[processed]
-				processed += 1
-			}
-		}
-		s.cond.Signal()
-		s.mu.Unlock()
 
-		for i := range processed {
-			embeddings[i] = <-seqs[i].embedding
+	seq := s.NewSequence(req.Content, NewSequenceParams{embedding: true})
+
+	// TODO (jessegross): Wait for a free slot instead of failing and blocking forever
+	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: embeddings,
+		Embedding: embedding,
 	}); err != nil {
 		log.Println("Failed to encode result:", err)
 		return
@@ -584,7 +575,7 @@ type HealthResponse struct {
 	Progress float32 `json:"progress"`
 }
 
-// TODO (jmorganca): is it safe to do this concurrently with decoding?
+// TODO (jmorganca): is it safe to do this concurrently with updating status?
 func (s *Server) health(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
 	if err := json.NewEncoder(w).Encode(&HealthResponse{
@@ -638,7 +629,7 @@ func main() {
 
 	// TODO actually implement...
 	if *embedding {
-		slog.Warn("embeddings not yet support")
+		slog.Warn("embeddings not yet supported")
 	}
 	if *logDisable {
 		slog.Info("ignoring --log-disable")
@@ -661,9 +652,11 @@ func main() {
 		batchSize: *batchSize,
 		parallel:  *parallel,
 		seqs:      make([]*Sequence, *parallel),
-		status:    "loading",
+		status:    "loading model",
 	}
 
+	// TODO (jessegross): This should be in a separate goroutine so we can report progress,
+	// otherwise Ollama can timeout for large model loads
 	// load the model
 	llama.BackendInit()
 	params := llama.NewModelParams(*nGpuLayers, *mainGpu, func(progress float32) {