infer.py 277 Bytes
Newer Older
chenpangpang's avatar
chenpangpang committed
1
2
3
4
5
from transformers import pipeline, set_seed
generator = pipeline('text-generation', model='openai-community/gpt2', device='cuda')
output = generator("Hello, I'm a language model,", truncation=True, max_length=30, num_return_sequences=1)[0][
    "generated_text"]
print(output)