File size: 2,346 Bytes
6c36a30
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
---
tags:
- mamba2
license: mit
---

# mamba2-1.3b-av

## Introduction
This is a mirror model to [mamba2-1.3b](https://huggingface.co/state-spaces/mamba2-1.3b) which is compatible with [mamba2-torch](https://github.com/vasqu/mamba2-torch), a Hugging Face compatible mamba2 library that is not dependent on the original cuda wheels of the [original mamba repo](https://github.com/state-spaces/mamba). Credit goes to the original authors of [Mamba2](https://arxiv.org/abs/2405.21060) and the [transformers](https://github.com/huggingface/transformers) library by Hugging Face. Without their work, this would not be possible.


NOTE: `mamba2-torch` offers different optimisation paths to use:
- Triton kernels and [causal-conv1d](https://github.com/Dao-AILab/causal-conv1d) ("fastest")
- Triton kernels only (default) 
- Pure PyTorch

## How to Get Started with the Model
You can follow the instructions in the [mamba2-torch repo](https://github.com/vasqu/mamba2-torch) for a more detailed explanation. First of all, you should install the mamba2-torch lib:
```bash
git clone https://github.com/vasqu/mamba2-torch.git
cd mamba2-torch
pip install .
```

Then you can download this repository here via git lfs and then use the files locally the following way (after installing mamba2-torch):
```python
from transformers import AutoTokenizer
from mamba2_torch import Mamba2Model, Mamba2ForCausalLM, Mamba2Config

device = "cuda"
mamba2_hf_path = "<path-to-converted-model>"

model = Mamba2ForCausalLM.from_pretrained(mamba2_hf_path, local_files_only=True).to(device)
tokenizer = AutoTokenizer.from_pretrained(mamba2_hf_path, local_files_only=True)

input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"].to(device)

# expected output (1.3b): `["Hey how are you doing? I'm doing good. I'm doing good."]`
out = model.generate(input_ids, max_new_tokens=10)
print(tokenizer.batch_decode(out))
``` 

## Citation

<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->

**BibTeX:**

```bibtex
@inproceedings{mamba2,
 title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
 author={Dao, Tri and Gu, Albert},
 booktitle={International Conference on Machine Learning (ICML)},
 year={2024}
}
```