Recover tokenization based on config

4 files changed, 30 insertions, 12 deletions

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)
+	if t.config.Recover {
+		defer func() {
+			if r := recover(); r != nil {
+				// defer UnknownRunes(t, s)
 
-			ids = []int{}
-		}
-	}()
+				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 {