summaryrefslogtreecommitdiff
path: root/llm/evaluator.go
diff options
context:
space:
mode:
Diffstat (limited to 'llm/evaluator.go')
-rw-r--r--llm/evaluator.go18
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
}