BasicTextGen / app.py
pcasale's picture
Update app.py
d895e1d verified
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()