Tokenizers and Text Tensors

Text and NLP helpers for TorchLean examples.

TorchLean’s executable runtime expects inputs as floating tensors, so runtime and autograd code can handle them with the same typed tensor APIs as parameters. For language models this means we commonly represent token ids as one-hot / token-distribution tensors of shape:

(batch × seqLen × vocab)

and implement “token embeddings” as a matrix multiply against an embedding table.

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).

• vocabSize : ℕ

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

• encode : String → List ℕ

  Encode a string into token ids.

• decode : List ℕ → String

  Decode token ids back into a string.

def TorchLean.text.Tokenizer.byteArrayOfIds (ids : List ℕ) : ByteArray

Convert token ids to bytes, truncating each id modulo 256.

def TorchLean.text.Tokenizer.decodeByteIds (ids : List ℕ) : String

Decode byte tokens as UTF-8 when possible, falling back to a byte-wise display mode for generated byte streams that are not valid UTF-8. For valid UTF-8 strings, $\operatorname{decode}(\operatorname{encode}(s))=s$; model output remains printable even when the byte stream is invalid UTF-8.

def TorchLean.text.Tokenizer.byte : Tokenizer

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

def TorchLean.text.Tokenizer.ofAlphabet (alphabet : Array Char) (unkId : ℕ := 0) (unkChar : 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.

Notes:

• This tokenizer is deterministic given alphabet; callers are responsible for choosing how to construct the alphabet (e.g. sorted(set(data))).
• Characters not present in the alphabet map to unkId (default 0), so encode is total.
• Ids outside [0, vocabSize) decode to the unkChar (default ?).

def TorchLean.text.Tokenizer.encodeVec (t : Tokenizer) (n : ℕ) (s : String) (padId : ℕ := 0) : Vector ℕ n

Encode and pad/truncate to a fixed length, returning a length-indexed Vector.

def TorchLean.text.Tokenizer.encodeBatchVec (t : Tokenizer) (batch seqLen : ℕ) (ss : List String) (padId : ℕ := 0) : Vector (Vector ℕ seqLen) batch

Encode a batch of strings, padding/truncating each to length seqLen.

One-Hot Token Tensors

def TorchLean.text.oneHotTokenFloat (vocab tokenId : ℕ) : Spec.Tensor Float (Spec.Shape.dim vocab Shape.scalar)

One-hot vector for a single token id (Vec vocab). Out-of-range ids map to all-zeros.

def TorchLean.text.tokensToOneHotMatFloat {seqLen vocab : ℕ} (tokens : Vector ℕ seqLen) : Spec.Tensor Float (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))

One-hot encode a fixed-length token sequence as a matrix (seqLen × vocab).

def TorchLean.text.tokensToOneHotBatchFloat {batch seqLen vocab : ℕ} (tokens : Vector (Vector ℕ seqLen) batch) : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))

One-hot encode a fixed-size batch of token sequences as (batch × seqLen × vocab).

Causal LM Samples

def TorchLean.text.causalLmXYOneHotMatFloat (seqLen vocab : ℕ) (tokens : List ℕ) (padId : ℕ := 0) : Spec.Tensor Float (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)) × Spec.Tensor Float (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))

Build a (x, y) pair for next-token prediction from a token stream.

$$ x[t] = \operatorname{oneHot}(\mathrm{tokens}[t]), \qquad y[t] = \operatorname{oneHot}(\mathrm{tokens}[t+1]). $$

If the stream is too short, we pad with padId.

def TorchLean.text.causalLmXYOneHotBatchRowsFloat (batch seqLen vocab : ℕ) (tokensAt : Fin batch → List ℕ) (padId : ℕ := 0) : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))) × Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))

Build a batched causal-LM (x, y) pair from one token window per batch row.

This is the text analogue of image/tabular minibatching:

• row $i$ receives its own token window tokensAt i;
• $x[i,t]$ is $\mathrm{tokensAt}(i)[t]$;
• $y[i,t]$ is $\mathrm{tokensAt}(i)[t+1]$;
• short rows are padded with padId.

GPT-style examples share this batching logic. The contract is explicit: a text batch is a typed tensor of shape (batch, seqLen, vocab), just like the vision loader collates rows into (batch, C, H, W).

