GPT2-Jokes / test.py
ameerazam08's picture
Create test.py
2d03482
# Load the trained model and tokenizer
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel,GPT2Model
from transformers import TFAutoModel
model_path = "./trained_gpt2_jokes/5/"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = GPT2LMHeadModel.from_pretrained(model_path, from_tf=True).to(device)
tokenizer = GPT2Tokenizer.from_pretrained(model_path)
# Put the model in eval mode
model.eval()
# Define a function for generating text
def generate_text(prompt_text):
# Encode the prompt text and move to GPU
input_ids = tokenizer.encode(prompt_text, return_tensors="pt").to(device)
# Generate text
with torch.no_grad():
output = model.generate(input_ids, max_length=100, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id, temperature=0.5)
# Decode the generated text
decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
return decoded_output
# Testing the model with a sample prompt
prompt = "Why man are "
generated_joke = generate_text(prompt)
print(generated_joke)
#result getting --->Why man are so upset about the new $20 bill? Because it's only worth $15 more than $20.