ronantakizawa/codereview-bench huggingface dataset:
Repository: ollama/ollama
File: sample/transforms_test.go (Go)

Code before:
			t.Errorf("%s: index %d: want %f, got %f", name, i, want[i], got[i].value)
		}
	}
}

func TestTemperature(t *testing.T) {
	input := []float32{1.0, 4.0, -2.0, 0.0}
	got := temperature(toTokens(input), 0.5)
	want := []float32{2.0, 8.0, -4.0, 0.0}
	compareLogits(t, "temperature(0.5)", want, got)

	got = temperature(toTokens(input), 1.0)
	want = []float32{1.0, 4.0, -2.0, 0.0}
	compareLogits(t, "temperature(1)", want, got)

	got = temperature(toTokens(input), 0.0)
	want = []float32{1e7, 4e7, -2e7, 0.0}
	compareLogits(t, "temperature(0)", want, got)
}

func TestSoftmax(t *testing.T) {
	tests := []struct {
		name  string
		input []float32
	}{
		{
			name:  "normal distribution",
			input: []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367},
		},
		{
			name:  "single value",
			input: []float32{1.0},
		},
		{
			name:  "identical values",
			input: []float32{0.9, 0.9, 0.9},
		},
		{
			name:  "large values",
			input: []float32{1000.0, 2000.0, 3000.0},
		},
		{
			name:  "small values",
			input: []float32{1e-6, 2e-6, 3e-6},
		},
		{
			name:  "negative values",
			input: []float32{-1.0, -2.0, -3.0},
		},
		{
			name:  "mixed values",

Review comment (suggestion):
Let's add a test with manually computed values

```
    name: "simple softmax",
    input: []float32{1, 2, 3},
    expected: []float32{0.09003057, 0.24472848, 0.66524094},
```

Code after:
			t.Errorf("%s: index %d: want %f, got %f", name, i, want[i], got[i].value)
		}
	}
}

func TestTemperature(t *testing.T) {
	input := []float32{1.0, 4.0, -2.0, 0.0}
	got := temperature(toTokens(input), 0.5)
	want := []float32{2.0, 8.0, -4.0, 0.0}
	compareLogits(t, "temperature(0.5)", want, got)

	got = temperature(toTokens(input), 1.0)
	want = []float32{1.0, 4.0, -2.0, 0.0}
	compareLogits(t, "temperature(1)", want, got)

	got = temperature(toTokens(input), 0.0)
	want = []float32{1e7, 4e7, -2e7, 0.0}
	compareLogits(t, "temperature(0)", want, got)
}

func TestSoftmax(t *testing.T) {
	tests := []struct {
		name     string
		input    []float32
		expected []float32
	}{
		{
			name:     "correctness softmax",
			input:    []float32{1, -2, 3, 0},
			expected: []float32{0.113550, 0.005653, 0.839024, 0.041773},
		},
		{
			name:  "normal distribution",
			input: []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367},
		},
		{
			name:  "single value",
			input: []float32{1.0},
		},
		{
			name:  "identical values",
			input: []float32{0.9, 0.9, 0.9},
		},
		{
			name:  "large values",
			input: []float32{1000.0, 2000.0, 3000.0},
		},
		{
			name:  "small values",
			input: []float32{1e-6, 2e-6, 3e-6},
		},
Repository: ollama/ollama
File: sample/samplers.go (Go)

Code before:
// sample returns the highest probability token from the tokens
// given sampler parameters. It also has side effects of modifying the tokens
func (s *Sampler) sample(tokens []token) (token, error) {
	if s.temperature == 0 {
		return greedy(tokens), nil
	}

	// topK also sorts the tokens in descending order of logits
	tokens = topK(tokens, s.topK)

	tokens = temperature(tokens, s.temperature)
	tokens = softmax(tokens)

	tokens = topP(tokens, s.topP)
	tokens = minP(tokens, s.minP)

	// fallback to greedy sampling if no tokens are left
	if len(tokens) == 0 {
		slog.Warn("sample: no tokens left after applying transforms, falling back to greedy sampling")
		return greedy(tokens), nil
	}

	var r float32
	if s.rng != nil {
		r = s.rng.Float32()
	} else {
		r = rand.Float32()
	}

	// 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

	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 {

Review comment (suggestion):
This will be problematic because `tokens` has been mutated! and is in fact probably empty. I suggest we fix this after figuring out the mutable/immutable story

Code after:

	return max
}

// sample returns the highest probability token from the tokens
// given sampler parameters. It also has side effects of modifying the tokens
func (s *Sampler) sample(tokens []token) (token, error) {
	if s.temperature == 0 {
		return greedy(tokens), nil
	}

	// topK also sorts the tokens in descending order of logits
	tokens = topK(tokens, s.topK)

	// scale and normalize the tokens in place
	temperature(tokens, s.temperature)
	softmax(tokens)

	tokens = topP(tokens, s.topP)
	tokens = minP(tokens, s.minP)

	var r float32
	if s.rng != nil {
		r = s.rng.Float32()
	} else {
		r = rand.Float32()
	}

	// 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

	idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
		if token.value < target {
			return -1
		}
		return 1
	})

	if math.IsNaN(float64(sum)) {
		return token{}, errors.New("sample: logits sum to NaN, check model output")
	}
	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 {
Repository: ollama/ollama
File: sample/samplers.go (Go)

Code before:
	// 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")
	}

	var r float32
	if s.rng != nil {
		r = s.rng.Float32()
	} else {
		r = rand.Float32()
	}

	// 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

	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))
	}
	if temperature < 0.0 {
		temperature = 0.0
	}

	if topP < 0.0 {
		topP = 0.0
	}
	if topP >= 1.0 {
		topP = 1.0
	}

