import gradio as gr
from transformers import pipeline

generator = pipeline(
    "text-generation",
    model="distilgpt2",
    device=-1
)

def generate_text(prompt, temperature, max_tokens):
    if not prompt or not prompt.strip():
        return "Please enter a prompt."

    result = generator(
        prompt,
        max_new_tokens=int(max_tokens),
        temperature=float(temperature),
        do_sample=True,
        pad_token_id=50256,
        num_return_sequences=1
    )
    return result[0]["generated_text"]

demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2),
        gr.Slider(0.1, 1.2, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(20, 80, value=40, step=10, label="Max Tokens")
    ],
    outputs=gr.Textbox(label="Generated Text", lines=6),
    title="Text Generator",
    description="Generate text with DistilGPT-2"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)