Recover tokenization based on config

2 files changed, 13 insertions, 3 deletions

diff --git a/research/lesci/cmd/lesci/main.go b/research/lesci/cmd/lesci/main.go
index c52013a..27c257a 100644
--- a/research/lesci/cmd/lesci/main.go
+++ b/research/lesci/cmd/lesci/main.go
@@ -45,8 +45,14 @@ func setup(control, treatment int) (*lesci.Experiment, *gpt2.Model) {
 	b := fmt.Sprintf(shelf.Abs("models/mbpe/gpt2_%d_m000_babylm_v2"), 100512)
 
 	m := must(model(path.Join(a, "model_eval.onnx"), control))
-	t := must(bpe.NewTokenizerFromFiles(path.Join(a, "vocab.json"), path.Join(a, "merges.txt")))
-	c := must(bpe.NewTokenizerFromFiles(path.Join(b, "vocab.json"), path.Join(b, "merges.txt")))
+
+	cfg := bpe.Config{
+		Recover: true,
+	}
+
+	t := must(bpe.NewTokenizerFromFiles(path.Join(a, "vocab.json"), path.Join(a, "merges.txt"), cfg))
+	c := must(bpe.NewTokenizerFromFiles(path.Join(b, "vocab.json"), path.Join(b, "merges.txt"), cfg))
+
 	d := must(dataset.NewFileReader(shelf.Abs("data/babylm/train_100M"), "*.train"))
 
 	o := fmt.Sprintf(shelf.Abs("results/lesci/m000/babylm_%d_%d"), control, treatment)
diff --git a/research/sander/cmd/sander/main.go b/research/sander/cmd/sander/main.go
index 6f68434..56f1488 100644
--- a/research/sander/cmd/sander/main.go
+++ b/research/sander/cmd/sander/main.go
@@ -28,7 +28,11 @@ func main() {
 }
 
 func run(src, dst string) error {
-	t := must(bpe.NewTokenizerFromFiles(path.Join(src, "vocab.json"), path.Join(src, "merges.txt")))
+	cfg := bpe.Config{
+		Recover: true,
+	}
+
+	t := must(bpe.NewTokenizerFromFiles(path.Join(src, "vocab.json"), path.Join(src, "merges.txt"), cfg))
 	e := must(sander.NewExperiment(dst, path.Join(src, "model.onnx"), sander.UnusedTokensMBPE(t)))
 
 	return e.Run()