btsv1 / app.py
TomMoule91's picture
Update app.py
c1ad901
import streamlit as st
from transformers import pipeline
# Load the text generation pipeline
generator = pipeline("text-generation")
def generate_bedtime_story(name, interests, friends, age):
story_prompt = f"I want to tell a bed time story to a child named {name}. I want the story to revolve around some of {name}'s main interests, which are {interests}. I also want the story to include some of {name}'s friends, who are called: {friends}. Make sure the story is suitable for a child of the age {age}."
story = generator(story_prompt, max_length=1000)[0]['generated_text']
return story
def main():
st.title("Bedtime Story Generator")
name = st.text_input("Enter the name of the person you're reading to:")
age = st.text_input("Enter the age of the person you're reading to:")
interests = st.text_input("Enter some of things the person is interested in:")
friends = st.text_input("Enter the names of some of their friends:")
if st.button("Generate Story"):
if name and interests and friends and age:
story = generate_bedtime_story(name, interests, age, friends.split(","))
st.write(story)
else:
st.warning("Please fill in all the inputs.")
if __name__ == "__main__":
main()