GPT-2-Style Model Helpers (API)

This module collects compact, reusable GPT-2-style building blocks for TorchLean examples:

• a single “causal LM over one-hot tokens” model constructor, and
• a small configuration record that keeps the hyperparameter inventory explicit.

These helpers live in the API layer so runnable examples can stay focused on: data prep, training loops, and text decoding, rather than repeating the same embedding → positional embedding → Transformer stack → LayerNorm → linear boilerplate.

Important scope note:

• This is not a pretrained checkpoint loader.
• These are compact example architectures shaped like GPT-2 blocks.
• Tokenizers live under NN.API.text / NN.API.text.Gpt2Bpe.

structure NN.API.nn.models.CausalOneHotConfig : Type

Configuration for a small GPT-2-style causal language model over one-hot token inputs.

The model has the common GPT-2 “shape”:

embedding → learned positional embedding → (masked self-attention + FFN)×layers → LayerNorm → linear

The input and output shapes are (batch × seqLen × vocab) one-hot/logit tensors.

• batch : ℕ
• seqLen : ℕ
• vocab : ℕ
• numHeads : ℕ
• headDim : ℕ
• ffnHidden : ℕ
• layers : ℕ
• seedStride : ℕ

  Seed stride used when initializing repeated blocks.

@[implicit_reducible]

instance NN.API.nn.models.instReprCausalOneHotConfig : Repr CausalOneHotConfig

def NN.API.nn.models.instReprCausalOneHotConfig.repr : CausalOneHotConfig → ℕ → Std.Format

def NN.API.nn.models.CausalOneHotConfig.dModel (cfg : CausalOneHotConfig) : ℕ

Transformer width implied by numHeads * headDim.

@[reducible, inline]

abbrev NN.API.nn.models.causalOneHotShape (cfg : CausalOneHotConfig) : Shape

Input/output tensor shape (batch × seqLen × vocab) for a one-hot causal LM.

@[reducible, inline]

abbrev NN.API.nn.models.causalEmbeddingShape (cfg : CausalOneHotConfig) : Shape

Embedded-token tensor shape (batch × seqLen × dModel).

def NN.API.nn.models.causalTransformerFromEmbeddings (cfg : CausalOneHotConfig) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (causalEmbeddingShape cfg) (causalOneHotShape cfg))

GPT-2-style causal Transformer body after token embeddings have already been computed.

This is the shared body used by both one-hot-token models and indexed-token experiments. Keeping it separate avoids duplicating the Transformer stack when callers use a different token representation: the input boundary changes, while positional embeddings, masked self-attention blocks, layer norm, and the language-model head stay the same.

def NN.API.nn.models.causalTransformerOneHot (cfg : CausalOneHotConfig) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (causalOneHotShape cfg) (causalOneHotShape cfg))

Build a GPT-2-style causal language model over one-hot tokens.

This is the shared constructor used by the runnable GPT-2 examples. It stays in nn.M so it composes with the rest of the API-layer model-building interface.

def NN.API.nn.models.causalTransformerTokenScalarModuleDefWithMode (mode : TorchLean.NN.Mode) (cfg : CausalOneHotConfig) (body : Sequential (causalEmbeddingShape cfg) (causalOneHotShape cfg)) (tokens targets : Spec.Tensor ℕ (Spec.Shape.dim (cfg.batch * cfg.seqLen) Spec.Shape.scalar)) (reduction : TorchLean.Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : Runtime.Autograd.TorchLean.ScalarModuleDef (Spec.Shape.dim cfg.vocab (Spec.Shape.dim cfg.dModel Spec.Shape.scalar) :: paramShapes body) []

Scalar loss for causal language modeling with integer token ids.

The public one-hot constructor above is useful for small teaching examples because the input is an ordinary Float tensor. File-backed tokenized datasets use the representation found in language-model training systems: token ids are Nats, the embedding table is a trainable Float parameter, and the loss gathers the target classes directly instead of building one-hot targets.

tokens and targets are flattened (batch * seqLen) vectors. This matches the backend gather ops and keeps dataset storage simple; the embedding helper reshapes gathered rows back to (batch, seqLen, dModel) before running the Transformer body.

def NN.API.nn.models.causalTransformerTokenScalarModuleDef (cfg : CausalOneHotConfig) (body : Sequential (causalEmbeddingShape cfg) (causalOneHotShape cfg)) (tokens targets : Spec.Tensor ℕ (Spec.Shape.dim (cfg.batch * cfg.seqLen) Spec.Shape.scalar)) (reduction : TorchLean.Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : Runtime.Autograd.TorchLean.ScalarModuleDef (Spec.Shape.dim cfg.vocab (Spec.Shape.dim cfg.dModel Spec.Shape.scalar) :: paramShapes body) []

Training-mode wrapper for integer-token causal language modeling.