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 TorchLean.text.GPT2BPE.Internal.VocabularyEntry : 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 TorchLean.text.GPT2BPE.Internal.instReprVocabularyEntry.repr : VocabularyEntry → ℕ → Std.Format

@[instance_reducible]

instance TorchLean.text.GPT2BPE.Internal.instReprVocabularyEntry : Repr VocabularyEntry

def TorchLean.text.GPT2BPE.Internal.instDecidableEqVocabularyEntry.decEq (x✝ x✝¹ : VocabularyEntry) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance TorchLean.text.GPT2BPE.Internal.instDecidableEqVocabularyEntry : DecidableEq VocabularyEntry

structure TorchLean.text.GPT2BPE.Internal.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 TorchLean.text.GPT2BPE.Internal.instReprMergeRank.repr : MergeRank → ℕ → Std.Format

@[instance_reducible]

instance TorchLean.text.GPT2BPE.Internal.instReprMergeRank : Repr MergeRank

def TorchLean.text.GPT2BPE.Internal.instDecidableEqMergeRank.decEq (x✝ x✝¹ : MergeRank) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance TorchLean.text.GPT2BPE.Internal.instDecidableEqMergeRank : DecidableEq MergeRank

structure TorchLean.text.GPT2BPE.Internal.Tokenizer : Type

Loaded GPT-2 BPE tokenizer.

• vocabulary : Array VocabularyEntry

  Token vocabulary as loaded from vocab.json.

• merges : Array MergeRank

  Ranked merge table from merges.txt.

• tokenIds : Std.HashMap String ℕ

  Fast token-to-id lookup derived from vocabulary.

• tokensById : Std.HashMap ℕ String

  Fast id-to-token lookup derived from vocabulary.

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

  Fast pair-to-rank lookup derived from merges.

structure TorchLean.text.GPT2BPE.Internal.MergeCandidate : Type

One merge that is currently applicable to a BPE symbol sequence.

• left : String
• right : String
• rank : ℕ

Byte Escaping

def TorchLean.text.GPT2BPE.Internal.bytesVisible : Array ℕ

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

def TorchLean.text.GPT2BPE.Internal.bytesLatin1A : Array ℕ

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

def TorchLean.text.GPT2BPE.Internal.bytesLatin1B : Array ℕ

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

def TorchLean.text.GPT2BPE.Internal.baseBytes : Array ℕ

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

def TorchLean.text.GPT2BPE.Internal.containsNat (xs : Array ℕ) (x : ℕ) : Bool

Boolean membership test used while constructing the byte escape table.

def TorchLean.text.GPT2BPE.Internal.byteCodeTable : Array ℕ

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

def TorchLean.text.GPT2BPE.Internal.byteToChar (b : UInt8) : Char

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

def TorchLean.text.GPT2BPE.Internal.charByteTable : Std.HashMap Char UInt8

Shared inverse byte-escape lookup, constructed once for the fixed GPT-2 byte alphabet.

def TorchLean.text.GPT2BPE.Internal.charToByte? (c : Char) : Option UInt8

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

def TorchLean.text.GPT2BPE.Internal.byteEncode (s : String) : String

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

def TorchLean.text.GPT2BPE.Internal.byteDecode? (s : String) : Option String

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

Pre-tokenization

inductive TorchLean.text.GPT2BPE.Internal.RegexClass : Type

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

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

@[instance_reducible]

instance TorchLean.text.GPT2BPE.Internal.instDecidableEqRegexClass : DecidableEq RegexClass

def TorchLean.text.GPT2BPE.Internal.isRegexOther (c : Char) : Bool

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

def TorchLean.text.GPT2BPE.Internal.matchesRegexClass (cls : RegexClass) (c : Char) : Bool

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

def TorchLean.text.GPT2BPE.Internal.consumeContraction? : List Char → Option (String × List Char)

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

def TorchLean.text.GPT2BPE.Internal.consumeClassRun? (cls : RegexClass) : List Char → Option (String × List Char)

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

