Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3tokenizer = AutoTokenizer.from_pretrained("AntonV/mamba2-130m-hf")
4model = AutoModelForCausalLM.from_pretrained("AntonV/mamba2-130m-hf")
5
6input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"]
7out = model.generate(input_ids, max_new_tokens=10)
8print(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}