diff options
Diffstat (limited to 'llm/evaluator.go')
| -rw-r--r-- | llm/evaluator.go | 28 |
1 files changed, 28 insertions, 0 deletions
diff --git a/llm/evaluator.go b/llm/evaluator.go index 0eb5418..8f2a295 100644 --- a/llm/evaluator.go +++ b/llm/evaluator.go @@ -1,7 +1,11 @@ package llm import ( + "errors" + "sync" "sync/atomic" + + "go.jknobloc.com/x/dataset" ) type Evaluator[R any] struct { @@ -35,3 +39,27 @@ func NewEvaluator[R any](model Causal, tokenizer Tokenizer, callback func(job Jo func (e *Evaluator[R]) Results() chan R { return e.results } + +func (e *Evaluator[R]) RunAndCollect(data dataset.Reader, window, stride int, callback func(R) error) error { + var wg sync.WaitGroup + + var collectErr error + + wg.Add(1) + + go func() { + defer wg.Done() + + for r := range e.results { + if err := callback(r); err != nil && collectErr == nil { + collectErr = err + } + } + }() + + runErr := e.Run(data, window, stride) + + wg.Wait() + + return errors.Join(runErr, collectErr) +} |
