Text Generation

Score filtering, top-k sampling, logit extraction, decoding, and causal masks used by language-model examples.

structure TorchLean.text.Internal.ScoredToken (vocabularySize : ℕ) : Type

A valid token id together with its unscaled logit.

• token : Fin vocabularySize

  Token id in the model vocabulary.

• score : Float

  Logit, before temperature scaling. Selection excludes NaNs.

def TorchLean.text.Internal.worseOrEqual {vocabularySize : ℕ} (left right : ScoredToken vocabularySize) : Bool

Lower scores rank worse; equal scores prefer the smaller token id.

def TorchLean.text.Internal.scoredTokenAt? {vocabularySize : ℕ} (scores : Tensor Float [vocabularySize]) (allowToken : Fin vocabularySize → Bool) (index : ℕ) : Option (ScoredToken vocabularySize)

Read one allowed, non-NaN logit without building an index tensor.

def TorchLean.text.Internal.topCandidates {vocabularySize : ℕ} (scores : Tensor Float [vocabularySize]) (k : ℕ) (allowToken : Fin vocabularySize → Bool) : Array (ScoredToken vocabularySize)

Keep only the best k logits in a bounded heap, then return them in descending order.

def TorchLean.text.Internal.bestCandidate? {vocabularySize : ℕ} (count : ℕ) (candidateAt? : ℕ → Option (ScoredToken vocabularySize)) : Option (ScoredToken vocabularySize)

Find the best indexed candidate in one pass, breaking ties by the smaller token id.

def TorchLean.text.topKTokens {vocabularySize : ℕ} (scores : Tensor Float [vocabularySize]) (k : ℕ) (allowToken : Fin vocabularySize → Bool := fun (x : Fin vocabularySize) => true) : Tensor (Option (Fin vocabularySize)) [min k vocabularySize]

Return up to k allowed, non-NaN score indices, largest first.

The result has min k vocabularySize slots; unused trailing slots contain none. Equal scores, including signed zeros and infinities, prefer smaller token ids. A bounded heap uses O(vocabularySize * log(k + 1)) time and O(min k vocabularySize) auxiliary storage.

The optional predicate filters token ids without replacing their scores by a finite sentinel. This matters when model logits are unbounded: a disallowed token must never become selectable merely because every allowed score is smaller than an arbitrary masking constant.

Example:

-- Highest scores first. `allowToken` filters rather than masking, so a banned token cannot become
-- selectable just because the masking constant happened to sit above every real score.
def best (scores : Tensor Float [256]) : Tensor (Option (Fin 256)) [8] :=
  text.topKTokens scores 8 (allowToken := fun token => token.val < 128)

def TorchLean.text.greedyToken? {vocabularySize : ℕ} (scores : Tensor Float [vocabularySize]) (allowToken : Fin vocabularySize → Bool := fun (x : Fin vocabularySize) => true) : Option (Fin vocabularySize)

Greedy argmax, or none when no allowed non-NaN token exists. Equal scores prefer smaller token ids. This scans the vocabulary once with constant auxiliary storage.

Example:

-- `none` means no allowed non-NaN score exists, which is a failure worth seeing rather than a
-- silent fall back to token zero.
def next (scores : Tensor Float [256]) : Option (Fin 256) :=
  text.greedyToken? scores

def TorchLean.text.Internal.penalizeRepeats {vocabularySize recentCount : ℕ} (scores : Tensor Float [vocabularySize]) (recentTokens : Tensor ℕ [recentCount]) (repeatPenalty : Float) : Tensor Float [vocabularySize]

Apply a repetition penalty by subtracting $\mathrm{repeatPenalty}\,\mathrm{count}(\mathrm{token})$ for tokens appearing in recent.

This is a local sampling heuristic; it is not the same as the presence or frequency penalties used by hosted APIs, but it gives examples a deterministic way to discourage immediate repetition.

Internal on purpose: chooseNextToken applies it for you, and calling it out of order (after the softmax rather than on the logits) would silently change the sampling distribution.

def TorchLean.text.isPrintableAscii (token : ℕ) : Bool

True for byte tokens that a terminal can print: the printable ASCII range plus newline.

Named with the is prefix that Lean core uses for Char.isAlpha and friends, so that reading if isPrintableAscii token at a call site tells you a Bool comes back.

def TorchLean.text.Internal.escapeByteToken (token : ℕ) : String

Escape one byte token for display inside a quoted string. Used only by formatByteTokens.

def TorchLean.text.formatByteTokens {n : ℕ} (tokens : Tensor ℕ [n]) : String

Escape byte ids as a one-line quoted display string.

Example:

-- Generated bytes are not always valid UTF-8, so display escapes them instead of guessing:
-- `#[104, 105, 10]` prints as `"hi\n"`.
def display {n : Nat} (tokens : Tensor Nat [n]) : String := text.formatByteTokens tokens

def TorchLean.text.Internal.samplingWeight (score maximum temperature : Float) : Float

Stable softmax weight relative to a finite maximum and a positive finite temperature.

def TorchLean.text.Internal.sampleCandidates? {vocabularySize : ℕ} (count : ℕ) (candidateAt? : ℕ → Option (ScoredToken vocabularySize)) (temperature : Float) (seed counter : ℕ) : Option (Fin vocabularySize)

Sample an indexed candidate sequence, storing weights once and never sorting the sequence.

