daking vchiley commited on
Commit
55363b1
1 Parent(s): ea983eb

add text-generation pipeline example with autocast (#40)

Browse files

- add text-generation pipeline example with autocast (8cee036ddd83d886eb157a72ab1c3648c0753bd8)


Co-authored-by: Vitaliy Chiley <vchiley@users.noreply.huggingface.co>

Files changed (1) hide show
  1. README.md +16 -0
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.