Tokenizers and Text Tensors

Text and NLP helpers for TorchLean examples.

Language models may keep token ids as Nat tensors and gather embedding rows directly. Small examples can instead use one-hot tensors of shape (batchSize × sequenceLength × vocabularySize). Both representations remain separate from floating-point model parameters at the API boundary.

This module provides:

• a tokenizer interface (with a byte-level tokenizer),
• helpers to turn token streams into one-hot tensors,
• “next-token prediction” sample builders used by GPT-style examples,
• display helpers for turning model logits back into readable token predictions.

Tokenizers

structure TorchLean.text.Tokenizer : Type

Tokenizer interface (encode/decode).

• vocabularySize : ℕ

  Vocabulary size (token ids are expected to be in [0, vocabularySize)).

• encode : String → Array ℕ

  Encode a string into a variable-length token buffer.

• decode : Array ℕ → String

  Decode a variable-length token buffer back into a string.

def TorchLean.text.Tokenizer.byte : Tokenizer

Byte-level UTF-8 tokenizer: each byte is one token in $[0,256)$.

Decoding uses UTF-8 when possible and falls back to byte-wise display for generated streams that are not valid UTF-8.

Example:

-- 256 tokens, no vocabulary file, nothing to train: this is where every text example starts.
def tokenizer : text.Tokenizer := text.Tokenizer.byte

-- Encode then decode returns the original string whenever it was valid UTF-8.
def roundTrip (line : String) : String :=
  tokenizer.decode (tokenizer.encode line)

def TorchLean.text.Tokenizer.fromAlphabet (alphabet : Array Char) (unknownTokenId : Fin alphabet.size) (unknownCharacter : Char := '?') : Tokenizer

Build a character-level tokenizer from an explicit alphabet.

The resulting encode/decode pair has the same role as the stoi/itos tables used in character-level GPT examples: encode maps characters to ids 0..alphabet.size-1, and decode maps ids back to characters.

The unknownTokenId argument proves that the alphabet is nonempty and identifies the token used for a character outside the alphabet. Ids outside [0, alphabet.size) decode to unknownCharacter. Repeated characters in the alphabet encode to their first index. A lookup table is shared across calls to the returned encoder.

Example:

-- The `stoi` and `itos` tables of character-level GPT tutorials, with the nonempty-alphabet
-- requirement carried by the unknown-token index instead of a runtime assertion.
def alphabet : Array Char := #['a', 'b', 'c', ' ']

def tokenizer : text.Tokenizer :=
  text.Tokenizer.fromAlphabet alphabet ⟨3, by decide⟩ (unknownCharacter := '?')

def TorchLean.text.Tokenizer.encodeFixed (tokenizer : Tokenizer) (sequenceLength : ℕ) (text : String) (paddingTokenId : ℕ := 0) : Tensor ℕ [sequenceLength]

Encode a string and pad or truncate it to exactly sequenceLength token ids.

Example:

-- Padded or truncated to the length the model expects, so the result carries a shape rather than
-- a length a caller has to check.
def tokens : Tensor Nat [16] :=
  text.Tokenizer.byte.encodeFixed 16 "hello world"

def TorchLean.text.Tokenizer.encodeFixedBatch {batchSize : ℕ} (tokenizer : Tokenizer) (sequenceLength : ℕ) (texts : Tensor String [batchSize]) (paddingTokenId : ℕ := 0) : Tensor ℕ [batchSize, sequenceLength]

Encode exactly batchSize strings, padding or truncating each row to sequenceLength token ids.

Example:

-- Two prompts become one `[2, 16]` batch, ready for a model with a batch axis in front.
def batch : Tensor Nat [2, 16] :=
  text.Tokenizer.byte.encodeFixedBatch 16 ["hello", "world"]

Byte-Corpus Windows

def TorchLean.text.Internal.byteAtOrPad (bytes : ByteArray) (index : ℕ) (paddingTokenId : ℕ := 0) : ℕ

Read one byte token from a raw corpus, returning paddingTokenId past the end.

This is byte-level rather than BPE-level: examples can train causal language models directly from a text file without depending on an external tokenizer artifact. GPT-2 BPE support lives in NN.API.Text.Bpe.

