Fix perplexity calculation

* Track total evaluated tokens
* Use first token as context

4 files changed, 31 insertions, 12 deletions

diff --git a/llm/evaluator.go b/llm/evaluator.go
index ccc554e..9faf483 100644
--- a/llm/evaluator.go
+++ b/llm/evaluator.go
@@ -6,14 +6,14 @@ type Evaluator struct {
 	mutex     sync.RWMutex
 	models    []Causal
 	tokenizer Tokenizer
-	results   []float64
+	results   []result
 	jobs      int
 }
 
 func NewEvaluator() *Evaluator {
 	return &Evaluator{
 		models:  make([]Causal, 0),
-		results: make([]float64, 0),
+		results: make([]result, 0),
 	}
 }
 
diff --git a/llm/job.go b/llm/job.go
index 3e7c5c9..204d324 100644
--- a/llm/job.go
+++ b/llm/job.go
@@ -4,7 +4,7 @@ type job struct {
 	positions []int
 	tokens    [][]int64
 	seen      []int
-	results   []float64
+	results   []result
 	debug     [][2]int
 }
 
@@ -13,7 +13,7 @@ func newJob(batchSize int) *job {
 		positions: make([]int, 0, batchSize),
 		tokens:    make([][]int64, 0, batchSize),
 		seen:      make([]int, 0, batchSize),
-		results:   make([]float64, 0, batchSize),
+		results:   make([]result, 0, batchSize),
 		debug:     make([][2]int, 0),
 	}
 }
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
 }
diff --git a/llm/result.go b/llm/result.go
new file mode 100644
index 0000000..c77f68d
--- /dev/null
+++ b/llm/result.go
@@ -0,0 +1,6 @@
+package llm
+
+type result struct {
+	nll float64
+	n   int
+}