Add tokenizer module

* Wrap MBPE tokenizer

4 files changed, 116 insertions, 0 deletions

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
+}
diff --git a/tokenizer/go.mod b/tokenizer/go.mod
new file mode 100644
index 0000000..e66e1f8
--- /dev/null
+++ b/tokenizer/go.mod
@@ -0,0 +1,3 @@
+module go.jknobloc.com/x/tokenizer
+
+go 1.25