Mamba Text Training

Runnable byte-level language-model training with the public Mamba API constructor.

The model is trainable end-to-end:

mamba(contextLength, vocabularySize, stateWidth) → linear(stateWidth → vocabularySize)

and the same code runs on CPU or CUDA through TorchLean autograd.

python3 scripts/datasets/download_example_data.py --tiny-shakespeare
lake -R -K cuda=true exe torchlean mamba --device cuda --tiny-shakespeare --steps 1 --windows 1 \
  --generate 0

def NN.Examples.Models.Sequence.Mamba.exeName : String

CLI subcommand name used in terminal banners and error messages.

def NN.Examples.Models.Sequence.Mamba.defaultLogPath : System.FilePath

Default JSON loss-curve path for this command.

def NN.Examples.Models.Sequence.Mamba.usage : String

Complete command help, including the text and training flags parsed after runtime selection.

def NN.Examples.Models.Sequence.Mamba.contextLength : ℕ

Training and generation context length for the Mamba text example.

def NN.Examples.Models.Sequence.Mamba.tokenizer : TorchLean.text.Tokenizer

Byte tokenizer used by this sequence model.

def NN.Examples.Models.Sequence.Mamba.vocabularySize : ℕ

Number of byte-token classes.

def NN.Examples.Models.Sequence.Mamba.modelWidth : ℕ

Width of each token embedding and of the Mamba block's output.

def NN.Examples.Models.Sequence.Mamba.paddingByte : ℕ

Byte value used to pad short text windows.

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.Mamba.modelConfig : TorchLean.nn.models.Mamba.Config

Mamba text-model configuration shared by shapes and the constructor.

theorem NN.Examples.Models.Sequence.Mamba.instNeZeroNatVocabularySize : NeZero vocabularySize

def NN.Examples.Models.Sequence.Mamba.byteIndex (id : ℕ) : Fin vocabularySize

Embed a byte id in the complete 256-entry byte vocabulary.

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.Mamba.input : Spec.Shape

Input shape: one sequence of one-hot byte tokens.

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.Mamba.output : Spec.Shape

Output shape: one vocabulary-logit row per input position.

def NN.Examples.Models.Sequence.Mamba.model : TorchLean.nn.Builder (TorchLean.nn.Sequential input output)

Public Mamba language-model constructor specialized to the example config.

structure NN.Examples.Models.Sequence.Mamba.Options : Type

Command-local training, sampling, and corpus-window controls.

• training : TorchLean.CLI.Training.OptimizerOptions

  Optimizer, step, batching, and logging controls.

• generation : TorchLean.text.GenerationOptions

  Prompt and token-sampling policy.

• window : TorchLean.text.WindowOptions

  Number of corpus windows exposed to training.

def NN.Examples.Models.Sequence.Mamba.instReprOptions.repr : Options → ℕ → Std.Format

@[instance_reducible]

instance NN.Examples.Models.Sequence.Mamba.instReprOptions : Repr Options

def NN.Examples.Models.Sequence.Mamba.Options.parse (args : List String) : Except String (Options × List String)

Parse the Mamba command's training and sampling flags.

def NN.Examples.Models.Sequence.Mamba.sampleFromTokenIds (ids : TorchLean.Tensor ℕ [contextLength + 1]) : TorchLean.Sample.Supervised Float input output

Convert a token window into the one-hot next-token sample consumed by the Mamba model.

def NN.Examples.Models.Sequence.Mamba.samplesFromCorpus (corpus : String) (windows : ℕ) : TorchLean.Data.SampleStream (TorchLean.Sample.Supervised Float input output)

Build a finite training set from approximately evenly spaced corpus windows.

def NN.Examples.Models.Sequence.Mamba.printPredictionReport (label prompt : String) (logits : TorchLean.Tensor Float output) : IO Unit

Print the current argmax prediction beside the prompt and shifted target text.

def NN.Examples.Models.Sequence.Mamba.inputTensorFromIds (ids : TorchLean.Tensor ℕ [contextLength]) : TorchLean.Tensor Float input

Convert a prompt window into the typed one-hot input tensor used during generation.

def NN.Examples.Models.Sequence.Mamba.generateSampled (predict : TorchLean.Tensor Float input → IO (TorchLean.Tensor Float output)) (generation : TorchLean.text.GenerationOptions) : IO String

Autoregressively extend a prompt using the trained Mamba parameters.

def NN.Examples.Models.Sequence.Mamba.trainOnText (runtime : TorchLean.Runtime.Config) (corpus : String) (options : Options) : IO (Float × Float)

Train the Mamba language model and print before/after prediction and generation reports.

def NN.Examples.Models.Sequence.Mamba.main (args : List String) : IO UInt32

CLI entrypoint for the Mamba text command.