From ac670c698540488c818a17c3749abf3ed4cb764a Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Tue, 24 Mar 2026 13:02:49 +0100 Subject: Add tokenizer module * Wrap MBPE tokenizer --- tokenizer/bpe/cmd/tokenize/main.go | 68 ++++++++++++++++++++++++++++++++++++++ tokenizer/bpe/tokenizer.go | 25 ++++++++++++++ tokenizer/bpe/utility.go | 20 +++++++++++ 3 files changed, 113 insertions(+) create mode 100644 tokenizer/bpe/cmd/tokenize/main.go create mode 100644 tokenizer/bpe/tokenizer.go create mode 100644 tokenizer/bpe/utility.go (limited to 'tokenizer/bpe') diff --git a/tokenizer/bpe/cmd/tokenize/main.go b/tokenizer/bpe/cmd/tokenize/main.go new file mode 100644 index 0000000..aaf10e7 --- /dev/null +++ b/tokenizer/bpe/cmd/tokenize/main.go @@ -0,0 +1,68 @@ +package main + +import ( + "log" + "sync/atomic" + "time" + + "go.jknobloc.com/x/dataset" + "go.jknobloc.com/x/llm" + "go.jknobloc.com/x/tokenizer/bpe" + "go.jknobloc.com/x/tui" +) + +func main() { + reader := data() + + t := tokenizer() + + pb := tui.NewProgressBar("Tokenize", 20, 1000, time.Now()) + + var processed atomic.Int64 + + pb.Start(1*time.Second, func() int { + return int(processed.Load()) + }) + + defer pb.Close() + + n := 0 + + for d := range reader.Texts("text") { + if n >= 1000 { + break + } + + tokens := t.Tokenize(d) + + _ = tokens + + processed.Add(1) + + n++ + } +} + +func data() *dataset.ParquetReader { + var simple *dataset.ParquetReader + + if r, err := dataset.NewParquetReader("dataset/cmd/dataset/tmp/wikipedia/simple/train"); err != nil { + log.Fatal(err) + } else { + simple = r + } + + return simple +} + +func tokenizer() llm.Tokenizer { + var tok *bpe.Tokenizer + + if t, err := bpe.NewTokenizerFromFiles("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"); err != nil { + log.Fatal(err) + } else { + tok = t + } + + return tok +} diff --git a/tokenizer/bpe/tokenizer.go b/tokenizer/bpe/tokenizer.go new file mode 100644 index 0000000..1b65de2 --- /dev/null +++ b/tokenizer/bpe/tokenizer.go @@ -0,0 +1,25 @@ +package bpe + +import "github.com/jonasknobloch/mbpe" + +type Tokenizer struct { + mbpe *mbpe.Tokenizer +} + +func NewTokenizer(mbpe *mbpe.Tokenizer) *Tokenizer { + return &Tokenizer{ + mbpe: mbpe, + } +} + +func (t *Tokenizer) Encode(s string) []int { + return t.mbpe.Tokenize(s) +} + +func (t *Tokenizer) Decode(ids []int) string { + panic("unimplemented") // TODO implement +} + +func (t *Tokenizer) Tokenize(s string) []int { + return t.Encode(s) +} diff --git a/tokenizer/bpe/utility.go b/tokenizer/bpe/utility.go new file mode 100644 index 0000000..8ad0ca2 --- /dev/null +++ b/tokenizer/bpe/utility.go @@ -0,0 +1,20 @@ +package bpe + +import "github.com/jonasknobloch/mbpe" + +func NewTokenizerFromFiles(vocab, merges string) (*Tokenizer, error) { + model := mbpe.NewMBPE() + + if err := model.Load(vocab, merges); err != nil { + return nil, err + } + + tokenizer := mbpe.NewTokenizer(model) + + pre := mbpe.NewByteLevel(false) + + tokenizer.SetPreTokenizer(pre) + tokenizer.SetDecoder(pre) + + return NewTokenizer(tokenizer), nil +} -- cgit v1.3.1