File size: 2,349 Bytes
aad51d1 5df603e 465f3e7 5df603e 465f3e7 5df603e 465f3e7 5df603e 465f3e7 5df603e 465f3e7 5df603e 465f3e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
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! 🌟") |