f4432e1dba
The `stop` option to the generate API is a list of sequences that should cause generation to stop. Although these are commonly called "stop tokens", they do not necessarily correspond to LLM tokens (per the LLM's tokenizer). For example, if the caller sends a generate request with `"stop":["\n"]`, then generation should stop on any token containing `\n` (and trim `\n` from the output), not just if the token exactly matches `\n`. If `stop` were interpreted strictly as LLM tokens, then it would require callers of the generate API to know the LLM's tokenizer and enumerate many tokens in the `stop` list. Fixes https://github.com/jmorganca/ollama/issues/295.
79 lines
2 KiB
Go
79 lines
2 KiB
Go
package llm
|
|
|
|
import (
|
|
"bytes"
|
|
"testing"
|
|
)
|
|
|
|
func TestCheckStopConditions(t *testing.T) {
|
|
tests := map[string]struct {
|
|
b string
|
|
stop []string
|
|
wantB string
|
|
wantStop bool
|
|
wantEndsWithStopPrefix bool
|
|
}{
|
|
"not present": {
|
|
b: "abc",
|
|
stop: []string{"x"},
|
|
wantStop: false,
|
|
wantEndsWithStopPrefix: false,
|
|
},
|
|
"exact": {
|
|
b: "abc",
|
|
stop: []string{"abc"},
|
|
wantStop: true,
|
|
wantEndsWithStopPrefix: false,
|
|
},
|
|
"substring": {
|
|
b: "abc",
|
|
stop: []string{"b"},
|
|
wantB: "a",
|
|
wantStop: true,
|
|
wantEndsWithStopPrefix: false,
|
|
},
|
|
"prefix 1": {
|
|
b: "abc",
|
|
stop: []string{"abcd"},
|
|
wantStop: false,
|
|
wantEndsWithStopPrefix: true,
|
|
},
|
|
"prefix 2": {
|
|
b: "abc",
|
|
stop: []string{"bcd"},
|
|
wantStop: false,
|
|
wantEndsWithStopPrefix: true,
|
|
},
|
|
"prefix 3": {
|
|
b: "abc",
|
|
stop: []string{"cd"},
|
|
wantStop: false,
|
|
wantEndsWithStopPrefix: true,
|
|
},
|
|
"no prefix": {
|
|
b: "abc",
|
|
stop: []string{"bx"},
|
|
wantStop: false,
|
|
wantEndsWithStopPrefix: false,
|
|
},
|
|
}
|
|
for name, test := range tests {
|
|
t.Run(name, func(t *testing.T) {
|
|
var b bytes.Buffer
|
|
b.WriteString(test.b)
|
|
stop, endsWithStopPrefix := handleStopSequences(&b, test.stop)
|
|
if test.wantB != "" {
|
|
gotB := b.String()
|
|
if gotB != test.wantB {
|
|
t.Errorf("got b %q, want %q", gotB, test.wantB)
|
|
}
|
|
}
|
|
if stop != test.wantStop {
|
|
t.Errorf("got stop %v, want %v", stop, test.wantStop)
|
|
}
|
|
if endsWithStopPrefix != test.wantEndsWithStopPrefix {
|
|
t.Errorf("got endsWithStopPrefix %v, want %v", endsWithStopPrefix, test.wantEndsWithStopPrefix)
|
|
}
|
|
})
|
|
}
|
|
}
|