Update app.py

app.py

@@ -15,9 +15,45 @@ def img2text(url):
 
 # text2story
 def text2story(text):
-    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2",max_new_tokens=150, min_new_tokens=130)
-    story_text = pipe(text)[0]['generated_text']
-    return story_text
+    # 定义提示词模板（包含变量占位符）
+    prompt_template = f"""Write a children's story for ages 3-10 based on: {text}
+    Requirements:
+    - Use simple words
+    - Include a happy ending
+    """
+    
+    # 填充模板中的变量
+    full_prompt = prompt_template.format(text=text)
+    
+    # 初始化生成管道
+    pipe = pipeline(
+        "text-generation", 
+        model="pranavpsv/genre-story-generator-v2",
+        max_new_tokens=180,
+        min_new_tokens=130,
+        temperature=0.7
+    )
+    
+    # 生成原始文本
+    raw_output = pipe(full_prompt, return_full_text=False)[0]['generated_text']
+    
+    # 增强版提示词移除功能
+    def clean_output(generated_text, prompt):
+        # 方法1：精确匹配移除
+        if generated_text.startswith(prompt):
+            return generated_text[len(prompt):].strip()
+        
+        # 方法2：正则表达式模糊匹配
+        import re
+        pattern = re.compile(r'Write a children\'s story.*?based on:.*?\n', re.DOTALL)
+        cleaned = re.sub(pattern, '', generated_text, count=1)
+        
+        # 移除残留的提示词片段
+        cleaned = cleaned.split("Requirements:")[0].strip()
+        return cleaned
+
+    # 返回处理后的干净文本
+    return clean_output(raw_output, full_prompt)
 
 
 # text2audio