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 300 --windows 128 \
  --temperature 0.85 --top-k 12 --sample-seed 7

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

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

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

Training/generation context length in byte tokens.

Mamba scales more gently with sequence length than attention, so the tutorial uses a 64-byte window. That is long enough to carry speaker tags and short phrases from Tiny Shakespeare while remaining fast in eager CUDA.

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

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

@[reducible, inline]

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

@[reducible, inline]

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

def NN.Examples.Models.Sequence.Mamba.mkModel : API.nn.M (API.nn.Sequential σ τ)

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

• base : API.Common.ModelTrainFlags
• windows : ℕ
• prompt : String
• generate : ℕ
• temperature : Float
• topK : ℕ
• seed : ℕ

@[implicit_reducible]

instance NN.Examples.Models.Sequence.Mamba.instReprTrainOptions : Repr TrainOptions

def NN.Examples.Models.Sequence.Mamba.instReprTrainOptions.repr : TrainOptions → ℕ → Std.Format

def NN.Examples.Models.Sequence.Mamba.TrainOptions.steps (train : TrainOptions) : ℕ

def NN.Examples.Models.Sequence.Mamba.TrainOptions.lr (train : TrainOptions) : Float

def NN.Examples.Models.Sequence.Mamba.TrainOptions.log (train : TrainOptions) : Runtime.Training.LogDestination

def NN.Examples.Models.Sequence.Mamba.TrainOptions.logPath (train : TrainOptions) : System.FilePath

def NN.Examples.Models.Sequence.Mamba.parseTrainOptions (args : List String) : Except String (TrainOptions × List String)

def NN.Examples.Models.Sequence.Mamba.castTensor {s : Shape} (t : Spec.Tensor Float s) : Spec.Tensor Float s

def NN.Examples.Models.Sequence.Mamba.sampleFromTokenIds (ids : List ℕ) : API.sample.Supervised Float σ τ

def NN.Examples.Models.Sequence.Mamba.samplesFromCorpus (input prompt : String) (windows : ℕ) : Array (API.sample.Supervised Float σ τ)

def NN.Examples.Models.Sequence.Mamba.firstSample (samples : Array (API.sample.Supervised Float σ τ)) : API.sample.Supervised Float σ τ

def NN.Examples.Models.Sequence.Mamba.printPredictionProbe (label prompt : String) (logits : Spec.Tensor Float τ) : IO Unit

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

def NN.Examples.Models.Sequence.Mamba.logitsArrayAt (logits : Spec.Tensor Float τ) (pos : ℕ) : Array Float

def NN.Examples.Models.Sequence.Mamba.greedyTokenAt (logits : Spec.Tensor Float τ) (pos : ℕ) : ℕ

def NN.Examples.Models.Sequence.Mamba.sampleFromLogitsAt (logits : Spec.Tensor Float τ) (pos : ℕ) (temperature : Float) (topK seed counter : ℕ) : ℕ

def NN.Examples.Models.Sequence.Mamba.generateSampled (opts : Runtime.Autograd.Torch.Options) (model : API.nn.Sequential σ τ) (params : Runtime.Autograd.Torch.ParamList Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model)) (prompt : String) (steps : ℕ) (temperature : Float) (topK seed : ℕ) : IO String

partial def NN.Examples.Models.Sequence.Mamba.generateSampled.loop (opts : Runtime.Autograd.Torch.Options) (model : API.nn.Sequential σ τ) (params : Runtime.Autograd.Torch.ParamList Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model)) (steps : ℕ) (temperature : Float) (topK seed : ℕ) (ids : List ℕ) : ℕ → IO (List ℕ)

def NN.Examples.Models.Sequence.Mamba.meanLossOnSamples (model : API.nn.Sequential σ τ) (m : Runtime.Autograd.TorchLean.ScalarModule Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model) [σ, τ]) (samples : Array (API.sample.Supervised Float σ τ)) : IO Float

def NN.Examples.Models.Sequence.Mamba.trainOnText (opts : Runtime.Autograd.Torch.Options) (input : String) (train : TrainOptions) : IO (Float × Float)

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