Text Generation

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

def TorchLean.text.topKIndices (scores : Array Float) (k : ℕ) : List ℕ

Return the indices of the top k scores (largest first).

This deterministic utility is used by the GPT-style examples. The direct O(k*vocab) implementation is adequate for the vocabulary sizes and top-k values used by these executable examples.

def TorchLean.text.greedyIndex (scores : Array Float) : ℕ

Greedy argmax index.

def TorchLean.text.penalizeRepeats (scores : Array Float) (recent : List ℕ) (repeatPenalty : Float) : Array Float

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.

def TorchLean.text.restrictScores (scores : Array Float) (allowId : ℕ → Bool) : Array Float

Mask scores by an allow-list predicate (disallowed ids get a very negative score).

This is mainly used by byte-level examples to optionally restrict output to printable ASCII.

def TorchLean.text.prepareScoresForGeneration (scores : Array Float) (recent : List ℕ) (repeatPenalty : Float) (allowId : ℕ → Bool := fun (x : ℕ) => true) : Array Float

Apply repeat penalty and an allow-list mask before sampling.

def TorchLean.text.printableAsciiByte (i : ℕ) : Bool

Printable ASCII bytes plus newline.

def TorchLean.text.escapeByteId (b : ℕ) : String

Escape one byte token for display inside a quoted string.

def TorchLean.text.escapeByteIdsForDisplay (ids : List ℕ) : String

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

def TorchLean.text.sampleTopKIndex (scores : Array Float) (temperature : Float) (topK seed counter : ℕ) : ℕ

Sample one token id from scores using temperature + top-k sampling.

The randomness is deterministic given (seed, counter), so a run with the same flags produces the same sampled text.

def TorchLean.text.chooseNextToken (scores : Array Float) (opts : GenerationOptions) (counter : ℕ) (recent : List ℕ := []) (allowId : ℕ → Bool := fun (x : ℕ) => true) : ℕ

Select the next token from prepared logits using greedy or temperature/top-k sampling.

def TorchLean.text.autoregressiveTokenIds (seqLen padId : ℕ) (promptIds : List ℕ) (opts : GenerationOptions) (scoreWindow : List ℕ → ℕ → IO (Array Float)) (allowId : ℕ → Bool := fun (x : ℕ) => true) (sanitize : ℕ → ℕ := fun (tok : ℕ) => tok) : IO (List ℕ)

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

The callback receives the padded context window and the sequence position whose logits should be used for the next token. The shared policy crops to the last seqLen tokens, pads, applies repeat penalties, samples by top-k/temperature, and appends one token per step.

partial def TorchLean.text.autoregressiveTokenIds.loop (seqLen padId : ℕ) (opts : GenerationOptions) (scoreWindow : List ℕ → ℕ → IO (Array Float)) (allowId : ℕ → Bool) (sanitize : ℕ → ℕ) (ids : List ℕ) : ℕ → IO (List ℕ)

def TorchLean.text.logitScoresAt {seqLen vocab : ℕ} (logits : Spec.Tensor Float (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))) (pos : ℕ) : Array Float

Extract the vocabulary-score row at one sequence position.

def TorchLean.text.batchLogitScoresAt {batch seqLen vocab : ℕ} (logits : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))) (batchIdx : Fin batch) (pos : ℕ) : Array Float

Extract a vocabulary-score row from batched logits.

def TorchLean.text.argmaxTokenIdsFromLogits {α : Type} [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen vocab : ℕ} (logits : Spec.Tensor α (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))) : List ℕ

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

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

def TorchLean.text.decodeArgmaxLogits {α : Type} [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (t : Tokenizer) {seqLen vocab : ℕ} (logits : Spec.Tensor α (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar))) : String

Decode (seqLen × vocab) logits as text using a tokenizer.

def TorchLean.text.argmaxTokenIdsFromBatchLogits {α : Type} [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {batch seqLen vocab : ℕ} (logits : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))) (batchIdx : Fin batch) : List ℕ

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

def TorchLean.text.decodeArgmaxBatchLogits {α : Type} [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (t : Tokenizer) {batch seqLen vocab : ℕ} (logits : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim vocab Shape.scalar)))) (batchIdx : Fin batch) : String

Decode one batch row of (batch × seqLen × vocab) logits as text.

def TorchLean.text.causalMask (seqLen : ℕ) : Spec.Tensor Bool (Spec.Shape.dim seqLen (Spec.Shape.dim seqLen Shape.scalar))

Causal (autoregressive) attention mask of shape (seqLen × seqLen).

Entry $(i,j)$ is true iff $j \leq i$, meaning position $i$ may attend to itself and earlier positions but not to future positions.