import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model and tokenizer
model_name = "mstftmk/shakespeare-gpt2"  # Replace with your model's repo
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Define the generation function
def generate_text(input_text, max_length, temperature):
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    outputs = model.generate(inputs, max_length=max_length, temperature=temperature)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your prompt..."),
        gr.Slider(50, 300, value=100, label="Max Length"),
        gr.Slider(0.1, 1.0, value=0.7, label="Temperature"),
    ],
    outputs=gr.Textbox(),
    title="Shakespeare GPT-2",
    description="Generate text inspired by Shakespeare.",
)

# Launch the Gradio app
interface.launch()