import gradio as gr
from transformers import pipeline

# Load the DistilGPT-2 text-generation pipeline
generator = pipeline("text-generation", model="distilgpt2")

def generate_text(prompt):
    result = generator(prompt, max_length=120, num_return_sequences=1)[0]["generated_text"]
    return result

# Two simple example prompts for students to try
example_prompts = [
    ["Once upon a time in a distant kingdom,"],
    ["The future of artificial intelligence depends on"]
]

with gr.Blocks(title="DistilGPT-2 Text Generator") as demo:
    
    gr.Markdown(
        """### DistilGPT-2 Text Generator  
Enter a short prompt and the model will continue it.  
Below the input box you will find two example prompts you can try immediately."""
    )
    
    prompt = gr.Textbox(
        lines=2,
        label="Prompt",
        placeholder="Type a short prompt to begin…"
    )
    
    output = gr.Textbox(
        lines=12,         # Increased height for the output area
        label="Generated Text"
    )
    
    btn = gr.Button("Generate")
    btn.click(generate_text, inputs=prompt, outputs=output)
    
    gr.Examples(
        examples=example_prompts,
        inputs=prompt,
        label="Example Prompts"
    )

demo.launch()