from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
import gradio as gr

# Load the tokenizer and the model
tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
model = GPT2LMHeadModel.from_pretrained("gpt2-medium")

# Function for generating text
def generate_text(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs.input_ids, max_length=100, num_return_sequences=1)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Define a function that will be used in the Gradio interface
def generate_from_prompt(prompt):
    return generate_text(prompt)

# Create the Gradio interface
gr.Interface(fn=generate_from_prompt, inputs="text", outputs="text", title="GPT-2 Medium Text Generator").launch()