add text-generation pipeline example with autocast
Browse files
README.md
CHANGED
@@ -91,6 +91,22 @@ from transformers import AutoTokenizer
|
|
91 |
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
92 |
```
|
93 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
94 |
## Community-Created Resources
|
95 |
|
96 |
These were not created by MosaicML, but you may find them useful. These links are not an endorsement of the creators or their content.
|
|
|
91 |
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
92 |
```
|
93 |
|
94 |
+
The model can then be used, for example, within a text-generation pipeline.
|
95 |
+
Note: when running Torch modules in lower precision, it is best practice to use the [torch.autocast context manager](https://pytorch.org/docs/stable/amp.html).
|
96 |
+
|
97 |
+
```python
|
98 |
+
from transformers import pipeline
|
99 |
+
|
100 |
+
pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, device='cuda:0')
|
101 |
+
|
102 |
+
with torch.autocast('cuda', dtype=torch.bfloat16):
|
103 |
+
print(
|
104 |
+
pipe('Here is a recipe for vegan banana bread:\n',
|
105 |
+
max_new_tokens=100,
|
106 |
+
do_sample=True,
|
107 |
+
use_cache=True))
|
108 |
+
```
|
109 |
+
|
110 |
## Community-Created Resources
|
111 |
|
112 |
These were not created by MosaicML, but you may find them useful. These links are not an endorsement of the creators or their content.
|