dtm95's picture
Update app.py
b6a7d80 verified
raw
history blame
2.36 kB
import streamlit as st
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
# Set up the Streamlit app
st.title("πŸ§šβ€β™€οΈ Magic Story Buddy πŸ“š")
st.markdown("Let's create a magical story just for you!")
def initialize_model():
# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
# Create the pipeline for text generation
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
return text_generator
# Initialize the model and pipeline
pipe = initialize_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):
messages = [{"role": "user", "content": prompt}]
result = pipe(messages, max_length=max_length, do_sample=True, temperature=0.7)
return result[0]['generated_text']
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! 🌟")