GPT-2 Byte-Pair Encoding

Lean-native support for GPT-2-style byte-level BPE tokenizers.

This module lives in NN.API.Text rather than a model file: any Transformer, diffusion LM, or verifier that wants GPT-2-compatible tokenization should share the same implementation. The implementation parses the standard vocab.json and merges.txt files directly in Lean.

The pre-tokenizer implements the GPT-2 regex shape:

's|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+

The Unicode \p{L}, \p{N}, and \s predicates are supplied by NN.API.Text.Unicode, rather than Lean's ASCII-oriented Char.isAlpha / Char.isDigit helpers.

Data

structure NN.API.text.Gpt2Bpe.VocabEntry : Type

One token-to-id entry from GPT-2's vocab.json.

• token : String

  Token spelling after GPT-2 byte-to-unicode escaping.

• id : ℕ

  Token id.

def NN.API.text.Gpt2Bpe.instReprVocabEntry.repr : VocabEntry → ℕ → Std.Format

@[implicit_reducible]

instance NN.API.text.Gpt2Bpe.instReprVocabEntry : Repr VocabEntry

def NN.API.text.Gpt2Bpe.instDecidableEqVocabEntry.decEq (x✝ x✝¹ : VocabEntry) : Decidable (x✝ = x✝¹)

@[implicit_reducible]

instance NN.API.text.Gpt2Bpe.instDecidableEqVocabEntry : DecidableEq VocabEntry

structure NN.API.text.Gpt2Bpe.MergeRank : Type

One ranked merge from GPT-2's merges.txt. Lower rank is applied earlier.

• left : String

  Left symbol.

• right : String

  Right symbol.

• rank : ℕ

  Merge priority.

def NN.API.text.Gpt2Bpe.instReprMergeRank.repr : MergeRank → ℕ → Std.Format

@[implicit_reducible]

instance NN.API.text.Gpt2Bpe.instReprMergeRank : Repr MergeRank

def NN.API.text.Gpt2Bpe.instDecidableEqMergeRank.decEq (x✝ x✝¹ : MergeRank) : Decidable (x✝ = x✝¹)

@[implicit_reducible]

instance NN.API.text.Gpt2Bpe.instDecidableEqMergeRank : DecidableEq MergeRank

structure NN.API.text.Gpt2Bpe.Tokenizer : Type

Loaded GPT-2 BPE tokenizer.

• vocab : Array VocabEntry

  Token vocabulary as loaded from vocab.json.

• merges : Array MergeRank

  Ranked merge table from merges.txt.

• vocabMap : Std.HashMap String ℕ

  Fast token-to-id lookup derived from vocab.

• idMap : Std.HashMap ℕ String

  Fast id-to-token lookup derived from vocab.

• mergeMap : Std.HashMap (String × String) ℕ

  Fast pair-to-rank lookup derived from merges.

Byte Escaping

def NN.API.text.Gpt2Bpe.bytesVisible : List ℕ

Bytes that GPT-2 leaves at their visible Unicode code points.

def NN.API.text.Gpt2Bpe.bytesLatin1A : List ℕ

First Latin-1 byte range kept visible by GPT-2 byte escaping.

def NN.API.text.Gpt2Bpe.bytesLatin1B : List ℕ

Second Latin-1 byte range kept visible by GPT-2 byte escaping.

def NN.API.text.Gpt2Bpe.baseBytes : List ℕ

Bytes that do not need synthetic code points in GPT-2 byte escaping.

def NN.API.text.Gpt2Bpe.containsNat (xs : List ℕ) (x : ℕ) : Bool

Boolean membership test used while constructing the byte escape table.

def NN.API.text.Gpt2Bpe.byteCodeTable : Array ℕ

GPT-2 byte-to-Unicode code-point table for all 256 byte values.

def NN.API.text.Gpt2Bpe.byteToChar (b : UInt8) : Char

GPT-2 byte-to-unicode escape for one byte.

def NN.API.text.Gpt2Bpe.charToByte? (c : Char) : Option UInt8

Inverse of byteToChar, used when decoding BPE token strings back to UTF-8.

def NN.API.text.Gpt2Bpe.byteEncode (s : String) : String

Reversible GPT-2 byte-to-unicode escape for a string fragment.

def NN.API.text.Gpt2Bpe.byteDecode? (s : String) : Option String

Decode GPT-2 byte-to-unicode escaped text back into a UTF-8 string.

Pre-tokenization

inductive NN.API.text.Gpt2Bpe.RegexClass : Type

Character classes used by the GPT-2 pre-tokenizer branches.

• letter : RegexClass
• number : RegexClass
• other : RegexClass

@[implicit_reducible]

instance NN.API.text.Gpt2Bpe.instDecidableEqRegexClass : DecidableEq RegexClass

def NN.API.text.Gpt2Bpe.isRegexOther (c : Char) : Bool

