Update perplexity evaluator

* Introduce results channel
* Introduce jobs channel
* Generalize evaluation via callback

1 files changed, 32 insertions, 11 deletions

diff --git a/llm/evaluator.go b/llm/evaluator.go
index 9faf483..24fe2d4 100644
--- a/llm/evaluator.go
+++ b/llm/evaluator.go
@@ -1,26 +1,47 @@
 package llm
 
-import "sync"
+import (
+	"sync/atomic"
+)
 
-type Evaluator struct {
-	mutex     sync.RWMutex
+type Evaluator[R any] struct {
 	models    []Causal
 	tokenizer Tokenizer
-	results   []result
-	jobs      int
+
+	batchSize  int
+	numWorkers int
+
+	scheduled atomic.Int64
+	completed atomic.Int64
+
+	jobs    chan batch
+	results chan R
+
+	callback func(job Job, logits [][]float32, tokens []int) R
 }
 
-func NewEvaluator() *Evaluator {
-	return &Evaluator{
-		models:  make([]Causal, 0),
-		results: make([]result, 0),
+func NewEvaluator[R any]() *Evaluator[R] {
+	return &Evaluator[R]{
+		models:     make([]Causal, 0),
+		batchSize:  1, // TODO arg
+		numWorkers: 4, // TODO arg
+		jobs:       make(chan batch, 1024),
+		results:    make(chan R, 1024),
 	}
 }
 
-func (e *Evaluator) AddModel(model Causal) {
+func (e *Evaluator[R]) AddModel(model Causal) {
 	e.models = append(e.models, model)
 }
 
-func (e *Evaluator) SetTokenizer(tokenizer Tokenizer) {
+func (e *Evaluator[R]) SetTokenizer(tokenizer Tokenizer) {
 	e.tokenizer = tokenizer
 }
+
+func (e *Evaluator[R]) SetCallback(callback func(job Job, logits [][]float32, tokens []int) R) {
+	e.callback = callback
+}
+
+func (e *Evaluator[R]) Results() chan R {
+	return e.results
+}