extras

sample/samplers.go

@@ -1,11 +1,10 @@
 package sample
 
 import (
-	"errors"
 	"math"
-	"math/rand/v2"
-	"slices"
+	"math/rand"
 	"sync"
+	"time"
 
 	"github.com/ollama/ollama/llama"
 )
@@ -87,53 +86,53 @@ func (s *Sampler) sample(tokens []token) (token, error) {
 	// topK also sorts the tokens in descending order of logits
 	tokens = topK(tokens, s.topK)
 
-	// token logit values are updated to probabilities
-	tokens = temperature(tokens, s.temperature)
-
 	tokens = topP(tokens, s.topP)
 	tokens = minP(tokens, s.minP)
 
-	// TODO: this should fall back to greedy sampling
-	// or topP, topK values etc should be such that
-	// there are always tokens to sample from
-	if len(tokens) == 0 {
-		return token{}, errors.New("no tokens to sample from")
-	}
+	// token logit values are updated to probabilities
+	temperature(tokens, s.temperature)
+	softmax(tokens)
+	return tokens[dist(tokens, s.rng.Int63())], nil
 
-	var r float32
-	if s.rng != nil {
-		r = s.rng.Float32()
-	} else {
-		r = rand.Float32()
-	}
+	// // TODO: this should fall back to greedy sampling
+	// // or topP, topK values etc should be such that
+	// // there are always tokens to sample from
+	// if len(tokens) == 0 {
+	// 	return token{}, errors.New("no tokens to sample from")
+	// }
 
-	// Calculate cumulative sum of probabilities
-	var sum float32
-	for i := range tokens {
-		sum += tokens[i].value
-		tokens[i].value = sum
-	}
-	r *= tokens[len(tokens)-1].value
+	// var r float32
+	// if s.rng != nil {
+	// 	r = s.rng.Float32()
+	// } else {
+	// 	r = rand.Float32()
+	// }
 
-	idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
-		if token.value < target {
-			return -1
-		}
-		return 1
-	})
+	// // Calculate cumulative sum of probabilities
+	// var sum float32
+	// for i := range tokens {
+	// 	sum += tokens[i].value
+	// 	tokens[i].value = sum
+	// }
+	// r *= tokens[len(tokens)-1].value
 