def TorchLean.text.causalLmXOneHotBatch {α : Type} [Context α] [Runtime.FromFloat α] (batch seqLen vocab : ℕ) (tokens : List ℕ) (padId : ℕ := 0) : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))

One-hot encode a causal-LM input window as a batched tensor.

Token ids are read from tokens, missing positions use padId, and every batch row receives the same window. Use causalLmSampleOneHotBatchRows when rows should come from different corpus offsets.

def TorchLean.text.causalLmXOneHotBatchRows {α : Type} [Context α] [Runtime.FromFloat α] (batch seqLen vocab : ℕ) (tokensAt : Fin batch → List ℕ) (padId : ℕ := 0) : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))

One-hot encode one causal-LM input window per batch row.

This is the input-only companion to causalLmSampleOneHotBatchRows, used by generation code that has prefixes but no shifted training targets.

def TorchLean.text.causalLmSampleOneHotBatch {α : Type} [Context α] [Runtime.FromFloat α] (batch seqLen vocab : ℕ) (tokens : List ℕ) (padId : ℕ := 0) : Sample.Supervised α (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))) (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))

Build a batched supervised next-token sample from a token stream.

The target is shifted by one position: $x[t]=\mathrm{tokens}[t]$ and $y[t]=\mathrm{tokens}[t+1]$. Every batch row receives the same window, which is useful for prompt evaluation, deterministic checks, and synthetic sequence tasks.

def TorchLean.text.causalLmSampleOneHotBatchRows {α : Type} [Context α] [Runtime.FromFloat α] (batch seqLen vocab : ℕ) (tokensAt : Fin batch → List ℕ) (padId : ℕ := 0) : Sample.Supervised α (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))) (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))

Build a batched supervised causal-LM sample from one token window per batch row.

Use this for GPT-style minibatches with distinct corpus windows. causalLmSampleOneHotBatch remains useful when every batch row should repeat a fixed prompt or synthetic sequence.

Byte-Corpus Windows

def TorchLean.text.byteAtD (bytes : ByteArray) (i : ℕ) (padId : ℕ := 0) : ℕ

Read one byte token from a raw corpus, returning padId 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.

def TorchLean.text.byteTokenWindow (bytes : ByteArray) (n : ℕ) (offset padId : ℕ := 0) : List ℕ

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) (minBytes seqLen : ℕ) : 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 minBytes to the scale they expect and require users to acknowledge smaller local files.

partial def TorchLean.text.Corpus.takeUtf8Input (exeName : String) (defaultPath : System.FilePath) (aliases : List (String × System.FilePath)) (missingHint : String) : 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.byteOffset (bytes : ByteArray) (i seqLen : ℕ) : ℕ

Deterministic sliding-window offset for a byte corpus.

def TorchLean.text.Corpus.tokenOffset (tokens : Array ℕ) (i seqLen : ℕ) : ℕ

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

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

Number of legal start positions for a (seqLen + 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.tokenArrayWindow (tokens : Array ℕ) (n offset : ℕ) (padId : ℕ := 0) : List ℕ

Extract a fixed token window from an array-backed token corpus.

def TorchLean.text.Corpus.randomBatchOffsets (tokenCount seqLen batch seed step : ℕ) : Fin batch → ℕ

Deterministic minGPT-style random offsets for one training batch.

The result is a function Fin batch → Nat: one corpus start offset per 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 DataLoader epoch.

def TorchLean.text.Corpus.randomBatchTokenWindows (tokens : Array ℕ) (batch seqLen seed step : ℕ) (padId : ℕ := 0) : Fin batch → List ℕ

Build token windows for one deterministic random text batch.

Each row gets seqLen + 1 ids so downstream causal-LM helpers can form both x and shifted y. The helper is token-array based, so byte, character, BPE, and synthetic tokenizers can all produce an Array Nat and reuse the same batching semantics.

def TorchLean.text.Corpus.startsWithAt (xs pat : Array ℕ) (off : ℕ) : Bool

Check whether pat occurs in xs at offset off.

def TorchLean.text.Corpus.findWindow? (xs pat : Array ℕ) : Option ℕ

Find the first offset where pat appears in xs.

def TorchLean.text.Corpus.promptAwareOffsets (tokenCount seqLen windows : ℕ) (promptOffset? : Option ℕ) : List ℕ

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.