import streamlit as st
from transformers import pipeline

# Set up the Streamlit app
st.title("🧚‍♀️ Magic Story Buddy 📚")
st.markdown("Let's create a magical story just for you!")

# Initialize the model
@st.cache_resource
def load_model():
    return pipeline("text-generation", model="ajibawa-2023/Young-Children-Storyteller-Mistral-7B")

model = load_model()

# User input
child_name = st.text_input("What's your name, young storyteller?")
story_theme = st.selectbox("What would you like your story to be about?", 
                           ["Space Adventure", "Magical Forest", "Underwater World", "Dinosaur Discovery"])

# Additional options
story_length = st.slider("How long should the story be?", 50, 200, 100)
include_moral = st.checkbox("Include a moral lesson?")

if st.button("Create My Story!"):
    if child_name and story_theme:
        # Construct the prompt
        prompt = f"[CHILDREN'S STORY] Once upon a time, in a {story_theme.lower()}, there was a brave child named {child_name}. "
        if include_moral:
            prompt += "This story teaches us that "
        
        # Generate the story
        story = model(prompt, max_length=story_length, num_return_sequences=1)[0]['generated_text']
        
        # Display the story
        st.markdown("## Your Magical Story")
        st.write(story)
        
        # Add a fun element
        st.balloons()
    else:
        st.warning("Please tell me your name and choose a story theme.")

# Add some child-friendly decorations
st.markdown("---")
st.markdown("🌟 Remember, you're the star of every story! 🌟")