Refactor tokenizer usage

1 files changed, 7 insertions, 12 deletions

diff --git a/llm/cmd/eval/main.go b/llm/cmd/eval/main.go
index 758adad..ec9b7e9 100644
--- a/llm/cmd/eval/main.go
+++ b/llm/cmd/eval/main.go
@@ -3,7 +3,7 @@ package main
 import (
 	"log"
 
-	"github.com/jonasknobloch/mbpe"
+	"go.jknobloc.com/x/tokenizer/bpe"
 
 	"go.jknobloc.com/x/dataset"
 	"go.jknobloc.com/x/gpt2"
@@ -36,19 +36,14 @@ func model() *gpt2.Model {
 	return m
 }
 
-func tokenizer() *mbpe.Tokenizer {
-	m := mbpe.NewMBPE()
+func tokenizer() *bpe.Tokenizer {
+	var tok *bpe.Tokenizer
 
-	if err := m.Load("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"); err != nil {
+	if t, err := bpe.NewTokenizerFromFiles("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"); err != nil {
 		log.Fatal(err)
+	} else {
+		tok = t
 	}
 
-	t := mbpe.NewTokenizer(m)
-
-	byteLevel := mbpe.NewByteLevel(false)
-
-	t.SetPreTokenizer(byteLevel)
-	t.SetDecoder(byteLevel)
-
-	return t
+	return tok
 }