mirror of
https://github.com/ollama/ollama.git
synced 2026-10-21 11:01:33 +00:00
feat: add System One scoring API (#18606)
This commit is contained in:
Parth Sareen
GitHub
parent
16b4376aea
commit
ffb6f220d0
8 changed files with 1557 additions and 0 deletions
228
decision/systemone.go
Normal file
228
decision/systemone.go
Normal file
|
|
@ -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
|
||||
}
|
||||
231
decision/systemone_test.go
Normal file
231
decision/systemone_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
76
decision/types.go
Normal file
76
decision/types.go
Normal file
|
|
@ -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"`
|
||||
}
|
||||
281
llm/llama_server_score.go
Normal file
281
llm/llama_server_score.go
Normal file
|
|
@ -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
|
||||
}
|
||||
421
llm/llama_server_score_test.go
Normal file
421
llm/llama_server_score_test.go
Normal file
|
|
@ -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, "<special>") {
|
||||
tokens = append(tokens[:len(tokens)-len("<special>")], 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{"<special>"}}}}},
|
||||
{"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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
27
llm/score.go
Normal file
27
llm/score.go
Normal file
|
|
@ -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.
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
204
server/systemone_test.go
Normal file
204
server/systemone_test.go
Normal file
|
|
@ -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, "<think>\n\n</think>\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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue