Recover tokenization based on config

11 files changed, 64 insertions, 21 deletions

diff --git a/llm/cmd/eval/main.go b/llm/cmd/eval/main.go
index d31bc32..bd734b3 100644
--- a/llm/cmd/eval/main.go
+++ b/llm/cmd/eval/main.go
@@ -52,7 +52,10 @@ func model() *gpt2.Model {
 func tokenizer() *bpe.Tokenizer {
 	var tok *bpe.Tokenizer
 
-	if t, err := bpe.NewTokenizerFromFiles(shelf.Abs("models/gpt2/vocab.json"), shelf.Abs("models/gpt2/merges.txt")); err != nil {
+	v := shelf.Abs("models/gpt2/vocab.json")
+	m := shelf.Abs("models/gpt2/merges.txt")
+
+	if t, err := bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()); err != nil {
 		log.Fatal(err)
 	} else {
 		tok = t
diff --git a/llmc/cmd/data/fineweb.go b/llmc/cmd/data/fineweb.go
index 3d3d51b..a1bf7df 100644
--- a/llmc/cmd/data/fineweb.go
+++ b/llmc/cmd/data/fineweb.go
@@ -26,7 +26,10 @@ func fineWeb() {
 
 	var tokenizer llm.Tokenizer
 
-	if t, err := bpe.NewTokenizerFromFiles(shelf.Abs("models/gpt2/vocab.json"), shelf.Abs("models/gpt2/merges.txt")); err != nil {
+	v := shelf.Abs("models/gpt2/vocab.json")
+	m := shelf.Abs("models/gpt2/merges.txt")
+
+	if t, err := bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()); err != nil {
 		log.Fatal(err)
 	} else {
 		tokenizer = t
diff --git a/llmc/cmd/data/tinyshakespeare.go b/llmc/cmd/data/tinyshakespeare.go
index 70d4fd4..d71151b 100644
--- a/llmc/cmd/data/tinyshakespeare.go
+++ b/llmc/cmd/data/tinyshakespeare.go
@@ -23,7 +23,10 @@ func tinyShakespeare() {
 
 	var tokenizer *bpe.Tokenizer
 
-	if t, err := bpe.NewTokenizerFromFiles(shelf.Abs("models/gpt2/vocab.json"), shelf.Abs("models/gpt2/merges.txt")); err != nil {
+	v := shelf.Abs("models/gpt2/vocab.json")
+	m := shelf.Abs("models/gpt2/merges.txt")
+
+	if t, err := bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()); err != nil {
 		log.Fatal(err)
 	} else {
 		tokenizer = t
diff --git a/llmc/cmd/debug/debug.go b/llmc/cmd/debug/debug.go
index d1a621a..f076949 100644
--- a/llmc/cmd/debug/debug.go
+++ b/llmc/cmd/debug/debug.go
@@ -16,7 +16,10 @@ func main() {
 	_ = must(llmc.Deserialize(shelf.Abs("llmc/tinyshakespeare/tiny_shakespeare_train.bin"), &data))
 	_ = must(llmc.Deserialize(shelf.Abs("test/llmc/tinyshakespeare/tiny_shakespeare_train.bin"), &gold))
 
-	t := must(bpe.NewTokenizerFromFiles(shelf.Abs("models/gpt2/vocab.json"), shelf.Abs("models/gpt2/merges.txt")))
+	v := shelf.Abs("models/gpt2/vocab.json")
+	m := shelf.Abs("models/gpt2/merges.txt")
+
+	t := must(bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()))
 
 	itoa := bpe.Itoa(t)
 
diff --git a/llmc/cmd/peek/main.go b/llmc/cmd/peek/main.go
index 0a5eeab..1bfcca8 100644
--- a/llmc/cmd/peek/main.go
+++ b/llmc/cmd/peek/main.go
@@ -18,7 +18,10 @@ func main() {
 	_ = must(llmc.Deserialize(shelf.Abs("llmc/edu_fineweb100B/edu_fineweb_val_000000.bin"), &a))
 	_ = must(llmc.Deserialize(shelf.Abs("test/llmc/edu_fineweb100B/edu_fineweb_val_000000.bin"), &b))
 
-	t := must(bpe.NewTokenizerFromFiles(shelf.Abs("models/gpt2/vocab.json"), shelf.Abs("models/gpt2/merges.txt")))
+	v := shelf.Abs("models/gpt2/vocab.json")
+	m := shelf.Abs("models/gpt2/merges.txt")
+
+	t := must(bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()))
 
 	s := decode(&a, 1024, t)
 	k := decode(&b, 1024, t)
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()
diff --git a/tokenizer/bpe/cmd/tokenize/main.go b/tokenizer/bpe/cmd/tokenize/main.go
index 13926b7..e6cab4b 100644
--- a/tokenizer/bpe/cmd/tokenize/main.go
+++ b/tokenizer/bpe/cmd/tokenize/main.go
@@ -100,7 +100,10 @@ func data() *dataset.ParquetReader {
 func tokenizer() llm.Tokenizer {
 	var tok *bpe.Tokenizer
 
-	if t, err := bpe.NewTokenizerFromFiles(shelf.Abs("models/gpt2/vocab.json"), shelf.Abs("models/gpt2/merges.txt")); err != nil {
+	v := shelf.Abs("models/gpt2/vocab.json")
+	m := shelf.Abs("models/gpt2/merges.txt")
+
+	if t, err := bpe.NewTokenizerFromFiles(v, m, bpe.DefaultConfig()); err != nil {
 		log.Fatal(err)
 	} else {
 		tok = t
diff --git a/tokenizer/bpe/config.go b/tokenizer/bpe/config.go
new file mode 100644
index 0000000..aef62c9
--- /dev/null
+++ b/tokenizer/bpe/config.go
@@ -0,0 +1,11 @@
+package bpe
+
+type Config struct {
+	Recover bool
+}
+
+func DefaultConfig() Config {
+	return Config{
+		Recover: false,
+	}
+}
diff --git a/tokenizer/bpe/tokenizer.go b/tokenizer/bpe/tokenizer.go
index 9bc351c..7e2459a 100644
--- a/tokenizer/bpe/tokenizer.go
+++ b/tokenizer/bpe/tokenizer.go
@@ -3,12 +3,14 @@ package bpe
 import "github.com/jonasknobloch/mbpe"
 
 type Tokenizer struct {
-	mbpe *mbpe.Tokenizer
+	mbpe   *mbpe.Tokenizer
+	config Config
 }
 
-func NewTokenizer(mbpe *mbpe.Tokenizer) *Tokenizer {
+func NewTokenizer(mbpe *mbpe.Tokenizer, cfg Config) *Tokenizer {
 	return &Tokenizer{
-		mbpe: mbpe,
+		mbpe:   mbpe,
+		config: cfg,
 	}
 }
 
@@ -29,13 +31,15 @@ func (t *Tokenizer) Decode(ids []int) string {
 }
 
 func (t *Tokenizer) Tokenize(s string) (ids []int) {
-	defer func() {
-		if r := recover(); r != nil {
-			// defer UnknownRunes(t, s)
-
-			ids = []int{}
-		}
-	}()
+	if t.config.Recover {
+		defer func() {
+			if r := recover(); r != nil {
+				// defer UnknownRunes(t, s)
+
+				ids = []int{}
+			}
+		}()
+	}
 
 	return t.Encode(s)
 }
diff --git a/tokenizer/bpe/utility.go b/tokenizer/bpe/utility.go
index fbe99fc..f8fcdbc 100644
--- a/tokenizer/bpe/utility.go
+++ b/tokenizer/bpe/utility.go
@@ -2,7 +2,7 @@ package bpe
 
 import "github.com/jonasknobloch/mbpe"
 
-func NewTokenizerFromFiles(vocab, merges string) (*Tokenizer, error) {
+func NewTokenizerFromFiles(vocab, merges string, cfg Config) (*Tokenizer, error) {
 	model := mbpe.NewMBPE()
 
 	if err := model.Load(vocab, merges); err != nil {
@@ -16,7 +16,7 @@ func NewTokenizerFromFiles(vocab, merges string) (*Tokenizer, error) {
 	tokenizer.SetPreTokenizer(pre)
 	tokenizer.SetDecoder(pre)
 
-	return NewTokenizer(tokenizer), nil
+	return NewTokenizer(tokenizer, cfg), nil
 }
 
 func Vocab(t *Tokenizer) []string {