diff options
| -rw-r--r-- | tokenizer/byte/tokenizer.go | 94 | ||||
| -rw-r--r-- | tokenizer/byte/tokenizer_test.go | 49 |
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})) +} |
