diff options
Diffstat (limited to 'research/lesci/experiment.go')
| -rw-r--r-- | research/lesci/experiment.go | 38 |
1 files changed, 29 insertions, 9 deletions
diff --git a/research/lesci/experiment.go b/research/lesci/experiment.go index d01d405..53db931 100644 --- a/research/lesci/experiment.go +++ b/research/lesci/experiment.go @@ -14,15 +14,17 @@ import ( ) type Experiment struct { - model llm.Causal - tokenizer llm.Tokenizer - counterfactual llm.Tokenizer - data dataset.Reader - name string - cutoff int - window int - config Config - options Options + model llm.Causal + tokenizer llm.Tokenizer + counterfactual llm.Tokenizer + data dataset.Reader + name string + cutoff int + window int + config Config + options Options + tokenBufferConfig llm.TokenBufferConfig + evaluatorConfig llm.EvaluatorConfig } func NewExperiment(model llm.Causal, tokenizer, counterfactual llm.Tokenizer, data dataset.Reader, name string, cutoff, window int, cfg Config, opts Options) (*Experiment, error) { @@ -38,9 +40,27 @@ func NewExperiment(model llm.Causal, tokenizer, counterfactual llm.Tokenizer, da options: opts, } + e.tokenBufferConfig = llm.TokenBufferConfig{ + Window: 1024, + Stride: 512, + } + + e.evaluatorConfig = llm.EvaluatorConfig{ + BatchSize: 32, + NumWorkers: 64, + } + return e, nil } +func (e *Experiment) SetTokenBufferConfig(cfg llm.TokenBufferConfig) { + e.tokenBufferConfig = cfg +} + +func (e *Experiment) SetEvaluatorConfig(cfg llm.EvaluatorConfig) { + e.evaluatorConfig = cfg +} + func (e *Experiment) Run() error { if err := os.MkdirAll(e.name, 0775); err != nil { log.Fatal(err) |