sample: temporarily use grammars for constrained generation in new engine (#9586)

This commit is contained in:
Jeffrey Morgan 2025-03-10 16:17:39 +01:00 committed by GitHub
parent a1cda80bcb
commit e093db92c4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 301 additions and 213 deletions

View file

@ -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()