Review comment (suggestion):
If all logits are NaNs then the sum will be NaN - we should error in this case. Added a test case as well

Code after:
	tokens = topP(tokens, s.topP)
	tokens = minP(tokens, s.minP)

	var r float32
	if s.rng != nil {
		r = s.rng.Float32()
	} else {
		r = rand.Float32()
	}

	// 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

	idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
		if token.value < target {
			return -1
		}
		return 1
	})

	if math.IsNaN(float64(sum)) {
		return token{}, errors.New("sample: logits sum to NaN, check model output")
	}
	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))
	}
	if temperature < 0.0 {
		temperature = 0.0
	}

	if topP < 0.0 {
		topP = 0.0
	}
	if topP >= 1.0 {
		topP = 1.0
	}
Repository: ollama/ollama
File: runner/ollamarunner/runner.go (Go)

Code before:
	port := fs.Int("port", 8080, "Port to expose the server on")
	threads := fs.Int("threads", runtime.NumCPU(), "Number of threads to use during generation")
	verbose := fs.Bool("verbose", false, "verbose output (default: disabled)")
	_ = fs.Bool("no-mmap", false, "do not memory-map model (slower load but may reduce pageouts if not using mlock)")
	_ = fs.Bool("mlock", false, "force system to keep model in RAM rather than swapping or compressing")
	tensorSplit := fs.String("tensor-split", "", "fraction of the model to offload to each GPU, comma-separated list of proportions")
	multiUserCache := fs.Bool("multiuser-cache", false, "optimize input cache algorithm for multiple users")

	var lpaths multiLPath
	fs.Var(&lpaths, "lora", "Path to lora layer file (can be specified multiple times)")

	fs.Usage = func() {
		fmt.Fprintf(fs.Output(), "Runner usage\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return err
	}
	level := slog.LevelInfo
	if *verbose {
		level = slog.LevelDebug
	}
	handler := slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level:     level,
		AddSource: true,
		ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
			if attr.Key == slog.SourceKey {
				source := attr.Value.Any().(*slog.Source)
				source.File = filepath.Base(source.File)
			}
			return attr
		},
	})
	slog.SetDefault(slog.New(handler))
	slog.Info("starting ollama engine")

	server := &Server{
		batchSize: *batchSize,
		status:    ServerStatusLoadingModel,
	}

	// TODO(jessegross): Parameters that need to be implemented:
	//	no-mmap
	//	mlock

	var tensorSplitFloats []float32
	if *tensorSplit != "" {
		splits := strings.Split(*tensorSplit, ",")
		tensorSplitFloats = make([]float32, len(splits))
		for i, s := range splits {
			f, _ := strconv.ParseFloat(s, 32)

Review comment (suggestion):
```suggestion
		http.Error(w, "this model does not support embeddings", http.StatusNotImplemented)
```

Code after:
		TensorSplit:    tensorSplitFloats,
		FlashAttention: *flashAttention,
	}

	server.ready.Add(1)
	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)

	server.cond = sync.NewCond(&server.mu)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go server.run(ctx)

	addr := "127.0.0.1:" + strconv.Itoa(*port)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		fmt.Println("Listen error:", err)
		return err
	}
	defer listener.Close()

	mux := http.NewServeMux()
	// TODO: support embeddings
	mux.HandleFunc("POST /embedding", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "this model does not support embeddings", http.StatusNotImplemented)
	})

	mux.HandleFunc("POST /completion", server.completion)
	mux.HandleFunc("GET /health", server.health)

	httpServer := http.Server{
		Handler: mux,
	}

	log.Println("Server listening on", addr)
	if err := httpServer.Serve(listener); err != nil {
		log.Fatal("server error:", err)
		return err
	}

	return nil
}
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
	"container/heap"
	"math"
	"slices"
)

// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
type tokenHeap []token

func (h tokenHeap) Len() int           { return len(h) }
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *tokenHeap) Push(x any) {
	*h = append(*h, x.(token))
}

func (h *tokenHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// temperature applies scaling to the logits
func temperature(ts *[]token, temp float32) {
	if ts == nil {
		return
	}

	// Ensure temperature clipping near 0 to avoid numerical instability
	temp = max(temp, 1e-7)
	for i := range *ts {
		(*ts)[i].value = (*ts)[i].value / temp
	}
}

// softmax applies normalization to the logits
func softmax(ts *[]token) {
	if ts == nil {
		return
	}

	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range *ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}


Review comment (suggestion):
```suggestion
func temperature(ts []token, temp float32) {
```

This can go back to passing in the slice directly. This is more inline with `sort.Sort` and other stdlib functions that mutate the values in the slice, but not the length or cap itself

Code after:
	"container/heap"
	"math"
	"slices"
)

// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
type tokenHeap []token

func (h tokenHeap) Len() int           { return len(h) }
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *tokenHeap) Push(x any) {
	*h = append(*h, x.(token))
}

func (h *tokenHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// temperature applies scaling to the logits
func temperature(ts []token, temp float32) {
	// Ensure temperature clipping near 0 to avoid numerical instability
	temp = max(temp, 1e-7)
	for i := range ts {
		ts[i].value = ts[i].value / temp
	}
}

// softmax applies normalization to the logits
func softmax(ts []token) {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Compute exp(x - max)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64(v.value - maxLogit)))
		sum += ts[i].value
	}

	// exp(x - max) / sum(exp(x - max))
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, (*ts)[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(*ts); i++ {
		if (*ts)[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, (*ts)[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, len(h))
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	*ts = result
}

// topP limits tokens to those with cumulative probability p
func topP(ts *[]token, p float32) {
	if ts == nil {
		return
	}

	if p == 1.0 {
		return
	}

	// Find cutoff index where cumulative sum exceeds p
	var sum float32
	for i, t := range *ts {
		sum += t.value
		if sum > float32(p) {
			*ts = (*ts)[:i+1]
			return
		}
	}
}

// minP filters tokens with probabilities >= p * max_prob
func minP(ts *[]token, p float32) {
	if ts == nil {
		return
	}


Review comment (suggestion):
`topP` and friends below can stay a regular slice parameter

```suggestion
func topP(ts []token, p float32) {
```

Code after:
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, len(h))
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result
}

// topP limits tokens to those with cumulative probability p
// requires ts to be sorted in descending order of probabilities
func topP(ts []token, p float32) []token {
	if p == 1.0 {
		return ts
	}

	// Find cutoff index where cumulative sum exceeds p
	var sum float32
	for i, t := range ts {
		sum += t.value
		if sum > float32(p) {
			return ts[:i+1]
		}
	}

	return ts
}

// minP filters tokens with probabilities >= p * max_prob
// requires ts to be sorted in descending order of probabilities
func minP(ts []token, p float32) []token {
	maxProb := ts[0].value

	threshold := maxProb * p

	for i, t := range ts {
		if t.value < threshold {
			return ts[:i]
		}
	}
	return ts
}
Repository: ollama/ollama
File: sample/samplers.go (Go)

Code before:
	return t.id, nil
}

// greedy returns the highest probability token from the tokens
func greedy(tokens []token) token {
	max := tokens[0]
	for i := 1; i < len(tokens); i++ {
		if tokens[i].value > max.value {
			max = tokens[i]
		}
	}

	return max
}

