Views
No views yet
mamba2-torch offers different optimisation paths to use:1git clone https://github.com/vasqu/mamba2-torch.git
2cd mamba2-torch
3pip install .1from transformers import AutoTokenizer
2from mamba2_torch import Mamba2Model, Mamba2ForCausalLM, Mamba2Config
3
4device = "cuda"
5mamba2_hf_path = "<path-to-converted-model>"
6
7model = Mamba2ForCausalLM.from_pretrained(mamba2_hf_path, local_files_only=True).to(device)
8tokenizer = AutoTokenizer.from_pretrained(mamba2_hf_path, local_files_only=True)
9
10input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"].to(device)
11
12# expected output (2.7b): `["Hey how are you doing? I'm doing good. I'm just trying to"]`
13out = model.generate(input_ids, max_new_tokens=10)
14print(tokenizer.batch_decode(out))1@inproceedings{mamba2,
2 title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
3 author={Dao, Tri and Gu, Albert},
4 booktitle={International Conference on Machine Learning (ICML)},
5 year={2024}
6}