mirror of
https://github.com/ollama/ollama.git
synced 2025-05-11 10:26:53 +02:00
sample: temporarily use grammars for constrained generation in new engine (#9586)
This commit is contained in:
parent
a1cda80bcb
commit
e093db92c4
10 changed files with 301 additions and 213 deletions
|
@ -7,10 +7,10 @@ import (
|
|||
)
|
||||
|
||||
// Helper to convert float64 slice to logit slice
|
||||
func toLogits(values []float64) []logit {
|
||||
tokens := make([]logit, len(values))
|
||||
func toTokens(values []float64) []token {
|
||||
tokens := make([]token, len(values))
|
||||
for i, v := range values {
|
||||
tokens[i] = logit{
|
||||
tokens[i] = token{
|
||||
id: int32(i),
|
||||
value: float32(v),
|
||||
}
|
||||
|
@ -19,7 +19,7 @@ func toLogits(values []float64) []logit {
|
|||
}
|
||||
|
||||
// Helper to compare logit slices
|
||||
func compareLogits(t *testing.T, name string, want []float64, got []logit) {
|
||||
func compareLogits(t *testing.T, name string, want []float64, got []token) {
|
||||
t.Helper()
|
||||
if len(want) != len(got) {
|
||||
t.Errorf("%s: length mismatch: want %d, got %d", name, len(want), len(got))
|
||||
|
@ -36,13 +36,13 @@ func TestTemperature(t *testing.T) {
|
|||
input := []float64{2, -1, 4, -3, 1, -2, 0}
|
||||
want := []float64{-4, -10, 0, -14, -6, -12, -8} // (logit - max logit) / temp
|
||||
|
||||
got := temperature(toLogits(input), 0.5)
|
||||
got := temperature(toTokens(input), 0.5)
|
||||
compareLogits(t, "Temperature", want, got)
|
||||
}
|
||||
|
||||
func TestSoftmax(t *testing.T) {
|
||||
input := []float64{-3, -2, -1, 0, 1, 2, 4}
|
||||
got := softmax(toLogits(input))
|
||||
got := softmax(toTokens(input))
|
||||
|
||||
// Check probabilities sum to 1
|
||||
var sum float32
|
||||
|
@ -65,7 +65,7 @@ func TestTopK(t *testing.T) {
|
|||
input := []float64{-3, -2, -1, 0, 1, 2, 4}
|
||||
|
||||
// Test k=3
|
||||
got := topK(toLogits(input), 3)
|
||||
got := topK(toTokens(input), 3)
|
||||
if len(got) != 3 {
|
||||
t.Errorf("topK(3): wrong length: want 3, got %d", len(got))
|
||||
}
|
||||
|
@ -74,13 +74,13 @@ func TestTopK(t *testing.T) {
|
|||
compareLogits(t, "topK(3)", want, got)
|
||||
|
||||
// Test k > len
|
||||
got = topK(toLogits(input), 10)
|
||||
got = topK(toTokens(input), 10)
|
||||
compareLogits(t, "topK(10)", input, got)
|
||||
}
|
||||
|
||||
func TestTopP(t *testing.T) {
|
||||
input := []float64{-3, -2, -1, 0, 1, 2, 4}
|
||||
tokens := toLogits(input)
|
||||
tokens := toTokens(input)
|
||||
|
||||
// First apply temperature and softmax to get probabilities
|
||||
tokens = temperature(tokens, 1)
|
||||
|
@ -99,7 +99,7 @@ func TestTopP(t *testing.T) {
|
|||
|
||||
func TestMinP(t *testing.T) {
|
||||
input := []float64{-3, -2, -1, 0, 1, 2, 4, 3}
|
||||
tokens := toLogits(input)
|
||||
tokens := toTokens(input)
|
||||
|
||||
// First apply temperature and softmax
|
||||
tokens = temperature(tokens, 1)
|
||||
|
@ -116,7 +116,7 @@ func TestMinP(t *testing.T) {
|
|||
|
||||
func TestSortLogits(t *testing.T) {
|
||||
input := []float64{3, 1, 4, 2, -1, 0, -2}
|
||||
tokens := toLogits(input)
|
||||
tokens := toTokens(input)
|
||||
|
||||
sortLogits(tokens)
|
||||
|
||||
|
@ -133,15 +133,15 @@ func TestSortLogits(t *testing.T) {
|
|||
|
||||
func BenchmarkTransforms(b *testing.B) {
|
||||
// Generate random logits
|
||||
tokens := make([]logit, 1<<16)
|
||||
tokens := make([]token, 1<<16)
|
||||
for i := range tokens {
|
||||
tokens[i] = logit{
|
||||
tokens[i] = token{
|
||||
id: int32(i),
|
||||
value: rand.Float32(),
|
||||
}
|
||||
}
|
||||
|
||||
tokensCopy := make([]logit, len(tokens))
|
||||
tokensCopy := make([]token, len(tokens))
|
||||
|
||||
b.Run("Temperature", func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue