Fix doctest more (for docs/source/en) (#30247)
* fix * fix --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
@@ -37,19 +37,21 @@ We also provide `StableLM Zephyr 3B`, an instruction fine-tuned version of the m
|
||||
The following code snippet demonstrates how to use `StableLM 3B 4E1T` for inference:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
>>> from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
|
||||
>>> device = "cuda" # the device to load the model onto
|
||||
|
||||
>>> set_seed(0)
|
||||
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t")
|
||||
>>> model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t")
|
||||
>>> model.to(device)
|
||||
>>> model.to(device) # doctest: +IGNORE_RESULT
|
||||
|
||||
>>> model_inputs = tokenizer("The weather is always wonderful in", return_tensors="pt").to(model.device)
|
||||
|
||||
>>> generated_ids = model.generate(**model_inputs, max_length=32, do_sample=True)
|
||||
>>> responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
>>> responses
|
||||
['The weather is always wonderful in Santa Barbara and, for visitors hoping to make the move to our beautiful seaside city, this town offers plenty of great places to...']
|
||||
['The weather is always wonderful in Costa Rica, which makes it a prime destination for retirees. That’s where the Pensionado program comes in, offering']
|
||||
```
|
||||
|
||||
## Combining StableLM and Flash Attention 2
|
||||
@@ -66,19 +68,21 @@ Now, to run the model with Flash Attention 2, refer to the snippet below:
|
||||
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
>>> from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
|
||||
>>> device = "cuda" # the device to load the model onto
|
||||
|
||||
>>> set_seed(0)
|
||||
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t")
|
||||
>>> model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
|
||||
>>> model.to(device)
|
||||
>>> model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2") # doctest: +SKIP
|
||||
>>> model.to(device) # doctest: +SKIP
|
||||
|
||||
>>> model_inputs = tokenizer("The weather is always wonderful in", return_tensors="pt").to(model.device)
|
||||
|
||||
>>> generated_ids = model.generate(**model_inputs, max_length=32, do_sample=True)
|
||||
>>> responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
>>> responses
|
||||
['The weather is always wonderful in Santa Barbara and, for visitors hoping to make the move to our beautiful seaside city, this town offers plenty of great places to...']
|
||||
>>> generated_ids = model.generate(**model_inputs, max_length=32, do_sample=True) # doctest: +SKIP
|
||||
>>> responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) # doctest: +SKIP
|
||||
>>> responses # doctest: +SKIP
|
||||
['The weather is always wonderful in Costa Rica, which makes it a prime destination for retirees. That’s where the Pensionado program comes in, offering']
|
||||
```
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user