Update app.py

app.py

@@ -10,22 +10,34 @@ def img2text(url):
     return text
 
 # text2story
-def text2story(text):
+def text2story(scenario):
+    child_prompt = f"Create a children's story suitable for 3-10 years old about {scenario}. " \
+                   "Include talking animals and magical elements. " \
+                   "Use simple words and happy ending. " \
+                   "Story structure: Beginning -> Problem -> Solution -> Happy ending.\n\n"
+    
     pipe = pipeline("text-generation", 
                     model="pranavpsv/genre-story-generator-v2",
-                    max_new_tokens=100,  
-                    truncation=True)      
-    story_text = pipe(text, 
-                      do_sample=True,     
-                      temperature=0.9,    
-                      top_k=50,           
-                      num_return_sequences=1)[0]['generated_text']
+                    max_new_tokens=200,
+                    truncation=True,
+                    temperature=0.7,  # 降低随机性
+                    top_p=0.9,
+                    repetition_penalty=1.2)
+    
+    full_story = pipe(child_prompt + scenario,
+                     do_sample=True,
+                     num_return_sequences=1)[0]['generated_text']
+    
+    # 过滤不合适内容
+    safe_story = full_story.replace("violence", "").replace("scary", "").replace("died", "solved")
+    
+    # 确保以快乐结局结束
+    happy_endings = ["happily ever after", "big smile", "learned a good lesson"]
+    if not any(ending in safe_story.lower() for ending in happy_endings):
+        safe_story += " And they all lived happily ever after!"
     
-    # Make sure the story is complete
-    last_punctuation = max(story_text.rfind("."), story_text.rfind("!"), story_text.rfind("?"))
-    if last_punctuation != -1:
-        story_text = story_text[:last_punctuation+1]
-    return story_text
+    # 添加表情符号装饰
+    return "🌈 " + safe_story[:safe_story.rfind(".")+1] + " 🎉"
 
 # text2audio
 def text2audio(story_text):