Skip to content

Commit

Permalink
quickfix documentation (#32566)
Browse files Browse the repository at this point in the history
* fix documentation

* update config
  • Loading branch information
molbap authored Aug 26, 2024
1 parent 3562772 commit 26f043b
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions docs/source/en/model_doc/mamba2.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,11 @@ The original code can be found [here](https://github.com/state-spaces/mamba).

### A simple generation example:
```python
from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
from transformers import Mamba2Config, Mamba2ForCausalLM, AutoTokenizer
import torch
model_id = 'mistralai/Mamba-Codestral-7B-v0.1'
tokenizer = AutoTokenizer.from_pretrained(model_id, revision='refs/pr/9', from_slow=True, legacy=False)
model = MambaForCausalLM.from_pretrained(model_id, revision='refs/pr/9')
model = Mamba2ForCausalLM.from_pretrained(model_id, revision='refs/pr/9')
input_ids = tokenizer("Hey how are you doing?", return_tensors= "pt")["input_ids"]

out = model.generate(input_ids, max_new_tokens=10)
Expand Down

0 comments on commit 26f043b

Please sign in to comment.