Lives in Internal on purpose: byteTokenWindow below is the only caller, and a padded single-byte read is not something a user of text should have to reason about. (private is not an option here. Every API module is inside @[expose] public section, so a private helper cannot be named from an exposed body; a nested Internal namespace is how the rest of the codebase says "plumbing".)

def TorchLean.text.byteTokenWindow (bytes : ByteArray) (length : ℕ) (offset paddingTokenId : ℕ := 0) : Tensor ℕ [length]

Extract a fixed-length byte-token window from a raw corpus.

offset is measured in bytes, as required for byte-level causal language modeling. This avoids hidden UTF-8 slicing assumptions.

Corpus Helpers

def TorchLean.text.Corpus.readUtf8File (exeName : String) (path : System.FilePath) (missingHint : String) : IO String

Read a UTF-8 text file with a caller-supplied preparation hint.

The examples pass their executable name and a concrete hint so failures point users to the exact download or conversion command for that dataset.

def TorchLean.text.Corpus.readByteFile (exeName : String) (path : System.FilePath) (allowSmallData : Bool) (minimumBytes sequenceLength : ℕ) : IO ByteArray

Read a raw byte corpus and optionally enforce a minimum size.

allowSmallData is an explicit override for bounded local runs. Corpus-training commands can set minimumBytes to the scale they expect and require users to acknowledge smaller local files.

def TorchLean.text.Corpus.takeUtf8Input (exeName : String) (defaultPath : System.FilePath) (aliases : List (String × System.FilePath)) (missingHint : String) (arguments : List String) : IO (String × List String)

Parse a text-corpus flag set and return (text, remainingArgs).

Supported forms:

• --data-file PATH
• any named alias in aliases, such as ("--tiny-shakespeare", path)
• no data flag, which uses defaultPath

def TorchLean.text.Corpus.usableTokenStarts (tokenCount sequenceLength : ℕ) : ℕ

Number of legal start positions for a (sequenceLength + 1) next-token window.

We return at least one start position so bounded corpora stay total; callers can still enforce a minimum corpus size before training.

def TorchLean.text.Corpus.byteOffset (bytes : ByteArray) (index sequenceLength : ℕ) : ℕ

Deterministic sliding-window offset for a byte corpus.

def TorchLean.text.Corpus.tokenOffset {tokenCount : ℕ} (_tokens : Tensor ℕ [tokenCount]) (index sequenceLength : ℕ) : ℕ

Deterministic sliding-window offset for an already-tokenized corpus.

def TorchLean.text.Corpus.evenlySpacedOffsets (tokenCount sequenceLength windowCount : ℕ) : Tensor ℕ [windowCount]

Choose deterministic, approximately evenly spaced starts for fixed-width token windows.

def TorchLean.text.Corpus.randomBatchOffsets (tokenCount sequenceLength batchSize seed step : ℕ) : Tensor ℕ [batchSize]

Deterministic minGPT-style random offsets for one training batch.

The result has one corpus start offset per batch row. We derive the random key from (seed, step) and then draw row offsets by the row index, so the run is reproducible without using ambient IO randomness. This is the text equivalent of a shuffled EpochLoader epoch.

def TorchLean.text.Corpus.randomTokenBatch {β : Type} [Storage β] {tokenCount : ℕ} (tokens : Tensor β [tokenCount]) (batchSize sequenceLength seed step : ℕ) (paddingTokenId : β) : Tensor β [batchSize, sequenceLength + 1]

Build token windows for one deterministic random text batch.

Each row gets sequenceLength + 1 ids so downstream causal-LM helpers can form both the input and shifted target. Byte, character, BPE, and synthetic tokenizers share the same tensor batching semantics.

def TorchLean.text.Corpus.promptAwareOffsets (tokenCount sequenceLength windowCount : ℕ) (promptOffset? : Option ℕ) : Tensor ℕ [windowCount]

Choose training-window offsets, biased toward a prompt occurrence when the corpus contains it.

If the prompt is present in the corpus, a portion of the sampled windows covers nearby text. That keeps generation reports tied to text the model actually saw during training.