def TorchLean.text.sampleTopKToken? {vocabularySize : ℕ} (scores : Tensor Float [vocabularySize]) (temperature : Float) (topK seed counter : ℕ) (allowToken : Fin vocabularySize → Bool := fun (x : Fin vocabularySize) => true) : Option (Fin vocabularySize)

Sample an allowed token id using a positive finite temperature and top-k sampling.

NaNs and disallowed tokens are excluded. Ties at the cutoff prefer smaller token ids. If the maximum is infinite, return its smallest token id directly (also when every allowed score is negative infinity), avoiding undefined softmax normalization.

topK = 0 or topK ≥ vocabularySize samples the full vocabulary in token-id order in linear time. A smaller positive cutoff samples the bounded heap's output in descending score order. Randomness is deterministic given the scores, filter, temperature, cutoff, seed, and counter; the two traversal orders can produce different tokens for the same random draw.

def TorchLean.text.chooseNextToken {vocabularySize recentCount : ℕ} (scores : Tensor Float [vocabularySize]) (options : GenerationOptions) (counter : ℕ) (recentTokens : Tensor ℕ [recentCount]) (allowToken : Fin vocabularySize → Bool := fun (x : Fin vocabularySize) => true) : Except String (Fin vocabularySize)

Select the next token, rejecting an empty allow-list, a non-finite or negative repetition penalty, or an invalid sampling temperature. Greedy decoding (topK = 1) ignores temperature.

Example:

-- One decoding policy shared by every text example: repeat penalty first, then greedy or top-k
-- sampling. `counter` keeps each step's randomness distinct while staying reproducible.
def next (scores : Tensor Float [256])
    (options : text.GenerationOptions) (step : Nat) :
    Except String (Fin 256) :=
  text.chooseNextToken scores options (counter := step) (recentTokens := Tensor.full [0] 0)

def TorchLean.text.autoregressiveTokenIds {vocabularySize promptLength : ℕ} (sequenceLength paddingTokenId : ℕ) (promptTokens : Tensor ℕ [promptLength]) (options : GenerationOptions) (scoreWindow : Tensor ℕ [sequenceLength] → Fin sequenceLength → IO (Tensor Float [vocabularySize])) (allowToken : Fin vocabularySize → Bool := fun (x : Fin vocabularySize) => true) : IO (Tensor ℕ [promptLength + options.newTokenCount])

Autoregressively extend token ids with a model-provided score callback.

The callback receives an exact-length context window and the sequence position whose logits should be used for the next token. The shared policy crops to the last sequenceLength tokens, pads, applies repeat penalties, samples by top-k/temperature, and writes one token per step. The output includes the prompt followed by exactly newTokenCount tokens. A zero-length model context is rejected when generation is requested.

Example:

-- The model arrives as a callback: given an exact-length window and the position whose logits to
-- read, return that position's scores. Cropping, padding, penalties, and sampling stay here.
def generate (options : text.GenerationOptions)
    (scoreWindow : Tensor Nat [16] → Fin 16 → IO (Tensor Float [256])) :
    IO (Tensor Nat [(text.Tokenizer.byte.encode options.prompt).size + options.newTokenCount]) :=
  text.autoregressiveTokenIds 16 (paddingTokenId := 0)
    (promptTokens := Tensor.from (text.Tokenizer.byte.encode options.prompt))
    (options := options) (scoreWindow := scoreWindow)

The next six declarations come in three pairs: an operation on (sequenceLength × vocabularySize) logits, and the same operation on a batch, which takes an extra batchIndex and works on one row. The batched member of a pair is the unbatched name with a batch prefix, always in that position, so knowing one spelling gives you the other.

def TorchLean.text.logitScoresAt {α : Type} [Storage α] {sequenceLength vocabularySize : ℕ} (logits : Tensor α [sequenceLength, vocabularySize]) (position : Fin sequenceLength) : Tensor α [vocabularySize]

Extract the vocabulary-score row at one statically valid sequence position.

def TorchLean.text.batchLogitScoresAt {α : Type} [Storage α] {batchSize sequenceLength vocabularySize : ℕ} (logits : Tensor α [batchSize, sequenceLength, vocabularySize]) (batchIndex : Fin batchSize) (position : Fin sequenceLength) : Tensor α [vocabularySize]

Extract a vocabulary-score row from batched logits.

def TorchLean.text.argmaxTokens {α : Type} [Storage α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {sequenceLength vocabularySize : ℕ} (logits : Tensor α [sequenceLength, vocabularySize]) : Tensor ℕ [sequenceLength]

Decode a matrix of token logits by taking argmax independently at each sequence position.

The shape is (sequenceLength × vocabularySize), i.e. one logits vector per token position. This helper is for inspection/debugging and is not differentiable.

def TorchLean.text.decodeArgmaxLogits {α : Type} [Storage α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (tokenizer : Tokenizer) {sequenceLength vocabularySize : ℕ} (logits : Tensor α [sequenceLength, vocabularySize]) : String

Decode (sequenceLength × vocabularySize) logits as text using a tokenizer.

def TorchLean.text.batchArgmaxTokens {α : Type} [Storage α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {batchSize sequenceLength vocabularySize : ℕ} (logits : Tensor α [batchSize, sequenceLength, vocabularySize]) (batchIndex : Fin batchSize) : Tensor ℕ [sequenceLength]

Extract batchIndex from batched logits and return the per-position argmax token ids.