import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# Load model (ensure Space is using GPU!)
model_id = "ajibawa-2023/Young-Children-Storyteller-Mistral-7B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_prompt(characters, theme, setting, scenario):
    return f"""You are a children's story writer.

Characters: {characters}
Theme: {theme}
Setting: {setting}
Scenario: {scenario}

Write a short story suitable for children aged 6–12. It should have a beginning, middle, and a happy ending. Use simple language, avoid extra characters, and include a positive message.

Story:"""

def generate_story(characters, theme, setting, scenario):
    prompt = generate_prompt(characters, theme, setting, scenario)
    result = generator(prompt, max_new_tokens=300, temperature=0.8, do_sample=True)[0]["generated_text"]
    return result[len(prompt):].strip()

# Gradio UI
with gr.Blocks(title="Genieverse: Story Generator for Kids") as app:
    gr.Markdown("## 📚 Genieverse – AI Story Creator for Children 6–12")
    with gr.Row():
        characters = gr.Dropdown(["Girl", "Boy", "Wolf", "Elephant", "Horse", "Dog", "Bird"], label="Select Characters", multiselect=True, value=["Girl", "Wolf"])
        theme = gr.Dropdown(["Bravery", "Friendship", "Kindness", "Helping Others"], label="Select Theme", value="Bravery")
    with gr.Row():
        setting = gr.Dropdown(["Forest", "Jungle", "Mountain", "Village", "River"], label="Select Setting", value="Forest")
        scenario = gr.Textbox(label="Enter a Scenario", placeholder="e.g., A girl saw a wolf coming toward her", value="A girl saw a wolf coming toward her")
    generate_btn = gr.Button("Generate Story ✨")
    story_output = gr.Textbox(label="Generated Story", lines=15)

    generate_btn.click(fn=generate_story, inputs=[characters, theme, setting, scenario], outputs=story_output)

app.launch()