A simple PyTorch re-implementation of Mamba in a single file. minMamba tries to be small, clean, interpretable and educational.

If you want to import minmamba into your project:

git clone https://github.com/lckr/minMamba.git
cd minMamba
pip install -e .

Here's how you'd load a pretrained Mamba model from Huggingface Hub:

import minmamba.model
pretrained_model = minmamba.model.MambaLMModel.from_pretrained("state-spaces/mamba-130m")

And here's how you'd run inference with it:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") # tokenizer used by "state-spaces/mamba-130m"

input_seq = tokenizer("A fish is a ", return_tensors="pt")["input_ids"]
gen_seq = pretrained_model.generate(input_seq, 100)
print(tokenizer.decode(gen_seq[0]))

Code:

• state-spaces/mamba the official Mamba implementation released by the authors
• karpathy/minGPT Andrej Kaparthys mingpt
• rjb7731/nanoMamba Ryan Bradys nanoMamba implementation

Paper:

Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Albert Gu*, Tri Dao*
Paper: https://arxiv.org/abs/2312.00752

MIT