Refactor evaluator initialization

1 files changed, 4 insertions, 14 deletions

diff --git a/llm/evaluator.go b/llm/evaluator.go
index 24fe2d4..0eb5418 100644
--- a/llm/evaluator.go
+++ b/llm/evaluator.go
@@ -20,28 +20,18 @@ type Evaluator[R any] struct {
 	callback func(job Job, logits [][]float32, tokens []int) R
 }
 
-func NewEvaluator[R any]() *Evaluator[R] {
+func NewEvaluator[R any](model Causal, tokenizer Tokenizer, callback func(job Job, logits [][]float32, tokens []int) R) *Evaluator[R] {
 	return &Evaluator[R]{
-		models:     make([]Causal, 0),
+		models:     []Causal{model}, // TODO multiple devices
+		tokenizer:  tokenizer,
 		batchSize:  1, // TODO arg
 		numWorkers: 4, // TODO arg
 		jobs:       make(chan batch, 1024),
 		results:    make(chan R, 1024),
+		callback:   callback,
 	}
 }
 
-func (e *Evaluator[R]) AddModel(model Causal) {
-	e.models = append(e.models, model)
-}
-
-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
 }