Character predicate for GPT-2's non-whitespace, non-letter, non-number regex branch.

def NN.API.text.Gpt2Bpe.matchesRegexClass (cls : RegexClass) (c : Char) : Bool

Test whether a character belongs to one of the GPT-2 regex classes.

def NN.API.text.Gpt2Bpe.takeWhileChars (p : Char → Bool) : List Char → List Char × List Char

Split a character list at the first character that fails p.

def NN.API.text.Gpt2Bpe.consumeContraction? : List Char → Option (String × List Char)

Consume one GPT-2 contraction token such as 's or 'll, if present.

def NN.API.text.Gpt2Bpe.consumeClassRun? (cls : RegexClass) : List Char → Option (String × List Char)

Consume one GPT-2 letter/number/other run, allowing a leading ASCII space.

def NN.API.text.Gpt2Bpe.consumeWhitespaceNotFollowedByNonspace? (xs : List Char) : Option (String × List Char)

Consume the GPT-2 branch \s+(?!\S).

Python's regex engine greedily takes a whitespace run but may backtrack so the negative lookahead sees either end-of-input or another whitespace character. For a whitespace run before a non-space token, this consumes all but the final whitespace; the final ASCII space can then attach to the next letter/number/punctuation branch, matching GPT-2's standard token boundaries.

def NN.API.text.Gpt2Bpe.consumeWhitespaceRun? (xs : List Char) : Option (String × List Char)

Consume a plain whitespace run when the lookahead-sensitive branch did not apply.

def NN.API.text.Gpt2Bpe.pretokenizeAux : ℕ → List Char → List String

Fuel-bounded worker for GPT-2 regex pre-token fragments before byte escaping and BPE merges.

The branch order mirrors GPT-2's tokenizer regex exactly: contractions, optional-space letter runs, optional-space number runs, optional-space non-space/non-letter/non-number runs, whitespace not followed by non-space, and finally a plain whitespace run. The fuel argument keeps this definition total; pretokenize supplies enough fuel for the whole input.

def NN.API.text.Gpt2Bpe.pretokenize (s : String) : List String

Split a string into GPT-2-style pre-token fragments.

BPE Merging

def NN.API.text.Gpt2Bpe.vocabId? (tok : Tokenizer) (s : String) : Option ℕ

Look up a token id in a loaded tokenizer.

def NN.API.text.Gpt2Bpe.tokenString? (tok : Tokenizer) (id : ℕ) : Option String

Look up the token spelling for a token id.

def NN.API.text.Gpt2Bpe.mergeRank? (tok : Tokenizer) (a b : String) : Option ℕ

Look up the merge rank for an adjacent pair of BPE symbols.

def NN.API.text.Gpt2Bpe.bestMerge? (tok : Tokenizer) : List String → Option (String × String × ℕ)

Find the lowest-ranked merge currently available in a symbol list.

def NN.API.text.Gpt2Bpe.applyMerge (target : String × String) : List String → List String

Apply one BPE merge everywhere it appears in the current symbol list.

def NN.API.text.Gpt2Bpe.bpeLoop (tok : Tokenizer) : ℕ → List String → List String

Fuel-bounded BPE merge loop for a single escaped pre-token fragment.

def NN.API.text.Gpt2Bpe.encodeFragment (tok : Tokenizer) (fragment : String) : Except String (List ℕ)

Apply BPE to one pre-tokenized fragment.

def NN.API.text.Gpt2Bpe.encode (tok : Tokenizer) (text : String) : Except String (List ℕ)

Encode text using the loaded GPT-2 BPE files.

def NN.API.text.Gpt2Bpe.decode? (tok : Tokenizer) (ids : List ℕ) : Except String String

Decode GPT-2 BPE ids back to text.

def NN.API.text.Gpt2Bpe.decodeD (tok : Tokenizer) (ids : List ℕ) : String

Total display-oriented decoder: invalid ids/UTF-8 decode to an empty string.

def NN.API.text.Gpt2Bpe.asTextTokenizer (tok : Tokenizer) : text.Tokenizer

Adapt a loaded GPT-2 BPE tokenizer to the generic text-tokenizer interface.

File Loading

def NN.API.text.Gpt2Bpe.parseVocab (j : Lean.Json) : Except String (Array VocabEntry)

Parse GPT-2 vocab.json as an array of (token, id) entries.

The standard GPT-2 vocab.json is a single flat JSON object from token strings to numeric ids. Using Lean's fully general JSON object parser is convenient but slow for interactive examples because it builds a 50k-entry tree before we immediately flatten it again. The small parser below recognizes exactly the JSON shape used by GPT-2 vocab files and decodes JSON string escapes, including \uXXXX escapes for byte-to-unicode code points.

def NN.API.text.Gpt2Bpe.charAtD (cs : Array Char) (i : ℕ) : Char

