1 files changed, 5 insertions, 18 deletions
diff --git a/llm/cmd/eval/ppl.go b/llm/cmd/eval/ppl.go
index 0a028d2..c03c67c 100644
--- a/llm/cmd/eval/ppl.go
+++ b/llm/cmd/eval/ppl.go
@@ -5,7 +5,6 @@ import (
 	"log"
 	"math"
 	"strings"
-	"sync"
 
 	"go.jknobloc.com/x/dataset"
 	"go.jknobloc.com/x/llm"
@@ -31,30 +30,18 @@ func perplexity() {
 		}
 	})
 
-	results := e.Results()
-
-	var wg sync.WaitGroup
-
 	total := float64(0)
 	n := 0
 
-	wg.Add(1)
-
-	go func() {
-		defer wg.Done()
+	if err := e.RunAndCollect(d, 1024, 512, func(r pplResult) error {
+		total += r.v
+		n += r.n
 
-		for r := range results {
-			total += r.v
-			n += r.n
-		}
-	}()
-
-	if err := e.Run(d, 1024, 512); err != nil {
+		return nil
+	}); err != nil {
 		log.Fatal(err)
 	}
 
-	wg.Wait()
-
 	avg := total / float64(n)
 	ppl := math.Exp(avg)