from transformers import pipeline
import gradio as gr

# Hugging Face model (lightweight GPT variant)
generator = pipeline("text-generation", model="distilgpt2")

def generate(prompt):
    result = generator(
        prompt,
        max_length=520,
        num_return_sequences=1,
        temperature=0.9
    )
    return result[0]["generated_text"]

def space_story():
    return generate("In the year 2050, humans discovered a new galaxy where...")

def horror_story():
    return generate("It was a dark night and suddenly I heard a strange sound inside the house..")

def motivation():
    return generate("Never stop trying because success comes after many failures and..")

with gr.Blocks() as demo:
    gr.Markdown("# AI Story Generator (HF Model)")

    with gr.Tab("Space Story"):
        gr.Button("Generate").click(space_story, outputs=gr.Textbox())

    with gr.Tab("Horror Story"):
        gr.Button("Generate").click(horror_story, outputs=gr.Textbox())

    with gr.Tab("Motivation"):
        gr.Button("Generate").click(motivation, outputs=gr.Textbox())

demo.launch()