WIP

2 files changed, 143 insertions, 0 deletions

diff --git a/tokenizer/byte/tokenizer.go b/tokenizer/byte/tokenizer.go
new file mode 100644
index 0000000..98982d0
--- /dev/null
+++ b/tokenizer/byte/tokenizer.go
@@ -0,0 +1,94 @@
+package byte
+
+import (
+	"encoding/json"
+	"io"
+
+	"github.com/jonasknobloch/mbpe"
+)
+
+// type parameter for ID type in bpe tokenzier ?!
+
+// Name? Byte is bad; it is a bijection but so are most if not all tokenizers; its the Alphabet?
+// It would be compativle with GenericTokenizer[uint8] without pretokenization
+
+// Tokenizer is a byte level tokenzier covering the all 2^8 bytes; Essentially each individual byte is mapped to a token ID.
+// Serialization follows the standard vocab.json (with HF byte replacements) while omitting merges.txt
+// Pre-tokenization is not necessary; however the tokenized strings can be much larger -> check allocations
+type Tokenizer struct {
+	atoi map[byte]uint8
+	itoa map[uint8]byte
+}
+
+func (t *Tokenizer) Encode(s string) []byte {
+	b := make([]byte, len(s))
+
+	for i := range len(s) {
+		v, ok := t.atoi[s[i]]
+
+		if !ok {
+			panic("unknown byte")
+		}
+
+		b[i] = v
+	}
+
+	return b
+}
+
+func (t *Tokenizer) Decode(ids []uint8) string {
+	b := make([]byte, len(ids))
+
+	for i, id := range ids {
+		v, ok := t.itoa[id]
+
+		if !ok {
+			panic("unknown token ID")
+		}
+
+		b[i] = v
+	}
+
+	return string(b)
+}
+
+func (t *Tokenizer) Tokenize(s string) []int {
+	ids := t.Encode(s)
+
+	r := make([]int, len(ids))
+
+	for i, id := range ids {
+		r[i] = int(id)
+	}
+
+	return r
+}
+
+func NewTokenizer(vocab io.Reader) (*Tokenizer, error) {
+	v := make(map[string]uint8)
+
+	decoder := json.NewDecoder(vocab)
+
+	if err := decoder.Decode(&v); err != nil {
+		return nil, err
+	}
+
+	if len(v) != 256 {
+		panic("vocabulary size != 256")
+	}
+
+	atoi := make(map[byte]uint8, len(v))
+	itoa := make(map[uint8]byte, len(v))
+
+	for char, id := range v {
+		b := mbpe.CharBytes[char]
+
+		atoi[b] = id
+		itoa[id] = b
+	}
+
+	return &Tokenizer{
+		atoi: atoi,
+		itoa: itoa,
+	}, nil
+}
diff --git a/tokenizer/byte/tokenizer_test.go b/tokenizer/byte/tokenizer_test.go
new file mode 100644
index 0000000..d2c8c6c
--- /dev/null
+++ b/tokenizer/byte/tokenizer_test.go
@@ -0,0 +1,49 @@
+package byte
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"go.jknobloc.com/x/shelf"
+)
+
+func TestNewTokenizer_Encode(t *testing.T) {
+	var file *os.File
+
+	if f, err := os.Open(shelf.Abs("tokenizers/minipile/tokenizer_gpt2_256_m000_minipile/vocab.json")); err != nil {
+		t.Fatal(err)
+	} else {
+		file = f
+	}
+
+	defer file.Close()
+
+	tok, err := NewTokenizer(file)
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fmt.Println(tok.Encode(" \nabc"))
+}
+
+func TestNewTokenizer_Decode(t *testing.T) {
+	var file *os.File
+
+	if f, err := os.Open(shelf.Abs("tokenizers/minipile/tokenizer_gpt2_256_m000_minipile/vocab.json")); err != nil {
+		t.Fatal(err)
+	} else {
+		file = f
+	}
+
+	defer file.Close()
+
+	tok, err := NewTokenizer(file)
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fmt.Println(tok.Decode([]uint8{220, 198, 64, 65, 66}))
+}