From ffb6f220d0d8ca5bbad02ce451e26240f09f2d03 Mon Sep 17 00:00:00 2001 From: Parth Sareen Date: Mon, 28 Sep 2026 13:21:01 -0700 Subject: [PATCH] feat: add System One scoring API (#18606) --- decision/systemone.go | 228 ++++++++++++++++++ decision/systemone_test.go | 231 ++++++++++++++++++ decision/types.go | 76 ++++++ llm/llama_server_score.go | 281 ++++++++++++++++++++++ llm/llama_server_score_test.go | 421 +++++++++++++++++++++++++++++++++ llm/score.go | 27 +++ server/routes.go | 89 +++++++ server/systemone_test.go | 204 ++++++++++++++++ 8 files changed, 1557 insertions(+) create mode 100644 decision/systemone.go create mode 100644 decision/systemone_test.go create mode 100644 decision/types.go create mode 100644 llm/llama_server_score.go create mode 100644 llm/llama_server_score_test.go create mode 100644 llm/score.go create mode 100644 server/systemone_test.go diff --git a/decision/systemone.go b/decision/systemone.go new file mode 100644 index 000000000..8f07da7e2 --- /dev/null +++ b/decision/systemone.go @@ -0,0 +1,228 @@ +// Package decision compiles typed decision requests into candidate-scoring +// prompts and converts the resulting logits into /v1/systemone responses. +package decision + +import ( + "bytes" + "encoding/json" + "fmt" + "math" + "slices" + "strconv" + "strings" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/internal/orderedmap" + "github.com/ollama/ollama/llm" +) + +type compiledField struct { + Field + typ string +} + +type Compiled struct { + Request llm.ScoreRequest + fields []compiledField + messages [][]api.Message +} + +func Compile(req Request) (*Compiled, error) { + if strings.TrimSpace(req.Model) == "" { + return nil, fmt.Errorf("model is required") + } + if req.Questions.Len() < 1 || req.Questions.Len() > 64 { + return nil, fmt.Errorf("questions must contain 1–64 fields") + } + context, err := content(req.State) + if err != nil { + return nil, fmt.Errorf("state: %w", err) + } + if strings.TrimSpace(context) == "" { + return nil, fmt.Errorf("state must not be empty") + } + c := &Compiled{} + for name, q := range req.Questions.All() { + f, err := compileField(name, q) + if err != nil { + return nil, fmt.Errorf("question %q: %w", name, err) + } + c.fields = append(c.fields, f) + } + data, err := json.Marshal(struct { + Context string `json:"context"` + Schema []compiledField `json:"schema"` + }{context, c.fields}) + if err != nil { + return nil, err + } + for _, f := range c.fields { + name, err := json.Marshal(f.Name) + if err != nil { + return nil, err + } + c.messages = append(c.messages, []api.Message{ + {Role: "user", Content: string(data) + "\n\nRequested field: " + string(name)}, + }) + var row llm.ScoreRow + for _, choice := range f.Choices { + row.Candidates = append(row.Candidates, choice.Code) + } + c.Request.Rows = append(c.Request.Rows, row) + } + return c, nil +} + +// Render prepares scoring prompts using the model's system prompt and chat template. +// The caller must set Request.MaxTokens to the loaded context size before scoring. +func (c *Compiled) Render(render func([]api.Message) (string, error)) error { + for i, messages := range c.messages { + prompt, err := render(messages) + if err != nil { + return err + } + c.Request.Rows[i].Prompt = prompt + } + return nil +} + +func content(raw json.RawMessage) (string, error) { + raw = bytes.TrimSpace(raw) + if len(raw) == 0 { + return "", fmt.Errorf("must be a string, object, or array") + } + switch raw[0] { + case '"': + var s string + err := json.Unmarshal(raw, &s) + return s, err + case '{', '[': + var b bytes.Buffer + err := json.Compact(&b, raw) + return b.String(), err + default: + return "", fmt.Errorf("must be a string, object, or array") + } +} + +func compileField(name string, q Question) (compiledField, error) { + f := compiledField{Field: Field{Name: name}, typ: q.Type} + if strings.TrimSpace(name) == "" { + return f, fmt.Errorf("field name must not be empty") + } + description, err := content(q.Instructions) + if err != nil || strings.TrimSpace(description) == "" { + return f, fmt.Errorf("instructions must be a nonempty string, object, or array") + } + f.Description = description + add := func(value any, description string) { + f.Choices = append(f.Choices, Choice{string(rune('A' + len(f.Choices))), value, description}) + } + switch q.Type { + case "noul": + criteria := orderedmap.New[string, *string]() + if len(q.Criteria) > 0 { + if err := json.Unmarshal(q.Criteria, criteria); err != nil || string(q.Criteria) == "null" { + return f, fmt.Errorf("noul criteria must be an object of true/false descriptions") + } + } + no, yes := "No", "Yes" + for key, value := range criteria.All() { + if value == nil { + return f, fmt.Errorf("noul descriptions must be strings") + } + switch key { + case "false": + no = *value + case "true": + yes = *value + default: + return f, fmt.Errorf("unknown noul criterion %q", key) + } + } + add(false, no) + add(true, yes) + case "choice": + criteria := orderedmap.New[string, *string]() + if err := json.Unmarshal(q.Criteria, criteria); err != nil { + return f, fmt.Errorf("choice criteria must map option keys to descriptions or null") + } + for key, value := range criteria.All() { + if strings.TrimSpace(key) == "" { + return f, fmt.Errorf("choice keys must not be empty") + } + description := key + if value != nil { + description = *value + } + add(key, description) + } + case "score": + var criteria []*string + if err := json.Unmarshal(q.Criteria, &criteria); err != nil { + return f, fmt.Errorf("score criteria must be an array of descriptions") + } + for i, description := range criteria { + if description == nil { + return f, fmt.Errorf("score descriptions must be strings") + } + add(strconv.Itoa(i), *description) + } + default: + return f, fmt.Errorf("type must be choice, noul, or score") + } + if len(f.Choices) < 2 || len(f.Choices) > 26 { + return f, fmt.Errorf("criteria must contain 2–26 candidates") + } + return f, nil +} + +func (c *Compiled) Answer(model string, result llm.ScoreResponse) (Response, error) { + response := Response{Model: model, Answers: &Answers{}, Usage: Usage{InputTokens: result.InputTokens, OutputTokens: result.OutputTokens}} + if len(result.Logits) != len(c.fields) { + return response, fmt.Errorf("scorer returned %d rows for %d questions", len(result.Logits), len(c.fields)) + } + for i, f := range c.fields { + logits := result.Logits[i] + if len(logits) != len(f.Choices) { + return response, fmt.Errorf("scorer returned the wrong number of candidates for %q", f.Name) + } + peak := slices.Max(logits) + p := make([]float64, len(logits)) + var sum float64 + for j, logit := range logits { + if math.IsNaN(float64(logit)) || math.IsInf(float64(logit), 0) { + return response, fmt.Errorf("scorer returned a non-finite logit for %q", f.Name) + } + p[j] = math.Exp(float64(logit) - float64(peak)) + sum += p[j] + } + var entropy, score float64 + for j := range p { + p[j] /= sum + score += float64(j) * p[j] + if p[j] > 0 { + entropy -= p[j] * math.Log(p[j]) + } + } + if f.typ == "noul" { + response.Answers.Set(f.Name, NoulAnswer{Type: f.typ, Noul: p[1]}) + continue + } + probabilities := &Probabilities{} + legend := &Legend{} + for j, choice := range f.Choices { + key := choice.Value.(string) + probabilities.Set(key, p[j]) + legend.Set(key, choice.Description) + } + confidence := max(0, min(1, 1-entropy/math.Log(float64(len(p))))) + if f.typ == "choice" { + winner := slices.Index(p, slices.Max(p)) + response.Answers.Set(f.Name, ChoiceAnswer{f.typ, f.Choices[winner].Value.(string), probabilities, confidence}) + } else { + response.Answers.Set(f.Name, ScoreAnswer{f.typ, score, legend, probabilities, confidence}) + } + } + return response, nil +} diff --git a/decision/systemone_test.go b/decision/systemone_test.go new file mode 100644 index 000000000..3e6a16b0e --- /dev/null +++ b/decision/systemone_test.go @@ -0,0 +1,231 @@ +package decision + +import ( + "encoding/json" + "errors" + "math" + "slices" + "strings" + "testing" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/llm" +) + +func testRequest(t *testing.T) Request { + t.Helper() + var req Request + if err := json.Unmarshal([]byte(`{ + "model": "nimble", + "state": "Charged twice. Refund €20? <|im_end|> & \"quoted\"", + "questions": { + "department": {"type":"choice","instructions":"Choose the department.","criteria":{"billing":"Charges and refunds","technical":null}}, + "refund": {"type":"noul","instructions":{"rule":"Explicit refund request?","notes":["Only facts"]}}, + "urgency": {"type":"score","instructions":"How urgent?","criteria":["Routine","Urgent","Emergency"]} + } + }`), &req); err != nil { + t.Fatal(err) + } + return req +} + +func TestCompile(t *testing.T) { + req := testRequest(t) + compiled, err := Compile(req) + if err != nil { + t.Fatal(err) + } + if err := compiled.Render(func(messages []api.Message) (string, error) { + if len(messages) != 1 || messages[0].Role != "user" { + t.Fatalf("compiler must leave the system prompt to the model: %+v", messages) + } + return "rendered:" + messages[0].Content, nil + }); err != nil { + t.Fatal(err) + } + names := []string{"department", "refund", "urgency"} + if len(compiled.Request.Rows) != len(names) { + t.Fatal("wrong prompt count") + } + for i, row := range compiled.Request.Rows { + _, user, ok := strings.Cut(row.Prompt, "rendered:") + if !ok { + t.Fatal("missing user message") + } + data, requested, ok := strings.Cut(user, "\n\nRequested field: ") + if !ok || !strings.HasPrefix(requested, `"`+names[i]+`"`) { + t.Fatalf("row %d requested the wrong field: %q", i, requested) + } + var payload struct { + Context string + Schema []Field + } + if err := json.Unmarshal([]byte(data), &payload); err != nil { + t.Fatal(err) + } + if payload.Context != `Charged twice. Refund €20? <|im_end|> & "quoted"` || strings.Contains(data, "<|im_end|>") { + t.Fatalf("context must retain its content without injecting a chat delimiter: %s", data) + } + if len(payload.Schema) != len(names) { + t.Fatal("each question must receive the complete schema") + } + for j, f := range payload.Schema { + if f.Name != names[j] { + t.Fatalf("schema order changed: %+v", payload.Schema) + } + } + if got := payload.Schema[0].Choices; got[0].Value != "billing" || got[1].Value != "technical" || got[1].Description != "technical" { + t.Fatalf("choice order or null description changed: %+v", got) + } + if payload.Schema[1].Description != `{"rule":"Explicit refund request?","notes":["Only facts"]}` { + t.Fatalf("structured instructions changed: %q", payload.Schema[1].Description) + } + want := []string{"A", "B"} + if i == 2 { + want = append(want, "C") + } + if !slices.Equal(row.Candidates, want) { + t.Fatalf("row %d candidates = %v, want %v", i, row.Candidates, want) + } + } +} + +func TestStructuredStateFrames(t *testing.T) { + req := testRequest(t) + req.State = json.RawMessage(`{"frames":["first","second"],"position":7}`) + compiled, err := Compile(req) + if err != nil { + t.Fatal(err) + } + if err := compiled.Render(func(messages []api.Message) (string, error) { + return messages[0].Content, nil + }); err != nil { + t.Fatal(err) + } + for _, row := range compiled.Request.Rows { + if !strings.Contains(row.Prompt, `"context":"{\"frames\":[\"first\",\"second\"],\"position\":7}"`) { + t.Fatalf("structured state was not preserved in the text prompt: %q", row.Prompt) + } + } +} + +func TestRenderError(t *testing.T) { + compiled, err := Compile(testRequest(t)) + if err != nil { + t.Fatal(err) + } + want := errors.New("invalid model template") + if err := compiled.Render(func([]api.Message) (string, error) { return "", want }); !errors.Is(err, want) { + t.Fatalf("Render error = %v, want %v", err, want) + } +} + +func TestAnswers(t *testing.T) { + req := testRequest(t) + c, err := Compile(req) + if err != nil { + t.Fatal(err) + } + result, err := c.Answer("nimble", llm.ScoreResponse{ + Logits: [][]float32{{0, 0}, {-1000, 1000}, {1000, 1000, 1000}}, InputTokens: 900, + }) + if err != nil { + t.Fatal(err) + } + choice, _ := result.Answers.Get(c.fields[0].Name) + got := choice.(ChoiceAnswer) + if got.Choice != "billing" || got.Confidence != 0 { + t.Fatalf("tie must choose the first candidate with zero concentration: %+v", got) + } + noul, _ := result.Answers.Get("refund") + if noul.(NoulAnswer).Noul != 1 { + t.Fatal("noul must report P(true), using stable softmax") + } + score, _ := result.Answers.Get("urgency") + if math.Abs(score.(ScoreAnswer).Score-1) > 1e-12 || score.(ScoreAnswer).Legend.Len() != 3 { + t.Fatalf("bad expected rubric index or legend: %+v", score) + } + if result.Usage.InputTokens != 900 || result.Usage.OutputTokens != 0 { + t.Fatalf("bad direct-scoring usage: %+v", result.Usage) + } + for _, logits := range [][][]float32{nil, {{1}, {1, 2}, {1, 2, 3}}, {{float32(math.NaN()), 0}, {0, 1}, {1, 2, 3}}} { + if _, err := c.Answer("nimble", llm.ScoreResponse{Logits: logits}); err == nil { + t.Errorf("accepted malformed runner result: %v", logits) + } + } +} + +func TestAnswerLogprobs(t *testing.T) { + c, err := Compile(testRequest(t)) + if err != nil { + t.Fatal(err) + } + logits := [][]float32{{0, 2}, {0, -2}, {0, 2, 4}} + logprobs := [][]float32{{-10, -8}, {-10, -12}, {-10, -8, -6}} + want, err := c.Answer("nimble", llm.ScoreResponse{Logits: logits, InputTokens: 123}) + if err != nil { + t.Fatal(err) + } + got, err := c.Answer("nimble", llm.ScoreResponse{Logits: logprobs, InputTokens: 123, OutputTokens: 7}) + if err != nil { + t.Fatal(err) + } + wantAnswers, _ := json.Marshal(want.Answers) + gotAnswers, _ := json.Marshal(got.Answers) + if string(wantAnswers) != string(gotAnswers) { + t.Fatalf("row offsets changed candidate probabilities: %s != %s", gotAnswers, wantAnswers) + } + if got.Usage.InputTokens != 123 || got.Usage.OutputTokens != 7 || want.Usage.OutputTokens != 0 { + t.Fatalf("incorrect backend usage: llama=%+v direct=%+v", got.Usage, want.Usage) + } +} + +func TestInvalidRequests(t *testing.T) { + for _, data := range []string{ + `{}`, `{"model":"nimble","state":"x","questions":{}}`, + `{"model":"nimble","state":" ","questions":{"x":{"type":"noul","instructions":"q"}}}`, + `{"model":"nimble","state":null,"questions":{"x":{"type":"noul","instructions":"q"}}}`, + `{"model":"nimble","state":"x","questions":{"x":{"type":"other","instructions":"q"}}}`, + `{"model":"nimble","state":"x","questions":{"x":{"type":"noul"}}}`, + `{"model":"nimble","state":"x","questions":{"":{"type":"noul","instructions":"q"}}}`, + `{"model":"nimble","state":"x","questions":{"x":{"type":"choice","instructions":"q","criteria":{"a":"a"}}}}`, + `{"model":"nimble","state":"x","questions":{"x":{"type":"score","instructions":"q","criteria":["a",null]}}}`, + `{"model":"nimble","state":"x","questions":{"x":{"type":"noul","instructions":"q","criteria":{"true":null}}}}`, + `{"model":"nimble","state":"x","questions":{"x":{"type":"noul","instructions":"q","criteria":{"yes":"Yes"}}}}`, + } { + t.Run(data, func(t *testing.T) { + var req Request + if err := json.Unmarshal([]byte(data), &req); err != nil { + t.Fatal(err) + } + if _, err := Compile(req); err == nil { + t.Fatal("accepted invalid request") + } + }) + } + req := testRequest(t) + q, _ := req.Questions.Get("urgency") + q.Criteria = json.RawMessage(`["x"` + strings.Repeat(`,"x"`, 26) + `]`) + req.Questions.Set("urgency", q) + if _, err := Compile(req); err == nil { + t.Fatal("accepted 27 candidates") + } +} + +func TestContent(t *testing.T) { + for _, tt := range []struct{ input, want string }{ + {` "line\nquoted \"text\"" `, "line\nquoted \"text\""}, + {` {"z": 9007199254740993, "a": 1e999} `, `{"z":9007199254740993,"a":1e999}`}, + {` [1.0, 1e6, -0, "€"] `, `[1.0,1e6,-0,"€"]`}, + } { + got, err := content(json.RawMessage(tt.input)) + if err != nil || got != tt.want { + t.Errorf("content(%s) = %q, %v; want %q", tt.input, got, err, tt.want) + } + } + for _, input := range []string{"", "null", "true", "42", `"unterminated`, `{"x":}`, `[1,]`} { + if _, err := content(json.RawMessage(input)); err == nil { + t.Errorf("accepted invalid content %q", input) + } + } +} diff --git a/decision/types.go b/decision/types.go new file mode 100644 index 000000000..a3dfb6481 --- /dev/null +++ b/decision/types.go @@ -0,0 +1,76 @@ +package decision + +import ( + "encoding/json" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/internal/orderedmap" +) + +// Questions preserves request order, which determines the order of scored fields. +// Its zero value is ready to use with Set. +type Questions = orderedmap.Map[string, Question] + +type Request struct { + Model string `json:"model"` + State json.RawMessage `json:"state"` + Questions *Questions `json:"questions"` + KeepAlive *api.Duration `json:"keep_alive,omitempty"` +} + +type Question struct { + Type string `json:"type"` + Instructions json.RawMessage `json:"instructions"` + Criteria json.RawMessage `json:"criteria"` +} + +// Choice maps a single-token answer code to a value in the question's schema. +type Choice struct { + Code string `json:"code"` + Value any `json:"value"` + Description string `json:"description"` +} + +// Field is a question's schema as presented to the model. +type Field struct { + Name string `json:"name"` + Description string `json:"description"` + Choices []Choice `json:"choices"` +} + +type ( + Answers = orderedmap.Map[string, any] + Probabilities = orderedmap.Map[string, float64] + Legend = orderedmap.Map[string, string] +) + +type Response struct { + Model string `json:"model"` + Answers *Answers `json:"answers"` + Usage Usage `json:"usage"` +} + +type Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` +} + +type NoulAnswer struct { + Type string `json:"type"` + Noul float64 `json:"noul"` +} + +type ChoiceAnswer struct { + Type string `json:"type"` + Choice string `json:"choice"` + Probabilities *Probabilities `json:"probabilities"` + Confidence float64 `json:"confidence"` +} + +type ScoreAnswer struct { + Type string `json:"type"` + Score float64 `json:"score"` + Legend *Legend `json:"legend"` + Probabilities *Probabilities `json:"probabilities"` + Confidence float64 `json:"confidence"` +} diff --git a/llm/llama_server_score.go b/llm/llama_server_score.go new file mode 100644 index 000000000..b5d64a292 --- /dev/null +++ b/llm/llama_server_score.go @@ -0,0 +1,281 @@ +package llm + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "math" + "net/http" + "slices" + + "github.com/ollama/ollama/api" +) + +var _ Scorer = (*llamaServerRunner)(nil) + +const ( + // Candidates are lifted above every other token so that top_k keeps exactly + // them; a shared bias does not change their softmax. A larger bias is + // retried if another token still outranks a candidate. + scoreBias = 100 + scoreRetryBias = 1000 +) + +// A candidate whose probability underflows float32 is omitted by llama-server. +var scoreFloorLogit = float32(math.Log(math.SmallestNonzeroFloat32)) + +type scoreRowTokens struct { + tokens, candidates []int +} + +func (s *llamaServerRunner) Score(ctx context.Context, input ScoreRequest) (ScoreResponse, error) { + var result ScoreResponse + badRequest := func(format string, args ...any) (ScoreResponse, error) { + return ScoreResponse{}, api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: fmt.Sprintf(format, args...)} + } + if len(input.Rows) < 1 || len(input.Rows) > 64 { + return badRequest("scoring requires 1–64 prompts") + } + if input.MaxTokens < 1 || input.MaxTokens > s.ContextLength() { + return badRequest("max_tokens must be between 1 and %d", s.ContextLength()) + } + if err := s.sem.Acquire(ctx, 1); err != nil { + return result, err + } + defer s.sem.Release(1) + status, err := s.getServerStatusRetry(ctx) + if err != nil { + return result, err + } else if status != ServerStatusReady { + return result, fmt.Errorf("unexpected server status: %s", status) + } + + rows := make([]scoreRowTokens, len(input.Rows)) + parseSpecial, ordinary := true, false + for i, row := range input.Rows { + if len(row.Candidates) < 1 || len(row.Candidates) > 26 { + return badRequest("prompt %d requires 1–26 candidates", i) + } + tokens, err := s.tokenize(ctx, row.Prompt, false, &parseSpecial) + if err != nil { + return result, err + } + if len(tokens) == 0 || len(tokens) > input.MaxTokens { + return badRequest("prompt %d has %d tokens; expected 1–%d (input is never truncated)", i, len(tokens), input.MaxTokens) + } + // llama-server marks the result truncated when prompt+1 reaches the + // context limit, so scoring needs two positions of headroom. + if len(tokens)+2 > s.ContextLength() { + return badRequest("prompt %d requires %d context tokens for scoring; model has %d", i, len(tokens)+2, s.ContextLength()) + } + rows[i].tokens = tokens + result.InputTokens += len(tokens) + for _, candidate := range row.Candidates { + joined, err := s.tokenize(ctx, row.Prompt+candidate, false, &parseSpecial) + if err != nil { + return result, err + } + literal, err := s.tokenize(ctx, candidate, false, &ordinary) + if err != nil { + return result, err + } + if len(joined) != len(tokens)+1 || !slices.Equal(joined[:len(tokens)], tokens) || len(literal) != 1 || literal[0] != joined[len(tokens)] { + return badRequest("candidate %q must append exactly one ordinary token to prompt %d", candidate, i) + } + id := joined[len(tokens)] + if id < 0 || slices.Contains(rows[i].candidates, id) { + return badRequest("prompt %d has invalid or duplicate candidate tokens", i) + } + rows[i].candidates = append(rows[i].candidates, id) + } + } + + result.OutputTokens, err = s.primeSharedPrefix(ctx, rows) + if err != nil { + return ScoreResponse{}, err + } + + // Rows are scored in order: the first row continues from the primer. + result.Logits = make([][]float32, len(rows)) + for i, row := range rows { + logits, outputTokens, err := s.scoreRow(ctx, row) + if err != nil { + return ScoreResponse{}, err + } + result.Logits[i] = logits + result.OutputTokens += outputTokens + } + return result, nil +} + +// llama-server saves a checkpoint of state it cannot roll back, from recurrent +// and sliding-window layers, this many tokens before the end of each prompt it +// evaluates. See https://github.com/ggml-org/llama.cpp/pull/20288. +const scoreCheckpointOffset = 4 + +// primeSharedPrefix makes llama-server save a checkpoint at the end of the +// tokens every row shares, so each row evaluates only its own last few tokens +// instead of its whole prompt. +// +// Rows share the rendered context and schema and differ only near the end. +// llama-server usually reuses a cached prefix by trimming its cache, but it +// cannot trim the state of recurrent layers, as in hybrid models like Qwen 3.5, +// or of sliding-window layers. Such models can only resume from a checkpoint, +// and the checkpoints llama-server saves on its own sit near the end of each +// prompt, past the point where the next row differs. Without a primer, every row would be +// evaluated from the start. +// +// The primer evaluates the shared tokens plus the first row's next +// scoreCheckpointOffset tokens into the cache, so llama-server saves its +// checkpoint exactly at the end of the shared tokens. The pinned server still +// generates one token with n_predict 0; it is discarded but counted in usage. +// The first row continues from the primer; later rows restore the checkpoint. +// With 8 questions, scoring is about 3x faster. +// +// Priming only changes speed: if llama-server stops saving that checkpoint, +// rows are evaluated from the start and scores stay the same. It can be +// removed once llama-server can save a checkpoint at a requested position or +// score rows as one batch. +func (s *llamaServerRunner) primeSharedPrefix(ctx context.Context, rows []scoreRowTokens) (int, error) { + if len(rows) < 2 { + return 0, nil + } + first := rows[0].tokens + shared := len(first) + for _, row := range rows[1:] { + shared = min(shared, len(row.tokens)) + for i := range shared { + if row.tokens[i] != first[i] { + shared = i + break + } + } + } + // Nothing is shared, or the first row is too short to continue from the + // primer. + if shared == 0 || len(first) <= shared+scoreCheckpointOffset { + return 0, nil + } + tokens := first[:shared+scoreCheckpointOffset] + output, err := s.scoreCompletion(ctx, scoreCompletionRequest{Prompt: tokens, CachePrompt: true, NPredict: 0}) + if err != nil { + return 0, err + } + if output.Truncated || output.TokensEvaluated != len(tokens) { + return 0, fmt.Errorf("scoring did not evaluate the complete shared prefix") + } + return output.TokensPredicted, nil +} + +// scoreCompletionRequest is the part of llama-server's /completion request +// that scoring uses. +type scoreCompletionRequest struct { + Prompt []int `json:"prompt"` + Stream bool `json:"stream"` + CachePrompt bool `json:"cache_prompt"` + NPredict int `json:"n_predict"` + NProbs int `json:"n_probs,omitempty"` + PostSamplingProbs bool `json:"post_sampling_probs"` + Samplers []string `json:"samplers,omitempty"` + TopK int `json:"top_k"` + Temperature float32 `json:"temperature"` + LogitBias [][2]float64 `json:"logit_bias,omitempty"` +} + +// scoreCompletionResponse is the part of llama-server's /completion response +// that scoring reads. +type scoreCompletionResponse struct { + TokensPredicted int `json:"tokens_predicted"` + TokensEvaluated int `json:"tokens_evaluated"` + Truncated bool `json:"truncated"` + Probabilities []struct { + TopProbs []struct { + ID *int `json:"id"` + Prob *float64 `json:"prob"` + } `json:"top_probs"` + } `json:"completion_probabilities"` +} + +// scoreRow reads every candidate from one next-token distribution. Biasing the +// candidates by the same amount keeps their relative logits, so the returned +// probabilities are the softmax over the candidates' original logits. +func (s *llamaServerRunner) scoreRow(ctx context.Context, row scoreRowTokens) ([]float32, int, error) { + var outputTokens int + for _, bias := range []float64{scoreBias, scoreRetryBias} { + req := scoreCompletionRequest{ + Prompt: row.tokens, CachePrompt: true, NPredict: 1, + NProbs: len(row.candidates), PostSamplingProbs: true, + Samplers: []string{"top_k", "temperature"}, TopK: len(row.candidates), Temperature: 1, + } + for _, id := range row.candidates { + req.LogitBias = append(req.LogitBias, [2]float64{float64(id), bias}) + } + output, err := s.scoreCompletion(ctx, req) + if err != nil { + return nil, 0, err + } + if output.Truncated || output.TokensEvaluated != len(row.tokens) || output.TokensPredicted != 1 || len(output.Probabilities) != 1 || len(output.Probabilities[0].TopProbs) == 0 { + return nil, 0, fmt.Errorf("scoring did not evaluate the complete prompt and return one distribution") + } + outputTokens += output.TokensPredicted + logits := slices.Repeat([]float32{scoreFloorLogit}, len(row.candidates)) + seen := make([]bool, len(row.candidates)) + outranked := false + for _, p := range output.Probabilities[0].TopProbs { + if p.ID == nil || p.Prob == nil || math.IsNaN(*p.Prob) || *p.Prob <= 0 || *p.Prob > 1 { + return nil, 0, fmt.Errorf("scoring returned an invalid candidate probability") + } + j := slices.Index(row.candidates, *p.ID) + if j < 0 { + outranked = true + break + } + if seen[j] { + return nil, 0, fmt.Errorf("scoring returned a candidate twice") + } + seen[j] = true + logits[j] = max(float32(math.Log(*p.Prob)), scoreFloorLogit) + } + if !outranked { + return logits, outputTokens, nil + } + } + return nil, 0, fmt.Errorf("scoring candidates were outranked by other tokens") +} + +func (s *llamaServerRunner) scoreCompletion(ctx context.Context, input scoreCompletionRequest) (scoreCompletionResponse, error) { + var output scoreCompletionResponse + data, err := json.Marshal(input) + if err != nil { + return output, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/completion", s.port), bytes.NewReader(data)) + if err != nil { + return output, err + } + req.Header.Set("Content-Type", "application/json") + res, err := s.httpClient().Do(req) + if err != nil { + if ctx.Err() != nil { + return output, ctx.Err() + } + if msg := s.lastErrMsg(); msg != "" { + return output, fmt.Errorf("scoring failed: %s: %w", msg, err) + } + return output, err + } + defer res.Body.Close() + body, err := io.ReadAll(res.Body) + if err != nil { + return output, err + } + if res.StatusCode != http.StatusOK { + return output, api.StatusError{StatusCode: res.StatusCode, ErrorMessage: s.statusErrorMessage(body)} + } + if err := json.Unmarshal(body, &output); err != nil { + return output, fmt.Errorf("invalid scoring response: %w", err) + } + return output, nil +} diff --git a/llm/llama_server_score_test.go b/llm/llama_server_score_test.go new file mode 100644 index 000000000..931eca103 --- /dev/null +++ b/llm/llama_server_score_test.go @@ -0,0 +1,421 @@ +package llm + +import ( + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "maps" + "math" + "net" + "net/http" + "net/http/httptest" + "reflect" + "slices" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/ollama/ollama/api" + "golang.org/x/sync/semaphore" +) + +func newLlamaScoreTestRunner(t *testing.T, completion http.HandlerFunc) *llamaServerRunner { + t.Helper() + mux := http.NewServeMux() + mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"status":"ok"}`) + }) + mux.HandleFunc("/tokenize", func(w http.ResponseWriter, r *http.Request) { + var req struct { + Content string `json:"content"` + AddSpecial bool `json:"add_special"` + ParseSpecial *bool `json:"parse_special"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Error(err) + return + } + if req.AddSpecial || req.ParseSpecial == nil { + t.Error("scoring must explicitly parse special tokens without adding BOS/EOS") + return + } + var tokens []int + for _, c := range req.Content { + tokens = append(tokens, int(c)) + } + if *req.ParseSpecial && strings.HasSuffix(req.Content, "") { + tokens = append(tokens[:len(tokens)-len("")], 1000) + } + // Model a tokenizer that merges the candidate with the prompt's last token. + if strings.HasSuffix(req.Content, "~") && len(tokens) > 1 { + tokens = append(tokens[:len(tokens)-2], 1001) + } + _ = json.NewEncoder(w).Encode(map[string]any{"tokens": tokens}) + }) + mux.HandleFunc("/completion", completion) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return &llamaServerRunner{ + port: srv.Listener.Addr().(*net.TCPAddr).Port, client: srv.Client(), + cmd: fakeRunningCmd(), options: api.Options{Runner: api.Runner{NumCtx: 8}}, + sem: semaphore.NewWeighted(1), + } +} + +type scoreCompletionTestRequest struct { + Prompt []int `json:"prompt"` + Grammar string `json:"grammar"` + Stream bool `json:"stream"` + CachePrompt bool `json:"cache_prompt"` + NPredict int `json:"n_predict"` + NProbs int `json:"n_probs"` + PostSamplingProbs *bool `json:"post_sampling_probs"` + Temperature float64 `json:"temperature"` + Samplers []string `json:"samplers"` + TopK int `json:"top_k"` + LogitBias [][2]float64 `json:"logit_bias"` +} + +// scoreCompletionTestResponse returns the candidates' probabilities in the +// order llama-server sorts them. +func scoreCompletionTestResponse(prompt int, probs map[int]float64) map[string]any { + ids := slices.Collect(maps.Keys(probs)) + slices.SortFunc(ids, func(a, b int) int { return cmp.Compare(probs[b], probs[a]) }) + var top []any + for _, id := range ids { + top = append(top, map[string]any{"id": id, "token": string(rune(id)), "prob": probs[id]}) + } + return map[string]any{ + "tokens_predicted": 1, "tokens_evaluated": prompt, + "completion_probabilities": []any{map[string]any{"id": ids[0], "token": string(rune(ids[0])), "prob": probs[ids[0]], "top_probs": top}}, + } +} + +func checkScoreRowRequest(t *testing.T, req scoreCompletionTestRequest, candidates []int, bias float64) { + t.Helper() + var biased []int + for _, b := range req.LogitBias { + if b[1] != bias { + t.Errorf("candidates must share bias %g: %v", bias, req.LogitBias) + } + biased = append(biased, int(b[0])) + } + if req.Stream || !req.CachePrompt || req.NPredict != 1 || req.Grammar != "" || req.NProbs != len(candidates) || req.PostSamplingProbs == nil || !*req.PostSamplingProbs || + req.Temperature != 1 || !slices.Equal(req.Samplers, []string{"top_k", "temperature"}) || req.TopK != len(candidates) || !slices.Equal(biased, candidates) { + t.Errorf("request changes the candidate distribution: %+v", req) + } +} + +func TestLlamaServerScore(t *testing.T) { + var calls int + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + var req scoreCompletionTestRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Error(err) + return + } + calls++ + switch { + case slices.Equal(req.Prompt, []int{'a', 'b', 'c'}): + checkScoreRowRequest(t, req, []int{'Z', 'A'}, scoreBias) + _ = json.NewEncoder(w).Encode(scoreCompletionTestResponse(len(req.Prompt), map[int]float64{'Z': 0.25, 'A': 0.75})) + case slices.Equal(req.Prompt, []int{'x', 'y'}): + checkScoreRowRequest(t, req, []int{'B'}, scoreBias) + _ = json.NewEncoder(w).Encode(scoreCompletionTestResponse(len(req.Prompt), map[int]float64{'B': 1})) + default: + t.Errorf("rendered prompt changed: %v", req.Prompt) + } + }) + // Model generation defaults must not affect scoring. + runner.options.Temperature = 0 + runner.options.Stop = []string{"A"} + runner.options.RepeatPenalty = 2 + runner.options.NumCtx = 5 // The longest prompt leaves exactly two positions of headroom. + result, err := runner.Score(t.Context(), ScoreRequest{MaxTokens: 5, Rows: []ScoreRow{ + {Prompt: "abc", Candidates: []string{"Z", "A"}}, + {Prompt: "xy", Candidates: []string{"B"}}, + }}) + if err != nil { + t.Fatal(err) + } + want := [][]float32{{float32(math.Log(0.25)), float32(math.Log(0.75))}, {0}} + if !reflect.DeepEqual(result.Logits, want) || result.InputTokens != 5 || result.OutputTokens != 2 || calls != 2 { + t.Fatalf("incomplete or reordered scores/usage: %+v, calls=%d", result, calls) + } +} + +func TestLlamaServerScoreSharedPrefix(t *testing.T) { + var prompts [][]int + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + var req scoreCompletionTestRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Error(err) + return + } + prompts = append(prompts, req.Prompt) + if len(prompts) == 1 { + // The native server generates one token even with n_predict: 0. + if req.NPredict != 0 || req.NProbs != 0 || len(req.LogitBias) != 0 || len(req.Samplers) != 0 || !req.CachePrompt { + t.Errorf("unexpected primer request: %+v", req) + } + _ = json.NewEncoder(w).Encode(map[string]any{"tokens_predicted": 1, "tokens_evaluated": len(req.Prompt)}) + return + } + candidate := int(req.LogitBias[0][0]) + checkScoreRowRequest(t, req, []int{candidate}, scoreBias) + _ = json.NewEncoder(w).Encode(scoreCompletionTestResponse(len(req.Prompt), map[int]float64{candidate: 1})) + }) + runner.options.NumCtx = 64 + result, err := runner.Score(t.Context(), ScoreRequest{MaxTokens: 32, Rows: []ScoreRow{ + {Prompt: "shared:first", Candidates: []string{"A"}}, + {Prompt: "shared:two", Candidates: []string{"B"}}, + {Prompt: "shared:x", Candidates: []string{"C"}}, + }}) + if err != nil { + t.Fatal(err) + } + // The primer evaluates the shared prefix plus 4 tokens of the first row, + // which continues from it; the other rows restore the checkpoint. + want := [][]int{scoreTestTokens("shared:firs"), scoreTestTokens("shared:first"), scoreTestTokens("shared:two"), scoreTestTokens("shared:x")} + if !reflect.DeepEqual(prompts, want) { + t.Fatalf("got prompts %v, want %v", prompts, want) + } + if !reflect.DeepEqual(result.Logits, [][]float32{{0}, {0}, {0}}) || result.InputTokens != 30 || result.OutputTokens != 4 { + t.Fatalf("unexpected result: %+v", result) + } +} + +func TestLlamaServerScorePrimer(t *testing.T) { + for _, tt := range []struct { + name string + rows []string + primer string + }{ + {"one row", []string{"abcdefgh"}, ""}, + {"no shared prefix", []string{"abcdefgh", "xbcdefgh"}, ""}, + {"first row too short", []string{"abcdefg", "abcdwxyzuv"}, ""}, + {"identical prompts", []string{"abcdefgh", "abcdefgh"}, ""}, + {"shared prefix", []string{"abcdefghij", "abcxyzuvwt", "abcdefgh"}, "abcdefg"}, + } { + t.Run(tt.name, func(t *testing.T) { + var primer []int + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + var req scoreCompletionTestRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Error(err) + return + } + if len(req.LogitBias) == 0 { + primer = req.Prompt + _ = json.NewEncoder(w).Encode(map[string]any{"tokens_predicted": 1, "tokens_evaluated": len(req.Prompt)}) + return + } + _ = json.NewEncoder(w).Encode(scoreCompletionTestResponse(len(req.Prompt), map[int]float64{'A': 1})) + }) + runner.options.NumCtx = 64 + input := ScoreRequest{MaxTokens: 32} + for _, row := range tt.rows { + input.Rows = append(input.Rows, ScoreRow{Prompt: row, Candidates: []string{"A"}}) + } + if _, err := runner.Score(t.Context(), input); err != nil { + t.Fatal(err) + } + if !slices.Equal(primer, scoreTestTokens(tt.primer)) { + t.Fatalf("primed %v, want %v", primer, scoreTestTokens(tt.primer)) + } + }) + } +} + +// scoreTestTokens tokenizes s the way the test tokenizer does: one token per rune. +func scoreTestTokens(s string) []int { + var tokens []int + for _, c := range s { + tokens = append(tokens, int(c)) + } + return tokens +} + +func TestLlamaServerScoreOutranked(t *testing.T) { + for _, recovers := range []bool{true, false} { + t.Run(fmt.Sprintf("recovers=%v", recovers), func(t *testing.T) { + var calls int + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + var req scoreCompletionTestRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Error(err) + return + } + calls++ + bias := float64(scoreBias) + if calls > 1 { + bias = scoreRetryBias + } + checkScoreRowRequest(t, req, []int{'A', 'B'}, bias) + probs := map[int]float64{'A': 0.6, 'x': 0.4} + if recovers && calls > 1 { + probs = map[int]float64{'A': 0.6, 'B': 0.4} + } + _ = json.NewEncoder(w).Encode(scoreCompletionTestResponse(len(req.Prompt), probs)) + }) + result, err := runner.Score(t.Context(), ScoreRequest{MaxTokens: 5, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A", "B"}}}}) + if calls != 2 { + t.Fatalf("got %d attempts, want 2", calls) + } + if !recovers { + if err == nil { + t.Fatal("accepted a distribution without every candidate") + } + return + } + if err != nil || result.OutputTokens != 2 || !reflect.DeepEqual(result.Logits, [][]float32{{float32(math.Log(0.6)), float32(math.Log(0.4))}}) { + t.Fatalf("unexpected result: %+v, %v", result, err) + } + }) + } +} + +func TestLlamaServerScoreUnderflow(t *testing.T) { + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + // llama-server omits candidates whose probability underflows. + _ = json.NewEncoder(w).Encode(scoreCompletionTestResponse(3, map[int]float64{'B': 1})) + }) + result, err := runner.Score(t.Context(), ScoreRequest{MaxTokens: 5, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A", "B"}}}}) + if err != nil || !reflect.DeepEqual(result.Logits, [][]float32{{scoreFloorLogit, 0}}) { + t.Fatalf("unexpected result: %+v, %v", result, err) + } +} + +func TestLlamaServerScoreValidation(t *testing.T) { + for _, tt := range []struct { + name string + input ScoreRequest + }{ + {"empty rows", ScoreRequest{MaxTokens: 7}}, + {"too many rows", ScoreRequest{MaxTokens: 7, Rows: make([]ScoreRow, 65)}}, + {"zero max tokens", ScoreRequest{Rows: []ScoreRow{{Prompt: "a", Candidates: []string{"A"}}}}}, + {"max exceeds context", ScoreRequest{MaxTokens: 9, Rows: []ScoreRow{{Prompt: "a", Candidates: []string{"A"}}}}}, + {"empty prompt", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Candidates: []string{"A"}}}}}, + {"oversize prompt", ScoreRequest{MaxTokens: 2, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A"}}}}}, + {"no scoring token space", ScoreRequest{MaxTokens: 8, Rows: []ScoreRow{{Prompt: "abcdefgh", Candidates: []string{"A"}}}}}, + {"no native scoring headroom", ScoreRequest{MaxTokens: 8, Rows: []ScoreRow{{Prompt: "abcdefg", Candidates: []string{"A"}}}}}, + {"empty candidates", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc"}}}}, + {"too many candidates", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: make([]string, 27)}}}}, + {"multiple tokens", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"AB"}}}}}, + {"special token", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{""}}}}}, + {"changed prefix", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"~"}}}}}, + {"duplicate token", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A", "A"}}}}}, + {"invalid later row", ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A"}}, {Prompt: "abc", Candidates: []string{"AB"}}}}}, + } { + t.Run(tt.name, func(t *testing.T) { + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + t.Error("started inference before validating every row") + }) + _, err := runner.Score(t.Context(), tt.input) + var status api.StatusError + if !errors.As(err, &status) || status.StatusCode != http.StatusBadRequest { + t.Fatalf("expected validation error, got %v", err) + } + }) + } +} + +func TestLlamaServerScoreResponse(t *testing.T) { + const valid = `{"tokens_predicted":1,"tokens_evaluated":3,"completion_probabilities":[{"id":65,"token":"A","prob":1,"top_probs":[{"id":65,"token":"A","prob":1}]}]}` + for _, tt := range []struct { + name string + body string + status int + }{ + {"missing probability", strings.Replace(valid, `"token":"A","prob":1}]`, `"token":"A"}]`, 1), 200}, + {"missing token id", strings.Replace(valid, `[{"id":65,"token":"A","prob":1}]`, `[{"token":"A","prob":1}]`, 1), 200}, + {"zero probability", strings.Replace(valid, `"token":"A","prob":1}]`, `"token":"A","prob":0}]`, 1), 200}, + {"probability above one", strings.Replace(valid, `"token":"A","prob":1}]`, `"token":"A","prob":1.5}]`, 1), 200}, + {"duplicate candidate", strings.Replace(valid, `[{"id":65,"token":"A","prob":1}]`, `[{"id":65,"token":"A","prob":0.5},{"id":65,"token":"A","prob":0.5}]`, 1), 200}, + {"empty top probabilities", strings.Replace(valid, `[{"id":65,"token":"A","prob":1}]`, `[]`, 1), 200}, + {"incomplete prompt", strings.Replace(valid, `"tokens_evaluated":3`, `"tokens_evaluated":2`, 1), 200}, + {"truncated prompt", strings.Replace(valid, `"tokens_predicted":1`, `"truncated":true,"tokens_predicted":1`, 1), 200}, + {"extra generated token", strings.Replace(valid, `"tokens_predicted":1`, `"tokens_predicted":2`, 1), 200}, + {"empty probabilities", `{"tokens_predicted":1,"tokens_evaluated":3,"completion_probabilities":[]}`, 200}, + {"malformed JSON", `{`, 200}, + {"upstream error", `{"error":{"message":"out of memory"}}`, 500}, + } { + t.Run(tt.name, func(t *testing.T) { + var failed atomic.Bool + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + if failed.CompareAndSwap(false, true) { + w.WriteHeader(tt.status) + fmt.Fprint(w, tt.body) + return + } + fmt.Fprint(w, valid) + }) + input := ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A"}}}} + if _, err := runner.Score(t.Context(), input); err == nil { + t.Fatal("accepted invalid scoring response") + } else if tt.status == 500 { + var status api.StatusError + if !errors.As(err, &status) || status.StatusCode != 500 || !IsOutOfMemory(err) { + t.Fatalf("lost upstream status/OOM diagnostic: %v", err) + } + } + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + if _, err := runner.Score(ctx, input); err != nil { + t.Fatalf("runner unusable after failed scoring: %v", err) + } + }) + } +} + +func TestLlamaServerScoreCancellation(t *testing.T) { + for _, queued := range []bool{true, false} { + t.Run(fmt.Sprintf("queued=%v", queued), func(t *testing.T) { + started := make(chan struct{}) + var block atomic.Bool + block.Store(!queued) + runner := newLlamaScoreTestRunner(t, func(w http.ResponseWriter, r *http.Request) { + if block.CompareAndSwap(true, false) { + _, _ = io.Copy(io.Discard, r.Body) + close(started) + <-r.Context().Done() + return + } + fmt.Fprint(w, `{"tokens_predicted":1,"tokens_evaluated":3,"completion_probabilities":[{"id":65,"token":"A","prob":1,"top_probs":[{"id":65,"token":"A","prob":1}]}]}`) + }) + input := ScoreRequest{MaxTokens: 7, Rows: []ScoreRow{{Prompt: "abc", Candidates: []string{"A"}}}} + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + if queued { + if err := runner.sem.Acquire(t.Context(), 1); err != nil { + t.Fatal(err) + } + cancel() + } else { + go func() { + select { + case <-started: + cancel() + case <-ctx.Done(): + } + }() + } + _, err := runner.Score(ctx, input) + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation lost: %v", err) + } + if queued { + runner.sem.Release(1) + } + next, done := context.WithTimeout(t.Context(), time.Second) + defer done() + if _, err := runner.Score(next, input); err != nil { + t.Fatalf("runner unusable after cancellation: %v", err) + } + }) + } +} diff --git a/llm/score.go b/llm/score.go new file mode 100644 index 000000000..e7d10ee76 --- /dev/null +++ b/llm/score.go @@ -0,0 +1,27 @@ +package llm + +import "context" + +// Scorer is an optional runner capability for next-token candidate scoring. +// Prompts are already rendered; no chat template or generation options apply. +type Scorer interface { + Score(context.Context, ScoreRequest) (ScoreResponse, error) +} + +type ScoreRequest struct { + Rows []ScoreRow `json:"rows"` + MaxTokens int `json:"max_tokens"` +} + +type ScoreRow struct { + Prompt string `json:"prompt"` + Candidates []string `json:"candidates"` +} + +type ScoreResponse struct { + // Logits may be log probabilities: a shared offset within a row does not + // affect the softmax over its candidates. + Logits [][]float32 `json:"logits"` + InputTokens int `json:"input_tokens"` // Sum of complete prompt lengths, including shared tokens. + OutputTokens int `json:"output_tokens,omitempty"` // Tokens generated internally to obtain candidate scores. +} diff --git a/server/routes.go b/server/routes.go index 6ca00a2e0..d56c98c69 100644 --- a/server/routes.go +++ b/server/routes.go @@ -33,6 +33,7 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/auth" + "github.com/ollama/ollama/decision" "github.com/ollama/ollama/discover" "github.com/ollama/ollama/envconfig" "github.com/ollama/ollama/format" @@ -837,6 +838,93 @@ func (s *Server) GenerateHandler(c *gin.Context) { streamResponse(c, ch) } +// SystemOneHandler compiles typed questions, scores their allowed answers, and +// returns probabilities. Callers must select weights trained for the prompt format. +func (s *Server) SystemOneHandler(c *gin.Context) { + // TODO(parthsareen): Check token limits before copying state and schema into + // each question's prompt. This byte cap limits memory use until then. + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 64<<10) + var req decision.Request + if err := c.ShouldBindJSON(&req); err != nil { + var sizeErr *http.MaxBytesError + if errors.As(err, &sizeErr) { + c.JSON(http.StatusRequestEntityTooLarge, gin.H{"error": "request body must not exceed 64 KiB"}) + return + } + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + compiled, err := decision.Compile(req) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + ref, err := parseAndValidateModelRef(req.Model) + if err != nil { + writeModelRefParseError(c, err, http.StatusNotFound, fmt.Sprintf("model '%s' not found", req.Model)) + return + } + if ref.Source == modelSourceCloud { + c.JSON(http.StatusBadRequest, gin.H{"error": "System One requires a local Nimble or Tev model"}) + return + } + name, err := getExistingName(ref.Name) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Model)}) + return + } + m, err := GetModel(name.String()) + if err != nil { + handleScheduleError(c, req.Model, err) + return + } + if !m.isGGUF() || m.Config.Renderer != "qwen3.5" && !(m.Config.Renderer == "" && m.Config.ModelFamily == "qwen35") { + c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("model %q is not supported by System One; use a local Nimble or Tev GGUF model", req.Model)}) + return + } + r, _, _, err := s.scheduleRunner(c.Request.Context(), m, []model.Capability{model.CapabilityCompletion}, nil, req.KeepAlive, nil) + if err != nil { + handleScheduleError(c, req.Model, err) + return + } + scorer, ok := r.(llm.Scorer) + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("model %q does not support System One scoring; use a local Nimble or Tev model with a scoring-capable runner", req.Model)}) + return + } + if err := compiled.Render(func(messages []api.Message) (string, error) { + if m.System != "" { + messages = append([]api.Message{{Role: "system", Content: m.System}}, messages...) + } + think := &api.ThinkValue{Value: false} + if m.HasChatTemplate && chatModeForModel(m) == chatExecutionModeNative { + return r.ApplyChatTemplate(c.Request.Context(), llm.ChatRequest{Messages: messages, Think: think}) + } + return renderPrompt(m, messages, nil, think) + }); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + compiled.Request.MaxTokens = r.ContextLength() + result, err := scorer.Score(c.Request.Context(), compiled.Request) + if err != nil { + s.sched.expireRunnersForRuntimeOOM(m, err) + status := http.StatusInternalServerError + var statusErr api.StatusError + if errors.As(err, &statusErr) { + status = statusErr.StatusCode + } + c.JSON(status, gin.H{"error": err.Error()}) + return + } + response, err := compiled.Answer(req.Model, result) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, response) +} + func (s *Server) EmbedHandler(c *gin.Context) { checkpointStart := time.Now() var req api.EmbedRequest @@ -1967,6 +2055,7 @@ func (s *Server) GenerateRoutes() (http.Handler, error) { r.POST("/api/chat", s.withInferenceRequestLogging("/api/chat", s.ChatHandler)...) r.POST("/api/embed", s.EmbedHandler) r.POST("/api/embeddings", s.EmbeddingsHandler) + r.POST("/v1/systemone", s.SystemOneHandler) // Inference (OpenAI compatibility) // TODO(cloud-stage-a): apply Modelfile overlay deltas for local models with cloud diff --git a/server/systemone_test.go b/server/systemone_test.go new file mode 100644 index 000000000..0bede8e56 --- /dev/null +++ b/server/systemone_test.go @@ -0,0 +1,204 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/decision" + gguftest "github.com/ollama/ollama/internal/testutil/gguf" + "github.com/ollama/ollama/llm" + "github.com/ollama/ollama/manifest" + "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/types/model" +) + +type systemOneTestRunner struct { + mockRunner + request llm.ScoreRequest + err error + calls int +} + +func (r *systemOneTestRunner) Score(ctx context.Context, input llm.ScoreRequest) (llm.ScoreResponse, error) { + r.calls++ + r.request = input + return llm.ScoreResponse{Logits: [][]float32{{0, 2}}, InputTokens: 123, OutputTokens: 2}, r.err +} + +func TestSystemOneHandler(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("OLLAMA_MODELS", t.TempDir()) + config := model.ConfigV2{ModelFormat: "safetensors", Renderer: "qwen3.5", Capabilities: []string{"completion"}} + createSafetensorsTestModel(t, "safetensors-decision", config, nil) + for _, modelConfig := range []struct { + name, architecture, renderer, system, template string + contextLength int + }{ + {"decision-test", "qwen35", "qwen3.5", "Model-specific scoring instructions.", "", 1024}, + {"gguf-decision", "qwen35", "", "Native model scoring instructions.", "", 4096}, + {"go-template", "qwen35", "", "Model-specific scoring instructions.", "custom:{{ range .Messages }}{{ .Role }}:{{ .Content }}\n{{ end }}answer:", 1024}, + {"no-system", "qwen35", "qwen3.5", "", "", 1024}, + {"gguf-wrong-architecture", "llama", "", "", "", 1024}, + {"gguf-wrong-renderer", "qwen35", "gemma4", "", "", 1024}, + } { + kv := gguftest.KV{"general.architecture": modelConfig.architecture} + if modelConfig.template == "" { + kv["tokenizer.chat_template"] = "{{ messages }}" + } + _, digest := createBinFile(t, kv, nil) + configLayer, err := createConfigLayer(model.ConfigV2{ + ModelFormat: "gguf", ModelFamily: modelConfig.architecture, + Renderer: modelConfig.renderer, Capabilities: []string{"completion"}, + }) + if err != nil { + t.Fatal(err) + } + layers := []manifest.Layer{{MediaType: "application/vnd.ollama.image.model", Digest: digest}} + for _, layer := range []struct{ content, mediaType string }{ + {modelConfig.system, "application/vnd.ollama.image.system"}, + {fmt.Sprintf(`{"num_ctx":%d}`, modelConfig.contextLength), "application/vnd.ollama.image.params"}, + } { + l, err := manifest.NewLayer(strings.NewReader(layer.content), layer.mediaType) + if err != nil { + t.Fatal(err) + } + layers = append(layers, l) + } + if modelConfig.template != "" { + l, err := manifest.NewLayer(strings.NewReader(modelConfig.template), "application/vnd.ollama.image.template") + if err != nil { + t.Fatal(err) + } + layers = append(layers, l) + } + if err := manifest.WriteManifest(model.ParseName(modelConfig.name), *configLayer, layers); err != nil { + t.Fatal(err) + } + } + + const prefix = `{"model":"decision-test","state":"` + const suffix = `","questions":{"refund":{"type":"noul","instructions":"Refund requested?"}}}` + stateLimit := (64 << 10) - len(prefix) - len(suffix) + for _, tt := range []struct { + name string + body string + err error + status int + calls int + expire bool + }{ + {"success", `{"model":"decision-test","state":"refund please","questions":{"refund":{"type":"noul","instructions":"Refund requested?"}}}`, nil, 200, 1, false}, + {"GGUF success without renderer", `{"model":"gguf-decision","state":"refund please","questions":{"refund":{"type":"noul","instructions":"Refund requested?"}}}`, nil, 200, 1, false}, + {"Modelfile template", `{"model":"go-template","state":"refund please","questions":{"refund":{"type":"noul","instructions":"Refund requested?"}}}`, nil, 200, 1, false}, + {"no system prompt", `{"model":"no-system","state":"refund please","questions":{"refund":{"type":"noul","instructions":"Refund requested?"}}}`, nil, 200, 1, false}, + {"template failure", `{"model":"gguf-decision","state":"refund please","questions":{"refund":{"type":"noul","instructions":"Refund requested?"}}}`, errors.New("invalid model template"), 500, 0, false}, + {"GGUF incompatible architecture", `{"model":"gguf-wrong-architecture","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, nil, 400, 0, false}, + {"GGUF conflicting renderer", `{"model":"gguf-wrong-renderer","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, nil, 400, 0, false}, + {"runner validation", `{"model":"decision-test","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, api.StatusError{StatusCode: 400, ErrorMessage: "prompt too long"}, 400, 1, false}, + {"runner failure", `{"model":"decision-test","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, errors.New("runner failed"), 500, 1, false}, + {"runtime OOM", `{"model":"decision-test","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, errors.New("out of memory"), 500, 1, true}, + {"invalid schema", `{"model":"decision-test","state":"x","questions":{}}`, nil, 400, 0, false}, + {"missing model", `{"model":"missing","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, nil, 404, 0, false}, + {"cloud", `{"model":"decision:cloud","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, nil, 400, 0, false}, + {"safetensors unsupported", `{"model":"safetensors-decision","state":"x","questions":{"x":{"type":"noul","instructions":"q"}}}`, nil, 400, 0, false}, + {"bad JSON", `{`, nil, 400, 0, false}, + {"at body limit", prefix + strings.Repeat("x", stateLimit) + suffix, nil, 200, 1, false}, + {"over body limit", prefix + strings.Repeat("x", stateLimit+1) + suffix, nil, 413, 0, false}, + } { + t.Run(tt.name, func(t *testing.T) { + runner := &systemOneTestRunner{err: tt.err} + runner.TemplateFn = func(_ context.Context, req llm.ChatRequest) (string, error) { + if len(req.Messages) != 2 || req.Messages[0].Role != "system" || req.Messages[0].Content != "Native model scoring instructions." || req.Messages[1].Role != "user" { + t.Fatalf("model system prompt was not passed to the native template: %+v", req.Messages) + } + if req.Think == nil || req.Think.Bool() { + t.Fatal("scoring must disable thinking") + } + return "native:" + req.Messages[1].Content, tt.err + } + ref := &runnerRef{llama: runner, refCount: 1, sessionDuration: time.Hour} + s := newServerWithMockRunner(t, &runner.mockRunner) + s.sched.loadFn = func(req *LlmRequest, _ ml.SystemInfo, _ []ml.DeviceInfo, _ bool) bool { + runner.contextLength = req.opts.NumCtx + ref.model = req.model + ref.modelKey = schedulerModelKey(req.model) + s.sched.loadedMu.Lock() + s.sched.loaded[ref.modelKey] = ref + s.sched.loadedMu.Unlock() + req.successCh <- ref + return false + } + router := gin.New() + router.POST("/v1/systemone", s.SystemOneHandler) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/systemone", strings.NewReader(tt.body)) + req.ContentLength = -1 + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + if w.Code != tt.status || runner.calls != tt.calls { + t.Fatalf("status=%d calls=%d body=%s", w.Code, runner.calls, w.Body) + } + if tt.calls == 0 && tt.err == nil && ref.model != nil { + t.Fatal("loaded a runner for a rejected request") + } + if tt.calls > 0 { + want := time.Hour + if tt.expire { + want = 0 + } + ref.refMu.Lock() + duration := ref.sessionDuration + ref.refMu.Unlock() + if duration != want { + t.Fatalf("runner keep-alive=%v, want %v", duration, want) + } + } + if tt.status == 200 { + var response struct { + Answers map[string]struct { + Noul float64 `json:"noul"` + } `json:"answers"` + Usage decision.Usage `json:"usage"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if response.Answers["refund"].Noul < 0.88 || response.Usage.InputTokens != 123 || response.Usage.OutputTokens != 2 { + t.Fatalf("incorrect scoring response: %s", w.Body) + } + prompt := runner.request.Rows[0].Prompt + wantContext := 1024 + if ref.model.HasGoTemplate { + if !strings.HasPrefix(prompt, "custom:system:Model-specific scoring instructions.\nuser:") || !strings.HasSuffix(prompt, "answer:") { + t.Fatalf("scorer did not receive the Modelfile template output: %q", prompt) + } + } else if ref.model.Config.Renderer == "" { + wantContext = 4096 + if !strings.HasPrefix(prompt, `native:{"context":`) { + t.Fatalf("scorer did not receive the native template output: %q", prompt) + } + } else { + wantPrefix := "<|im_start|>user\n" + if ref.model.System != "" { + wantPrefix = "<|im_start|>system\nModel-specific scoring instructions.<|im_end|>\n<|im_start|>user\n" + } + if !strings.HasPrefix(prompt, wantPrefix) || !strings.HasSuffix(prompt, "\n\n\n\n") { + t.Fatalf("scorer did not receive the model's rendered prompt: %q", prompt) + } + } + if runner.request.MaxTokens != wantContext { + t.Fatalf("scoring context = %d, want model num_ctx %d", runner.request.MaxTokens, wantContext) + } + } + }) + } +}