import gradio as gr
from transformers import pipeline
import torch

# Load GPT-2 once when the app starts
generator = pipeline(
    "text-generation",
    model="gpt2",
    device=-1  # use CPU on free Spaces
)

def generate_text(prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
    prompt = prompt.strip()
    if not prompt:
        return "Please enter a prompt."

    try:
        result = generator(
            prompt,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
            repetition_penalty=float(repetition_penalty),
            num_return_sequences=1,
            return_full_text=False,
            pad_token_id=generator.tokenizer.eos_token_id,
        )
        return result[0]["generated_text"].strip()

    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# GPT-2 Text Generator")
    gr.Markdown("Enter a prompt and generate text using the pretrained GPT-2 model.")

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            placeholder="Type your prompt here...",
            lines=4
        )

    with gr.Row():
        max_new_tokens = gr.Slider(20, 200, value=80, step=1, label="Max New Tokens")
        temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
    
    with gr.Row():
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")
        top_k = gr.Slider(0, 100, value=50, step=1, label="Top-k")

    repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.1, label="Repetition Penalty")

    submit = gr.Button("Generate")
    output = gr.Textbox(label="Generated Text", lines=10)

    submit.click(
        fn=generate_text,
        inputs=[prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
        outputs=output
    )

demo.launch()