Mamba Text Training

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

The model is trainable end-to-end:

mamba(seqLen, vocab, stateDim) → linear(stateDim → vocab)

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

python3 scripts/datasets/download_example_data.py --tiny-shakespeare
lake exe -K cuda=true torchlean mamba --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.defaultLogJson : System.FilePath

Default JSON loss-curve path for this command.

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

Training and generation context length for the Mamba text example.

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

Byte tokenizer used by this sequence model.

def NN.Examples.Models.Sequence.Mamba.cfg : TorchLean.nn.models.MambaTextConfig

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

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.Mamba.σ : Shape

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

@[reducible, inline]

abbrev NN.Examples.Models.Sequence.Mamba.τ : Shape

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

def NN.Examples.Models.Sequence.Mamba.model : TorchLean.nn.M (TorchLean.nn.Sequential σ τ)

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

def NN.Examples.Models.Sequence.Mamba.sampleFromTokenIds (ids : List ℕ) : TorchLean.SupervisedSample Float σ τ

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

def NN.Examples.Models.Sequence.Mamba.samplesFromCorpus (input _prompt : String) (windows : ℕ) : Array (TorchLean.SupervisedSample Float σ τ)

Build a finite cyclic training set from corpus text, biased toward the prompt when present.

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

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

def NN.Examples.Models.Sequence.Mamba.inputTensorFromIds (ids : List ℕ) : TorchLean.Tensor.T Float σ

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

def NN.Examples.Models.Sequence.Mamba.generateSampled (predict : TorchLean.Tensor.T Float σ → IO (TorchLean.Tensor.T Float τ)) (prompt : String) (steps : ℕ) (temperature : Float) (topK seed : ℕ) : IO String

Autoregressively extend a prompt using the trained Mamba parameters.

def NN.Examples.Models.Sequence.Mamba.trainOnText (opts : TorchLean.Options) (input : String) (train : API.text.WindowedTrainGenerationOptions) : 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.