add text-generation pipeline example with autocast (#47)
Browse files- add text-generation pipeline example with autocast (2faa761bc10cdff64021541e93fa9b2f67482bf6)
Co-authored-by: Vitaliy Chiley <[email protected]>
README.md
CHANGED
|
@@ -102,6 +102,22 @@ from transformers import AutoTokenizer
|
|
| 102 |
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
| 103 |
```
|
| 104 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 105 |
### Formatting
|
| 106 |
|
| 107 |
This model was trained on data formatted in the dolly-15k format:
|
|
|
|
| 102 |
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
| 103 |
```
|
| 104 |
|
| 105 |
+
The model can then be used, for example, within a text-generation pipeline.
|
| 106 |
+
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).
|
| 107 |
+
|
| 108 |
+
```python
|
| 109 |
+
from transformers import pipeline
|
| 110 |
+
|
| 111 |
+
pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, device='cuda:0')
|
| 112 |
+
|
| 113 |
+
with torch.autocast('cuda', dtype=torch.bfloat16):
|
| 114 |
+
print(
|
| 115 |
+
pipe('Here is a recipe for vegan banana bread:\n',
|
| 116 |
+
max_new_tokens=100,
|
| 117 |
+
do_sample=True,
|
| 118 |
+
use_cache=True))
|
| 119 |
+
```
|
| 120 |
+
|
| 121 |
### Formatting
|
| 122 |
|
| 123 |
This model was trained on data formatted in the dolly-15k format:
|