Add flag to configure estimation windowmain

1 files changed, 5 insertions, 1 deletions

diff --git a/research/lesci/cmd/lesci/main.go b/research/lesci/cmd/lesci/main.go
index c19d17e..db3e894 100644
--- a/research/lesci/cmd/lesci/main.go
+++ b/research/lesci/cmd/lesci/main.go
@@ -20,6 +20,8 @@ type Config struct {
 	TokObs string
 	TokCtf string
 
+	RdWindow int
+
 	Chkpt string
 
 	OutPath string
@@ -50,6 +52,8 @@ func init() {
 	flag.StringVar(&cfg.TokObs, "tok-obs", "", "observed tokenizer")
 	flag.StringVar(&cfg.TokCtf, "tok-ctf", "", "counterfactual tokenizer")
 
+	flag.IntVar(&cfg.RdWindow, "rd-window", 5000, "estimation window")
+
 	flag.StringVar(&cfg.Chkpt, "c", "", "checkpoint")
 
 	flag.StringVar(&cfg.OutPath, "o", "", "output path")
@@ -133,7 +137,7 @@ func main() {
 	expCfg.ClampRulesBeforeFilter = true
 
 	cutoff := len(bpe.Vocab(t))
-	window := 5000
+	window := cfg.RdWindow
 
 	e := must(lesci.NewExperiment(m, t, c, d, o, cutoff, window, expCfg, expOpts))