# 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()