File size: 526 Bytes
5fa1a76 |
1 2 3 4 5 6 7 8 9 10 |
For example, you can use the [TextStreamer] class to stream the output of generate() into
your screen, one word at a time:
thon
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
inputs = tok(["An increasing sequence: one,"], return_tensors="pt")
streamer = TextStreamer(tok)
Despite returning the usual output, the streamer will also print the generated text to stdout. |