Skip to content

Instantly share code, notes, and snippets.

@MichelNivard
Created March 15, 2023 08:59
Show Gist options
  • Save MichelNivard/ab45563f34b52238607a13b30615ee62 to your computer and use it in GitHub Desktop.
Save MichelNivard/ab45563f34b52238607a13b30615ee62 to your computer and use it in GitHub Desktop.
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
# Load the tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('./results')
# Set the device to GPU if available, otherwise use CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
# Generate some text
prompt = 'Once upon a time'
input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
output = model.generate(input_ids=input_ids, max_length=100, do_sample=True, temperature=0.7)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment