tokenizer.TrainBPE — Train BPE Merge Rules from Text

TrainBPE implements the Byte Pair Encoding training algorithm, iteratively finding and merging the most frequent adjacent token pair in a corpus until the target vocabulary size is reached. The resulting merge rules are stored in a DefaultDict[Pair] and can be persisted to disk with Save and restored with Load for use with NewBPETokenizer.

TrainBPE()

func TrainBPE(inputText string, vocabSize int, endToken ...string) *DefaultDict[Pair]

Trains a BPE vocabulary on a text corpus and returns the learned merge rules.

inputText

string

required

The full training corpus as a single string. Multiple documents can be concatenated with <|endoftext|> separators to prevent merges from crossing document boundaries.

vocabSize

int

required

Target vocabulary size. Must be greater than 257 (256 base byte tokens + 1 end token) to produce at least one merge. The number of merge operations performed is vocabSize - 256 - 1.

endToken

...string

Optional end-of-text token string used as a document separator. Defaults to "<|endoftext|>". Occurrences in the corpus are treated as boundaries: no pair is ever merged across an end token.

Returns: *DefaultDict[Pair] — the trained merge rules, where each key is a Pair of adjacent token IDs and the value is the new merged token ID assigned during training.

Algorithm

Split on end token — the corpus is split on the end token string using strings.Split. This prevents any merge from bridging two documents.
Pre-tokenize — each segment is broken into pre-tokens with the GPT-2 style regex `'(?:[sdmt]|ll|ve|re)| ?\pL+| ?\pN+| ?[^\s\pL\pN]+|\s+`, then each pre-token is converted to a sequence of byte IDs (0–255).
Build initial pair counts — all adjacent byte ID pairs across all pre-tokens are counted, weighted by the frequency of each unique pre-token string. A Cache is built that maps each pair to the set of pre-token keys that contain it.
Iterative merging (vocabSize - 256 - 1 steps):
- Find the Pair with the highest count in pairCounts.
- Assign it the next available ID: 256 + step (starting at 256).
- Record the rule in mergeRules.
- For every pre-token affected (looked up via the cache), apply the merge, then subtract the old pair counts and add the new pair counts produced by the merged sequence.
Return mergeRules — a DefaultDict[Pair] whose insertion order reflects the merge priority used by BPETokenizer.Encode.

sample := "Hello world!<|endoftext|>This is BPE training."
mergeRules := tokenizer.TrainBPE(sample, 260)
for pair, newID := range mergeRules.Seq2() {
    fmt.Println(pair, newID)
}
// [105 115] 256
// [105 110] 257
// [72 101] 258

The Writer package-level variable controls where the training progress bar is written. It defaults to io.Discard (silent). Set it to os.Stdout to watch training progress:

tokenizer.Writer = os.Stdout
mergeRules := tokenizer.TrainBPE(data, 1000)

Save() and Load()

Merge rules are serialized using Go’s encoding/gob package. The DefaultDict[Pair] struct — including both its Dict map and its insertion-ordered Order slice — is encoded in full, so the merge priority order is preserved exactly across save/load cycles.

Save()

func Save(path string, dict *DefaultDict[Pair]) error

Serializes mergeRules to a gob-encoded file at path. Creates the file (truncating any existing content). Returns an error if the file cannot be created or if gob encoding fails.

path

string

required

Destination file path. Conventionally uses the .gob extension.

dict

*DefaultDict[Pair]

required

The trained merge rules to persist, as returned by TrainBPE.

Returns: error

Load()

func Load(path string) (*DefaultDict[Pair], error)

Deserializes merge rules from a gob-encoded file previously written by Save. Returns a pointer to the reconstructed DefaultDict[Pair] with the original insertion order intact.

path

string

required

Path to the gob file to read.

Returns: (*DefaultDict[Pair], error) — returns an error if the file cannot be opened or the gob stream cannot be decoded.

Usage

data, err := os.ReadFile("corpus.txt")
if err != nil {
    log.Fatal(err)
}

// Train and save
mergeRules := tokenizer.TrainBPE(string(data), 1000)
if err := tokenizer.Save("merge_rules.gob", mergeRules); err != nil {
    log.Fatal(err)
}

// Load and build tokenizer
mergeRules, err = tokenizer.Load("merge_rules.gob")
if err != nil {
    log.Fatal(err)
}
tknizer := tokenizer.NewBPETokenizer(mergeRules)
fmt.Println(tknizer.VocabSize) // 1000

Pair type

type Pair [2]int  // [tokenID_left, tokenID_right]

A Pair represents two adjacent token IDs. It is the key type for the merge rules dictionary returned by TrainBPE. Pair{a, b} means: whenever token ID a is immediately followed by token ID b, replace both with the merged token ID stored as the corresponding value.

// Manually constructing merge rules for testing
mergeRules := tokenizer.NewDefaultDict[tokenizer.Pair]()
mergeRules.Set(tokenizer.Pair{105, 115}, 256) // bytes for 'i','s' → token 256
mergeRules.Set(tokenizer.Pair{256, 32},  257) // token 256 + space  → token 257

Because Pair is an array (not a slice), it is directly usable as a map key and is gob-encodable without any custom serialization.

DefaultDict[K]

DefaultDict[K comparable] is an ordered generic map that assigns integer values to keys of type K. It preserves insertion order via its Order slice, which is critical for merge rules: the order encodes training priority.

type DefaultDict[T comparable] struct {
    Dict  map[T]int
    Order []T
}

Method	Signature	Description
`NewDefaultDict`	`NewDefaultDict[K]() *DefaultDict[K]`	Creates an empty `DefaultDict` with initialized internal structures.
`Set`	`Set(key K, value int)`	Sets an absolute value for `key`. Appends to `Order` on first insertion.
`Get`	`Get(key K) (int, bool)`	Returns the value and an existence flag.
`Value`	`Value(key K) int`	Returns the value for `key`, or `0` if absent (no-error variant).
`Incr`	`Incr(key K, delta int)`	Adds `delta` to the current value. Useful for counting pair frequencies.
`Delete`	`Delete(key K)`	Removes `key` from both `Dict` and `Order`.
`Len`	`Len() int`	Returns the number of entries currently in the dictionary.
`Seq2`	`Seq2() iter.Seq2[K, int]`	Returns a range-compatible iterator that yields `(key, value)` pairs in insertion order.

Iterate over merge rules in training order using Seq2 with Go 1.23+ range-over-function syntax:

for pair, newID := range mergeRules.Seq2() {
    fmt.Printf("merge %v → %d\n", pair, newID)
}

DefaultDict is not safe for concurrent reads and writes. If you share a DefaultDict across goroutines, protect it with a mutex or use separate instances per goroutine.

Model

Tokenizer

Layers

Data & Utilities

tokenizer.TrainBPE — Train BPE Merge Rules from Text

TrainBPE()

Algorithm

Save() and Load()

Save()

Load()

Usage

Pair type

DefaultDict[K]

Build docs developers (and LLMs) love

Model

Tokenizer

Layers

Data & Utilities

Documentation Index

​TrainBPE()

​Algorithm

​Save() and Load()

​Save()

​Load()

​Usage

​Pair type

​DefaultDict[K]

Build docs developers (and LLMs) love

TrainBPE()

Algorithm

Save() and Load()

Save()

Load()

Usage

Pair type

DefaultDict[K]