app.py

import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 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():
    model = AutoModelForCausalLM.from_pretrained("ajibawa-2023/Young-Children-Storyteller-Mistral-7B", torch_dtype=torch.float16)
    tokenizer = AutoTokenizer.from_pretrained("ajibawa-2023/Young-Children-Storyteller-Mistral-7B")
    return model, tokenizer

model, tokenizer = 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?")

def generate_story(prompt, max_length=500):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=max_length, num_return_sequences=1, do_sample=True, temperature=0.7)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

if st.button("Create My Story!"):
    if child_name and story_theme:
        # Construct the prompt
        prompt = f"""Create a short children's story with the following details:
        - Main character: {child_name}
        - Theme: {story_theme}
        - Length: About {story_length} words
        - Audience: Children aged 5-10
        - Tone: Friendly, educational, and imaginative

        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 = generate_story(prompt, max_length=story_length)
        
        # 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! 🌟")