From e5d7c530ffd7ab88131988a05f66109ec880b809 Mon Sep 17 00:00:00 2001
From: Jonas Knobloch <jonas.knobloch@t-online.de>
Date: Mon, 9 Mar 2026 19:14:50 +0100
Subject: Fix perplexity calculation

* Track total evaluated tokens
* Use first token as context
---
 llm/perplexity.go | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

(limited to 'llm/perplexity.go')

diff --git a/llm/perplexity.go b/llm/perplexity.go
index 58bb886..3c09f60 100644
--- a/llm/perplexity.go
+++ b/llm/perplexity.go
@@ -82,15 +82,17 @@ func (e *Evaluator) Perplexity(data *dataset.Reader, window, stride int) (float6
 
 	<-done
 
-	total := float64(0)
+	totalNLL := float64(0)
+	totalTokens := 0
 
-	for _, nll := range e.results {
-		total += nll
+	for _, r := range e.results {
+		totalNLL += r.nll
+		totalTokens += r.n
 	}
 
-	average := total / float64(len(e.results)) * float64(window-1)
+	average := totalNLL / float64(totalTokens)
 
-	return math.Exp(total / average), nil
+	return math.Exp(average), nil
 }
 
 func (e *Evaluator) schedule(tokens []int64, contextSize, stride, batchSize int) error {
@@ -139,7 +141,7 @@ func (e *Evaluator) schedule(tokens []int64, contextSize, stride, batchSize int)
 	// 1024 to 2047: 1536 upwards
 	// ...
 
-	seen := 0
+	seen := 1 // first token as context
 	n := 0
 
 	// for i := 0; i < len(tokens); i += stride {
@@ -156,6 +158,7 @@ func (e *Evaluator) schedule(tokens []int64, contextSize, stride, batchSize int)
 		j.seen = append(j.seen, seen-i)
 
 		seen = i + contextSize
+
 		n++
 	}
 
@@ -175,6 +178,10 @@ func (e *Evaluator) execute(j *job, device int) {
 		panic("unimplemented")
 	}
 
+	if j.seen[0] < 1 {
+		panic("empty context")
+	}
+
 	m := e.models[device]
 
 	logits := make([][]float32, 0, len(j.tokens[0]))
@@ -185,9 +192,15 @@ func (e *Evaluator) execute(j *job, device int) {
 		panic(err) // TODO handle
 	}
 
-	nll := negLogLikelihood(logits[j.seen[0]:len(logits)-1], toInt(j.tokens[0][1+j.seen[0]:]))
+	nllLogits := logits[j.seen[0]-1 : len(logits)-1]
+	nnlTargets := toInt(j.tokens[0][j.seen[0]:])
+
+	nll := negLogLikelihood(nllLogits, nnlTargets)
 
-	j.results = append(j.results, nll)
+	j.results = append(j.results, result{
+		nll: nll,
+		n:   len(nnlTargets),
+	})
 
 	return
 }
-- 
cgit v1.3.1