dtm95 commited on
Commit
465f3e7
Β·
verified Β·
1 Parent(s): 42c720f

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +45 -0
main.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import as st
2
+ from transformers import pipeline
3
+
4
+ # Set up the Streamlit app
5
+ st.title("πŸ§šβ€β™€οΈ Magic Story Buddy πŸ“š")
6
+ st.markdown("Let's create a magical story just for you!")
7
+
8
+ # Initialize the model
9
+ @st.cache_resource
10
+ def load_model():
11
+ return pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator")
12
+
13
+ model = load_model()
14
+
15
+ # User input
16
+ child_name = st.text_input("What's your name, young storyteller?")
17
+ story_theme = st.selectbox("What would you like your story to be about?",
18
+ ["Space Adventure", "Magical Forest", "Underwater World", "Dinosaur Discovery"])
19
+
20
+ # Additional options
21
+ story_length = st.slider("How long should the story be?", 50, 200, 100)
22
+ include_moral = st.checkbox("Include a moral lesson?")
23
+
24
+ if st.button("Create My Story!"):
25
+ if child_name and story_theme:
26
+ # Construct the prompt
27
+ prompt = f"[CHILDREN'S STORY] Once upon a time, in a {story_theme.lower()}, there was a brave child named {child_name}. "
28
+ if include_moral:
29
+ prompt += "This story teaches us that "
30
+
31
+ # Generate the story
32
+ story = model(prompt, max_length=story_length, num_return_sequences=1)[0]['generated_text']
33
+
34
+ # Display the story
35
+ st.markdown("## Your Magical Story")
36
+ st.write(story)
37
+
38
+ # Add a fun element
39
+ st.balloons()
40
+ else:
41
+ st.warning("Please tell me your name and choose a story theme.")
42
+
43
+ # Add some child-friendly decorations
44
+ st.markdown("---")
45
+ st.markdown("🌟 Remember, you're the star of every story! 🌟")