import gradio as gr
from transformers import pipeline

# Use the same text-generation model as the existing project
generator = pipeline("text-generation", model="gpt2")


def generate_text(prompt: str, max_length: int = 150):
    prompt = prompt or ""
    if not prompt.strip():
        return "Please enter a prompt to generate text."

    result = generator(
        prompt,
        max_length=max_length,
        do_sample=False,
        truncation=True,
    )
    return result[0]["generated_text"]


demo = gr.Blocks()

with demo:
    gr.Markdown("# Text Generator")
    gr.Markdown(
        "Enter a prompt below and click Generate Text to produce output using GPT-2."
    )

    with gr.Row():
        prompt_input = gr.Textbox(
            label="Prompt",
            placeholder="Enter your prompt here...",
            lines=5,
        )
        max_length = gr.Slider(
            minimum=50,
            maximum=500,
            step=10,
            value=150,
            label="Max generated length",
        )

    output_text = gr.Textbox(label="Generated Text", lines=12)

    generate_button = gr.Button("Generate Text")
    generate_button.click(
        generate_text,
        inputs=[prompt_input, max_length],
        outputs=output_text,
    )


demo.launch()