BasicTextGen / app.py
pcasale's picture
Create app.py
a90147a verified
Raw
History Blame
1.59 kB
# Import Gradio and the Hugging Face Transformers pipeline
import gradio as gr
from transformers import pipeline
# Load a text-generation pipeline with the DistilGPT-2 model.
# DistilGPT-2 is a distilled (compressed) version of GPT-2, so it's faster and lightweight:contentReference[oaicite:3]{index=3}.
generator = pipeline("text-generation", model="distilgpt2") # This downloads the model weights if not already available
# Define a function that uses the generator to produce text based on the input prompt.
def generate_text(prompt):
# Use the text-generation pipeline to continue the prompt.
# We set a max_length to limit the output length for practicality.
result = generator(prompt, max_length=100, num_return_sequences=1)[0]["generated_text"]
return result
# Set up the Gradio interface:
# - Input: a textbox for the prompt (single-line or a short prompt, so we use lines=2).
# - Output: a textbox for the generated text.
# - We also add a title and description for user guidance.
input_prompt = gr.Textbox(lines=2, label="Prompt", placeholder="Enter a text prompt...")
output_text = gr.Textbox(label="Generated Text")
demo = gr.Interface(
fn=generate_text,
inputs=input_prompt,
outputs=output_text,
title="🤖 DistilGPT-2 Text Generator",
description="**Description:** Enter a prompt, and the DistilGPT-2 language model will continue the text. "
"This demonstrates basic text generation using a small pre-trained GPT-2 model."
)
# Launch the app (in a Hugging Face Space, this will run automatically).
demo.launch()