// sample returns the highest probability token from the tokens
// given sampler parameters. It also has side effects of modifying the tokens
func (s *Sampler) sample(tokens []token) (token, error) {
	if s.temperature == 0 {
		return greedy(tokens), nil
	}

	// topK also sorts the tokens in descending order of logits
	tokens = topK(tokens, s.topK)

	tokens = temperature(tokens, s.temperature)
	tokens = softmax(tokens)

	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")
	}

	var r float32
	if s.rng != nil {
		r = s.rng.Float32()
	} else {
		r = rand.Float32()
	}

	// Calculate cumulative sum of probabilities
	var sum float32
	for i := range tokens {
		sum += tokens[i].value
		tokens[i].value = sum
	}

Review comment (suggestion):
This is getting into bike-shed territory, but what do you think of a pattern like this? Its more clear to me that the tokens are being modified in this case.
```
tokens.ApplyTemperature(s.temperature)
tokens.ApplySoftmax()
```

Code after:
	return t.id, nil
}

// greedy returns the highest probability token from the tokens
func greedy(tokens []token) token {
	max := tokens[0]
	for i := 1; i < len(tokens); i++ {
		if tokens[i].value > max.value {
			max = tokens[i]
		}
	}

	return max
}

// sample returns the highest probability token from the tokens
// given sampler parameters. It also has side effects of modifying the tokens
func (s *Sampler) sample(tokens []token) (token, error) {
	if s.temperature == 0 {
		return greedy(tokens), nil
	}

	// topK also sorts the tokens in descending order of logits
	tokens = topK(tokens, s.topK)

	// scale and normalize the tokens in place
	temperature(tokens, s.temperature)
	softmax(tokens)

	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")
	}

	var r float32
	if s.rng != nil {
		r = s.rng.Float32()
	} else {
		r = rand.Float32()
	}

	// Calculate cumulative sum of probabilities
	var sum float32
	for i := range tokens {
		sum += tokens[i].value
		tokens[i].value = sum
Repository: ollama/ollama
File: model/input/input.go (Go)

Code before:
package input

// Input represents one token in the input stream
type Input struct {
	// Token is a single element of text.
	Token int32

	// Multimodal is opaque data representing a non-text
	// element such as an image (or part of one if the image
	// can be processed in pieces). It may be either together
	// with Token or on its own.
	Multimodal any

	// MultimodalHash is a unique representation of the data
	// stored in Multimodal, used for caching and comparing
	// equality.
	MultimodalHash uint64

	// SameBatch causes the follow number of tokens to be forced
	// into a single batch, breaking and extending batches as needed.
	// Useful for things like images that must be processed in one
	// shot.
	SameBatch int
}

// MultimodalIndex is a multimodal element (such as an image)
// together with an index into the slice of Inputs with the
// corresponding token. Note that the index is not the same
// as the position - to find that use the index with the
// Positions slice.
type MultimodalIndex struct {
	Index      int
	Multimodal any
}

// Options contains the inputs for a model forward pass
type Options struct {
	Inputs     []int32
	Multimodal []MultimodalIndex
	Positions  []int32
	Sequences  []int
	Outputs    []int32
}


Review comment (performance):
```suggestion
	// SameBatch forces the following number of tokens to be processed
	// in a single batch, breaking and extending batches as needed.
	// Useful for things like images that must be processed in one
	// shot.
	SameBatch int
```

nit: I found this wording a bit confusing

Code after:
package input

// Input represents one token in the input stream
type Input struct {
	// Token is a single element of text.
	Token int32

	// Multimodal is opaque data representing a non-text
	// element such as an image (or part of one if the image
	// can be processed in pieces). It may be either together
	// with Token or on its own.
	Multimodal any

	// MultimodalHash is a unique representation of the data
	// stored in Multimodal, used for caching and comparing
	// equality.
	MultimodalHash uint64

	// SameBatch forces the following number of tokens to be processed
	// in a single batch, breaking and extending batches as needed.
	// Useful for things like images that must be processed in one
	// shot.
	SameBatch int
}

// MultimodalIndex is a multimodal element (such as an image)
// together with an index into the slice of Inputs with the
// corresponding token. Note that the index is not the same
// as the position - to find that use the index with the
// Positions slice.
type MultimodalIndex struct {
	Index      int
	Multimodal any
}

// Options contains the inputs for a model forward pass
type Options struct {
	Inputs     []int32
	Multimodal []MultimodalIndex
	Positions  []int32
	Sequences  []int
	Outputs    []int32
}