A Tokenizer

// bpe.go — train a byte-pair encoder, encode, decode, and check the round trip.
package main

import (
	"fmt"
	"strings"
	"unicode/utf8"
)

type Pair [2]int32

type BPE struct {
	Merges map[Pair]int32 // pair → new token ID; lower ID = learned earlier
	Vocab  [][]byte       // token ID → bytes
}

// replace rewrites ids, replacing each occurrence of pair p with id, in place.
func replace(ids []int32, p Pair, id int32) []int32 {
	out := ids[:0]
	for i := 0; i < len(ids); i++ {
		if i+1 < len(ids) && ids[i] == p[0] && ids[i+1] == p[1] {
			out = append(out, id)
			i++
		} else {
			out = append(out, ids[i])
		}
	}
	return out
}

func Train(corpus string, merges int) *BPE {
	t := &BPE{Merges: map[Pair]int32{}}
	for b := 0; b < 256; b++ {
		t.Vocab = append(t.Vocab, []byte{byte(b)})
	}
	ids := make([]int32, len(corpus))
	for i := 0; i < len(corpus); i++ {
		ids[i] = int32(corpus[i])
	}
	for m := 0; m < merges; m++ {
		counts := map[Pair]int{}
		for i := 0; i+1 < len(ids); i++ {
			counts[Pair{ids[i], ids[i+1]}]++
		}
		best, bestN := Pair{}, 1
		for p, n := range counts {
			// most frequent; ties broken deterministically so training is reproducible
			if n > bestN || (n == bestN && n > 1 && (p[0] < best[0] || (p[0] == best[0] && p[1] < best[1]))) {
				best, bestN = p, n
			}
		}
		if bestN < 2 {
			break // nothing repeats any more
		}
		id := int32(len(t.Vocab))
		t.Merges[best] = id
		t.Vocab = append(t.Vocab, append(append([]byte{}, t.Vocab[best[0]]...), t.Vocab[best[1]]...))
		ids = replace(ids, best, id)
	}
	return t
}

// Encode appends the tokens for text to dst (caller-owned buffer: no allocation when reused).
func (t *BPE) Encode(dst []int32, text string) []int32 {
	start := len(dst)
	for i := 0; i < len(text); i++ {
		dst = append(dst, int32(text[i]))
	}
	ids := dst[start:]
	for len(ids) >= 2 {
		// find the pair in the sequence that was learned earliest
		best, bestID := Pair{}, int32(-1)
		for i := 0; i+1 < len(ids); i++ {
			if id, ok := t.Merges[Pair{ids[i], ids[i+1]}]; ok && (bestID < 0 || id < bestID) {
				best, bestID = Pair{ids[i], ids[i+1]}, id
			}
		}
		if bestID < 0 {
			break
		}
		ids = replace(ids, best, bestID)
	}
	return dst[:start+len(ids)]
}

func (t *BPE) Decode(ids []int32) string {
	var sb strings.Builder
	for _, id := range ids {
		sb.Write(t.Vocab[id])
	}
	return sb.String()
}

// show renders a token for printing: text when it is valid UTF-8, hex when it is a fragment
// of a multi-byte character.
func show(b []byte) string {
	if utf8.Valid(b) {
		return strings.ReplaceAll(string(b), " ", "_")
	}
	return fmt.Sprintf("0x%x", b)
}

func main() {
	corpus := strings.Repeat("the lower the latency the lower the cost. low latency, low cost, lowest loss. ", 20) +
		"tokens are the unit of cost and of latency. "

	tok := Train(corpus, 60)
	fmt.Printf("vocabulary: 256 bytes + %d merges\n", len(tok.Merges))
	fmt.Print("first learned tokens: ")
	for id := 256; id < 256+12 && id < len(tok.Vocab); id++ {
		fmt.Printf("%q ", tok.Vocab[id])
	}
	fmt.Println()

	for _, text := range []string{
		"the lowest latency",
		"low cost tokens",
		"unseen wörds and 🙂 still encode",
	} {
		ids := tok.Encode(nil, text)
		fmt.Printf("\n%q\n  %d bytes → %d tokens: ", text, len(text), len(ids))
		for _, id := range ids {
			fmt.Printf("[%s]", show(tok.Vocab[id]))
		}
		back := tok.Decode(ids)
		fmt.Printf("\n  round trip exact: %v\n", back == text)
	}

	// Reusing the output buffer: the hot-path shape from III.05.
	buf := make([]int32, 0, 256)
	total := 0
	for i := 0; i < 1000; i++ {
		buf = tok.Encode(buf[:0], "the lowest latency and the lowest cost")
		total += len(buf)
	}
	fmt.Printf("\nencoded 1,000 prompts into one reused buffer: %d tokens\n", total)
}