Safe character lookup used by the specialized vocab.json parser.

def NN.API.text.Gpt2Bpe.skipJsonWs (cs : Array Char) (i : ℕ) : ℕ

Skip JSON whitespace in the specialized GPT-2 vocabulary parser.

def NN.API.text.Gpt2Bpe.hexVal? (c : Char) : Option ℕ

Interpret one hexadecimal digit from a JSON unicode escape.

def NN.API.text.Gpt2Bpe.parseHex4? (cs : Array Char) (i : ℕ) : Option ℕ

Parse four hexadecimal digits starting at i.

def NN.API.text.Gpt2Bpe.combineSurrogate (hi lo : ℕ) : ℕ

Combine a JSON UTF-16 surrogate pair into one Unicode code point.

def NN.API.text.Gpt2Bpe.parseJsonStringAux (cs : Array Char) : ℕ → ℕ → List Char → Except String (String × ℕ)

Fuel-bounded worker for JSON string parsing with escape handling.

def NN.API.text.Gpt2Bpe.parseJsonStringAt (cs : Array Char) (i : ℕ) : Except String (String × ℕ)

Parse a JSON string beginning at index i.

def NN.API.text.Gpt2Bpe.parseNatAt (cs : Array Char) (i : ℕ) : Except String (ℕ × ℕ)

Parse a natural-number literal beginning at index i.

def NN.API.text.Gpt2Bpe.parseVocabTextLoop (cs : Array Char) : ℕ → ℕ → Array VocabEntry → Except String (Array VocabEntry)

Fuel-bounded loop for the specialized GPT-2 vocab.json object parser.

def NN.API.text.Gpt2Bpe.parseVocabText (s : String) : Except String (Array VocabEntry)

Parse GPT-2 vocab.json directly from text.

def NN.API.text.Gpt2Bpe.vocabMapOf (vocab : Array VocabEntry) : Std.HashMap String ℕ

Build the token-to-id lookup table stored in a loaded GPT-2 BPE tokenizer.

def NN.API.text.Gpt2Bpe.idMapOf (vocab : Array VocabEntry) : Std.HashMap ℕ String

Build the id-to-token lookup table stored in a loaded GPT-2 BPE tokenizer.

def NN.API.text.Gpt2Bpe.mergeMapOf (merges : Array MergeRank) : Std.HashMap (String × String) ℕ

Build the pair-to-rank lookup table stored in a loaded GPT-2 BPE tokenizer.

def NN.API.text.Gpt2Bpe.mkTokenizer (vocab : Array VocabEntry) (merges : Array MergeRank) : Tokenizer

Assemble a tokenizer and its lookup maps from parsed GPT-2 vocabulary and merge tables.

def NN.API.text.Gpt2Bpe.parseMergeLine? (rank : ℕ) (line : String) : Option MergeRank

Parse one non-comment merges.txt line with its rank.

def NN.API.text.Gpt2Bpe.parseMerges (s : String) : Array MergeRank

Parse GPT-2 merges.txt. Invalid non-comment lines are ignored conservatively.

def NN.API.text.Gpt2Bpe.load (vocabJson mergesTxt : System.FilePath) : IO Tokenizer

Load GPT-2 BPE files directly in Lean.

def NN.API.text.Gpt2Bpe.loadWithProgress (tag : String) (vocabJson mergesTxt : System.FilePath) : IO Tokenizer

Load GPT-2 BPE files while printing progress for larger vocab.json / merges.txt assets.

Compact vocabulary projections

structure NN.API.text.LocalBpeVocab : Type

Projection from original tokenizer ids to a compact working vocabulary.

Small runnable language-model examples can keep a real tokenizer while restricting the model head to ids observed in the local corpus or prompt.

• originals : Array ℕ

  Original tokenizer id for each local id.

• toLocalMap : Std.HashMap ℕ ℕ

  Reverse lookup from original tokenizer id to local id.

def NN.API.text.LocalBpeVocab.size (lv : LocalBpeVocab) : ℕ

Number of live entries in a compact BPE projection.

def NN.API.text.LocalBpeVocab.toLocal (lv : LocalBpeVocab) (id : ℕ) : ℕ

Map an original tokenizer id into the compact local vocabulary, using local id 0 as OOV.

def NN.API.text.LocalBpeVocab.toOriginal (lv : LocalBpeVocab) (localId : ℕ) : ℕ

Map a compact local id back to its original tokenizer id.

def NN.API.text.buildLocalBpeVocab (maxVocab : ℕ) (corpusIds promptIds : Array ℕ) : LocalBpeVocab

Build a compact working vocabulary from corpus ids and prompt ids.

def NN.API.text.localizeBpeTokens (lv : LocalBpeVocab) (tokens : Array ℕ) : Array ℕ

Apply a compact BPE projection to an array of original tokenizer ids.