def TorchLean.text.GPT2BPE.Internal.consumeLookaheadWhitespace? (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 TorchLean.text.GPT2BPE.Internal.consumeWhitespaceRun? (xs : List Char) : Option (String × List Char)

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

def TorchLean.text.GPT2BPE.Internal.pretokenizeWithFuel : ℕ → 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 TorchLean.text.GPT2BPE.Internal.pretokenize (s : String) : List String

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

BPE Merging

def TorchLean.text.GPT2BPE.Internal.tokenId? (tokenizer : Tokenizer) (token : String) : Option ℕ

Look up a token id in a loaded tokenizer.

def TorchLean.text.GPT2BPE.Internal.tokenString? (tokenizer : Tokenizer) (id : ℕ) : Option String

Look up the token spelling for a token id.

def TorchLean.text.GPT2BPE.Internal.mergeRank? (tokenizer : Tokenizer) (left right : String) : Option ℕ

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

def TorchLean.text.GPT2BPE.Internal.bestMerge? (tokenizer : Tokenizer) : List String → Option MergeCandidate

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

def TorchLean.text.GPT2BPE.Internal.applyMerge (targetLeft targetRight : String) : List String → List String

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

def TorchLean.text.GPT2BPE.Internal.bpeLoop (tokenizer : Tokenizer) : ℕ → List String → List String

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

def TorchLean.text.GPT2BPE.Internal.encodeFragment (tokenizer : Tokenizer) (fragment : String) : Except String (Array ℕ)

Apply BPE to one pre-tokenized fragment.

def TorchLean.text.GPT2BPE.Internal.tokenIdsFromVocabulary (vocabulary : Array VocabularyEntry) : Std.HashMap String ℕ

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

def TorchLean.text.GPT2BPE.Internal.tokensByIdFromVocabulary (vocabulary : Array VocabularyEntry) : Std.HashMap ℕ String

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

def TorchLean.text.GPT2BPE.Internal.mergeRanksFromMerges (merges : Array MergeRank) : Std.HashMap (String × String) ℕ

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

def TorchLean.text.GPT2BPE.Internal.buildTokenizer (vocabulary : Array VocabularyEntry) (merges : Array MergeRank) : Tokenizer

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

structure TorchLean.text.GPT2BPE.Tokenizer : Type

Loaded GPT-2 BPE tokenizer.

The representation is intentionally hidden. In particular, callers cannot construct a tokenizer whose lookup tables disagree with its vocabulary or merge list; use load.

• representation : TorchLean.text.GPT2BPE.Internal.Tokenizer
• isCanonical : representation self = TorchLean.text.GPT2BPE.Internal.buildTokenizer (representation self).vocabulary (representation self).merges

opaque TorchLean.text.GPT2BPE.Tokenizer.Internal.create (representation : Internal.Tokenizer) (isCanonical : representation = Internal.buildTokenizer representation.vocabulary representation.merges) : Tokenizer

Wrap a canonical GPT-2 BPE representation at the public API boundary.

opaque TorchLean.text.GPT2BPE.Tokenizer.Internal.view (tokenizer : Tokenizer) : Internal.Tokenizer

Reveal the canonical tokenizer representation only to implementation code.

def TorchLean.text.GPT2BPE.Tokenizer.vocabularySize (tokenizer : Tokenizer) : ℕ

Number of tokens in a loaded GPT-2 vocabulary.

def TorchLean.text.GPT2BPE.Tokenizer.tokenId? (tokenizer : Tokenizer) (token : String) : Option ℕ

Look up a complete vocabulary token, including special tokens, without splitting it as text.

def TorchLean.text.GPT2BPE.encode (tokenizer : Tokenizer) (text : String) : Except String (Array ℕ)

Encode text using the loaded GPT-2 BPE files.

def TorchLean.text.GPT2BPE.decode (tokenizer : Tokenizer) (tokens : Array ℕ) : Except String String

Decode GPT-2 BPE ids back to text.

File Loading

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 TorchLean.text.GPT2BPE.Internal.charAtOrNull (cs : Array Char) (i : ℕ) : Char

Read a character for the specialized vocab.json parser, using NUL past the input boundary.

def TorchLean.text.GPT2BPE.Internal.skipJsonWs (cs : Array Char) (i : ℕ) : ℕ

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

def TorchLean.text.GPT2BPE.Internal.hexVal? (c : Char) : Option ℕ

Interpret one hexadecimal digit from a JSON unicode escape.

def TorchLean.text.GPT2BPE.Internal.parseHex4? (cs : Array Char) (i : ℕ) : Option ℕ

Parse four hexadecimal digits starting at i.

def TorchLean.text.GPT2BPE.Internal.combineSurrogate (hi lo : ℕ) : ℕ

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

def TorchLean.text.GPT2BPE.Internal.parseJsonStringWithFuel (cs : Array Char) : ℕ → ℕ → List Char → Except String (String × ℕ)

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

def TorchLean.text.GPT2BPE.Internal.parseJsonStringAt (cs : Array Char) (i : ℕ) : Except String (String × ℕ)

Parse a JSON string beginning at index i.

def TorchLean.text.GPT2BPE.Internal.parseNatAt (cs : Array Char) (i : ℕ) : Except String (ℕ × ℕ)

Parse a natural-number literal beginning at index i.

def TorchLean.text.GPT2BPE.Internal.finishVocabulary (cs : Array Char) (closing : ℕ) (entries : Array VocabularyEntry) : Except String (Array VocabularyEntry)

Finish the object only when its closing brace is followed by JSON whitespace.

def TorchLean.text.GPT2BPE.Internal.parseVocabularyTextLoop (cs : Array Char) : ℕ → ℕ → Array VocabularyEntry → Except String (Array VocabularyEntry)

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

def TorchLean.text.GPT2BPE.Internal.parseVocabularyText (s : String) : Except String (Array VocabularyEntry)

Parse GPT-2 vocab.json, requiring unique tokens and contiguous token ids from zero.

def TorchLean.text.GPT2BPE.Internal.parseMergeLine (rank : ℕ) (line : String) : Except String (Option MergeRank)

Parse one merges.txt line, skipping the version header and blank lines.

def TorchLean.text.GPT2BPE.Internal.parseMerges (s : String) : Except String (Array MergeRank)

Parse GPT-2 merges.txt, retaining hash-prefixed symbols and rejecting malformed pairs.

def TorchLean.text.GPT2BPE.load (vocabularyFile mergesFile : System.FilePath) (progress : Bool := false) (label : String := "GPT2BPE") : IO Tokenizer

Load GPT-2 BPE files directly in Lean. Vocabulary tokens and ids must be unique, with ids covering 0 .. vocabularySize - 1.

Set progress := true to print progress for larger vocab.json and merges.txt assets. The optional label prefixes those messages.