# 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.