diff options
Diffstat (limited to 'llm')
| -rw-r--r-- | llm/cmd/eval/main.go | 19 |
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 } |