-	return tokens[idx], nil
+	// idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
+	// 	if token.value < target {
+	// 		return -1
+	// 	}
+	// 	return 1
+	// })
+
+	// return tokens[idx], nil
 }
 
 // TODO(parthsareen): update sampler interface to use json unmarshal https://github.com/ollama/ollama/issues/9278
 func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int, grammar *Grammar) Sampler {
 	var rng *rand.Rand
 	if seed != -1 {
-		// PCG requires two parameters: sequence and stream
-		// Use original seed for sequence
-		sequence := uint64(seed)
-		// Use golden ratio hash to generate statistically independent seeds
-		rng = rand.New(rand.NewPCG(sequence, sequence^0x9E3779B9))
+		rng = rand.New(rand.NewSource(int64(seed)))
+	} else {
+		rng = rand.New(rand.NewSource(time.Now().UnixNano()))
 	}
 	if temperature < 0.0 {
 		temperature = 0.0

sample/testdata/logits.bin

@@ -0,0 +1 @@
+[-0.14049109816551208, 0.06803577393293381, 0.08947011083364487, -0.15352056920528412, 0.16544274985790253, -0.22324655950069427, -0.07907228916883469, -0.11153314262628555, 0.16745619475841522, -0.08845115453004837, -0.05397915467619896, 0.09702621400356293, 0.1040145754814148, 0.034065235406160355, -0.023494448512792587, 0.023317236453294754, 0.06201161816716194, -0.14594611525535583, 0.05035833641886711, -0.06913720071315765, -0.0850675031542778, 0.0007670316845178604, -0.08676620572805405, 0.020724542438983917, 0.18044163286685944, 0.04924239218235016, -0.15723371505737305, 0.03611345216631889, -0.17997989058494568, -0.068993479013443, 0.05774486064910889, 0.08498939871788025, 0.12943890690803528, -0.006152433808892965, 0.14178094267845154, 0.1473844349384308, -0.012766982428729534, 0.045800212770700455, 0.19246874749660492, -0.09747152775526047, 0.006072793155908585, 0.017143171280622482, -0.02200348488986492, -0.027592593804001808, -0.04569224640727043, 0.026975959539413452, 0.0077785165049135685, -0.17419104278087616, 0.2779255509376526, -0.030980607494711876, 0.11484438180923462, -0.20746532082557678, -0.06505528837442398, -0.019896341487765312, 0.02105427347123623, 0.20272167026996613, 0.04019889235496521, -0.017126401886343956, -0.0002831551246345043, -0.07687067240476608, 0.012454139068722725, -0.15010742843151093, -0.07807016372680664, 0.01188774686306715, 0.13062332570552826, -0.08224914968013763, -0.08921745419502258, 0.013551718555390835, -0.02187133952975273, 0.005148813594132662, 0.008781243115663528, 0.10651258379220963, 0.015292760916054249, 0.03104853443801403, 0.017170460894703865, 0.09132644534111023, 0.06599127501249313, 0.11407657712697983, 0.2534945011138916, -0.1733740270137787, 0.0020254752598702908, 0.05714576318860054, -0.19750545918941498, -0.10138755291700363, -0.15401461720466614, -0.011985579505562782, 0.04546714201569557, -0.0640132948756218, -0.030880948528647423, 0.11062195152044296, 0.14438383281230927, -0.07269343733787537, 0.08082295954227448, 0.028231162577867508, -0.0299157053232193, 0.015319091267883778, -0.07127209007740021, -0.004352295771241188, -0.08830961585044861, -0.15393312275409698, -0.05910560116171837, -0.044640619307756424, -0.04869438707828522, -0.1213475689291954, -0.0413384735584259, 0.11684834957122803, 0.02147417515516281, -0.19815197587013245, -0.027197321876883507, 0.09237757325172424, 0.016114329919219017, -0.07321144640445709, 0.013616911135613918, 0.1633937656879425, 0.2265119105577469, 0.04887519031763077, -0.09309308975934982, -0.10166775435209274, 0.034362174570560455, -0.016704145818948746, 0.26136764883995056, 0.024831606075167656, 0.019148696213960648, -0.10227427631616592, 0.05994879826903343, -0.028703754767775536, -0.08792375773191452, -0.09701353311538696, 0.006381520070135593, 0.09135088324546814, -0.03421201556921005, 0.027725202962756157, 0.10950426012277603, 0.08471143990755081, -0.06850480288267136, -0.2098373919725418, 0.024568362161517143, 0.016832124441862106, -0.09213020652532578, 0.07319895178079605, -0.2254962921142578, 0.09601815789937973, -0.14963679015636444, 0.13309000432491302, -0.06646179407835007, 0.007633576635271311, -0.006754335947334766, -0.12666666507720947, -0.06185764819383621, -0.03350026533007622, -0.07640213519334793, 0.11266022175550461, 0.05684903264045715, -0.021437257528305054, -0.057359494268894196, 0.022800210863351822, 0.042317431420087814, 0.010130883194506168, 0.12166910618543625, -0.0677243322134018, -0.05597330257296562, 0.10634074360132217, -0.14240652322769165, -0.07514529675245285, 0.11910972744226456, -0.0635921061038971, -0.12317600101232529, 0.07619466632604599, -0.05430442839860916, 0.01882130652666092, -0.013355006463825703, 0.13947084546089172, 0.13103191554546356, 0.21173261106014252, 0.012381404638290405, 0.10963138192892075, 0.20157527923583984, 0.04969905689358711, -0.07988657802343369, -0.12305821478366852, 0.030198488384485245, -0.12815065681934357, -0.1493828296661377, 0.16609710454940796, 0.03147311136126518, -0.05435416102409363, -0.06691042333841324, 0.08390865474939346, 0.1344001740217209, -0.10086281597614288, -0.07271518558263779, -0.05838662385940552, -0.10467743873596191, 0.047540608793497086, -0.06734462082386017, -0.13437806069850922, -0.05936319753527641, 0.10448285192251205, -0.11318164318799973, 0.018219152465462685, -0.1846952736377716, -0.07407224178314209, 0.07009526342153549, 0.11060779541730881, 0.0008995900861918926, 0.031116077676415443, 0.005915861111134291, -0.022779200226068497, -0.059181224554777145, 0.07091255486011505, 0.07433564215898514, 0.056968193501234055, 0.11839757114648819, 0.15867207944393158, -0.15376228094100952, -0.21046996116638184, -0.0966387465596199, -0.0876750648021698, -0.21115104854106903, -0.07812130451202393, 0.013138437643647194, 0.037374746054410934, -0.001360200229100883, -0.01634298823773861, -0.04794382303953171, -0.22696244716644287, -0.09551295638084412, 0.09177322685718536, -0.07562369853258133, 0.018677910789847374, 0.098722Line truncated

sample/transforms.go

@@ -3,6 +3,7 @@ package sample
 import (
 	"container/heap"
 	"math"
+	"math/rand"
 	"slices"
 )
 
@@ -25,32 +26,6 @@ func (h *tokenHeap) Pop() any {
 	return x
 }
 
-// temperature applies scaling and softmax to the logits
-func temperature(ts []token, temp float32) []token {
-	// Find max logit for numerical stability
-	maxLogit := float32(math.Inf(-1))
-	for _, t := range ts {
-		if t.value > maxLogit {
-			maxLogit = t.value
-		}
-	}
-
-	// Apply temperature and compute exp(x - max)
-	temp = max(temp, 1e-7)
-	var sum float32
-	for i, v := range ts {
-		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
-		sum += ts[i].value
-	}
-
-	// Normalize
-	for i := range ts {
-		ts[i].value /= sum
-	}
-
-	return ts
-}
-
 // topK limits the number of tokens considered to the k highest logits
 func topK(ts []token, k int) []token {
 	if k >= len(ts) || k <= 0 {
@@ -134,3 +109,59 @@ func minP(ts []token, p float32) []token {
 	ts = validTokens
 	return ts
 }
+
+func temperature(ts []token, temp float32) {
+	for i := range ts {
+		ts[i].value /= temp
+	}
+}
+
+func softmax(ts []token) {
+	if len(ts) == 0 {
+		return
+	}
+
+	// Find max logit for numerical stability
+	maxLogit := ts[0].value
+	for _, t := range ts {
+		if t.value > maxLogit {
+			maxLogit = t.value
+		}
+	}
+
+	// Compute exp(logit - maxLogit) and sum them
+	var sumExp float32
+	for i, t := range ts {
+		expVal := float32(math.Exp(float64(t.value - maxLogit)))
+		ts[i].value = expVal
+		sumExp += expVal
+	}
+
+	// Normalize probabilities
+	for i := range ts {
+		ts[i].value /= sumExp
+	}
+}
+
+// applyDist selects a token based on probabilities and seed
+func dist(ts []token, seed int64) int {
+	rng := rand.New(rand.NewSource(seed))
+
+	cdf := make([]float32, len(ts))
+	var cumSum float32
+	for i, t := range ts {
+		cumSum += t.value
+		cdf[i] = cumSum
+	}
+
+	r := rng.Float32() * cumSum
+
+	// Select token based on CDF
+	for i, probSum := range cdf {
+		if r < probSum {
+			return i
+		}
+	}
+
+	return len(ts) - 1
+}

sample/transforms_test.go

@@ -1,8 +1,13 @@
 package sample
 
 import (
+	"encoding/binary"
+	"errors"
 	"math"
 	"math/rand/v2"
+	"os"
+	"path/filepath"
+	"runtime"
 	"testing"
 )
 
@@ -143,6 +148,98 @@ func TestSortLogits(t *testing.T) {
 	compareLogits(t, "sortLogits", want, tokens)
 }
 
+// TestSortLogitsWithRealData tests sorting behavior using real model logit distributions
+func TestSortLogitsWithRealData(t *testing.T) {
+	// This will be populated from testdata/logits.bin
+	// Format: 32-bit float array in binary format
+	logits, err := loadTestLogits(t)
+	if err != nil {
+		t.Skipf("Skipping real logit test: %v", err)
+		return
+	}
+
+	tokens := toTokens(logits)
+	sortLogits(tokens)
+
+	// Calculate n for verification
+	n := int(math.Sqrt(float64(len(tokens)))) + 1
+	if n > 1000 {
+		n = 1000
+	} else if n < 100 {
+		n = 100
+	}
+
+	t.Logf("Testing with %d tokens, partial sorting top %d", len(tokens), n)
+
+	// Only verify the top n elements are sorted (which is what we guarantee)
+	// This is much faster than checking the entire array
+	topN := tokens[:n]
+	for i := 1; i < len(topN); i++ {
+		if topN[i].value > topN[i-1].value {
+			t.Fatalf("top %d tokens not properly sorted at index %d: %.15f > %.15f",
+				n, i, topN[i].value, topN[i-1].value)
+		}
+	}
+
+	// Verify we didn't lose any high value tokens by checking that
+	// all tokens after position n are <= the nth token
+	// Do this in chunks to avoid timeouts on large arrays
+	nthValue := tokens[n-1].value
+	const chunkSize = 1000
+
+	for start := n; start < len(tokens); start += chunkSize {
+		end := min(start+chunkSize, len(tokens))
+		for i := start; i < end; i++ {
+			if tokens[i].value > nthValue {
+				t.Fatalf("found higher value token after position %d: tokens[%d].value = %.15f > %.15f",
+					n, i, tokens[i].value, nthValue)
+			}
+		}
+	}
+}
+
+// loadTestLogits loads logit test data from testdata/logits.bin
+func loadTestLogits(t *testing.T) ([]float32, error) {
+	t.Helper()
+
+	_, currFile, _, ok := runtime.Caller(0)
+	if !ok {
+		return nil, errors.New("could not determine test file path")
+	}
+	testDataPath := filepath.Join(filepath.Dir(currFile), "testdata", "logits.bin")
+
+	file, err := os.Open(testDataPath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	stat, err := file.Stat()
+	if err != nil {
+		return nil, err
+	}
+
+	numFloats := stat.Size() / 4 // each float32 is 4 bytes
+	if numFloats*4 != stat.Size() {
+		return nil, errors.New("logits.bin has invalid size: not a multiple of 4 bytes")
+	}
+
+	logits := make([]float32, numFloats)
+	for i := range logits {
+		var val uint32
+		if err := binary.Read(file, binary.LittleEndian, &val); err != nil {
+			return nil, err
+		}
+		logits[i] = math.Float32frombits(val)
+	}
+
+	if len(logits) == 0 {
+		return nil, errors.New("logits.bin is empty")
+	}
+
+	return logits, nil
+}
+
 func BenchmarkTransforms(b *testing.B) {
 	// Generate random logits
 	tokens := make([]token, 1<<16)