import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

MODEL_NAME = "gpt2"

tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
model = GPT2LMHeadModel.from_pretrained(MODEL_NAME)

def generate_text(prompt, max_length, temperature):
    if not prompt.strip():
        return "Please enter a prompt."

    inputs = tokenizer.encode(prompt, return_tensors="pt")

    outputs = model.generate(
        inputs,
        max_length=max_length,
        do_sample=True,
        temperature=temperature,
        top_k=50,
        top_p=0.95,
        num_return_sequences=1,
        pad_token_id=tokenizer.eos_token_id
    )

    generated_text = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True
    )

    return generated_text

demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(
            lines=4,
            placeholder="Enter a prompt...",
            label="Prompt"
        ),
        gr.Slider(
            minimum=50,
            maximum=300,
            value=100,
            step=10,
            label="Maximum Length"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=0.8,
            step=0.1,
            label="Temperature"
        )
    ],
    outputs=gr.Textbox(
        label="Generated Text",
        lines=10
    ),
    title="GPT-2 Text Generation",
    description="Generate text using the pre-trained GPT-2 model with adjustable temperature."
)

if __name__ == "__main__":
    demo.launch()