Update app.py

app.py

@@ -3,6 +3,7 @@ import streamlit as st
 from transformers import pipeline
 from gtts import gTTS
 import io
+import re
 
 
 # function part
@@ -15,61 +16,15 @@ def img2text(url):
 
 # text2story
 def text2story(text):
-    # 儿童故事提示模板
-    child_prompt = f"""Write a children's story for ages 3-10 based on: {text}
-    Requirements:
-    - Use simple words (1st-3rd grade level)
-    - Main character: friendly animal
-    - Story elements: magic, friendship, happy ending
-    - Moral lesson: sharing is caring
-    - Length: 100-120 words
-
-    Story structure:
-    1. Introduce characters and setting
-    2. A problem occurs
-    3. Magical solution with teamwork
-    4. Happy ending with a lesson
-    """
-    
-    pipe = pipeline(
-        "text-generation", 
-        model="pranavpsv/genre-story-generator-v2",
-        max_new_tokens=160,    # 160 tokens ≈ 120 words
-        min_new_tokens=134,    # 134 tokens ≈ 100 words
-        temperature=0.7,       # 平衡创造性与稳定性
-        top_p=0.9,
-        repetition_penalty=1.2,
-        pad_token_id=50256,    # 对齐模型config的bos/eos_token_id
-        num_return_sequences=1
-    )
-    
-    # 生成故事
-    raw_story = pipe(child_prompt, return_full_text=False)[0]['generated_text']
-    
-    # 智能后处理
-    def format_story(text):
-        # 移除技术性标记
-        text = text.replace("<|endoftext|>", "").strip()
-        # 寻找自然结尾点
-        end_marks = ['.', '!', '?', '...']
-        last_positions = [text.rfind(mark) for mark in end_marks]
-        valid_positions = [pos for pos in last_positions if pos != -1]
-        cutoff = max(valid_positions) + 1 if valid_positions else len(text)
-        return text[:cutoff]
-    
-    formatted_story = format_story(raw_story)
-    
-    # 字数验证
-    word_count = len(formatted_story.split())
-    if word_count < 100:
-        formatted_story += " And they lived happily ever after."
-    
-    return formatted_story
+    prompt = f"Generate a children's story for 3-10 ages based on {text}"
+    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2", max_new_tokens=160, min_new_tokens=130, num_return_sequences=1)
+    story_text = pipe(prompt)[0]['generated_text']
+    return story_text
 
 
 # text2audio
 def text2audio(story_text):
-    tts = gTTS(text=story_text, lang='en')
+    tts = gTTS(text = story_text, lang='en')
     audio_bytes = io.BytesIO()
     tts.write_to_fp(audio_bytes)
     audio_bytes.seek(0)
@@ -120,20 +75,20 @@ def main():
         progress_bar = st.progress(0)
 
         # Stage 1: Image to Text
-        with status_container.status("🔮 **Step 1/3**: Decoding picture magic...", expanded=True) as status:
+        with status_container.status("🔮 **Step 1/3**: Decoding picture magic...", expanded=True) as status:  # Add progress bar components to improve experience
             progress_bar.progress(33)
             scenario = img2text(uploaded_file.name)
             status.update(label="✅ Picture decoded!", state="complete")
             st.write(f"**What I see:** {scenario}")
 
-        # Stage 2: Text to Story
+        #Stage 2: Text to Story
         with status_container.status("📚 **Step 2/3**: Writing your fairy tale...", expanded=True) as status:
             progress_bar.progress(66)
             story = text2story(scenario)
             status.update(label="✅ Story created!", state="complete")
             st.write(f"**Your Story:**\n{story}")
 
-        # Stage 3: Story to Audio data
+        #Stage 3: Story to Audio data
         with status_container.status("🎵 **Step 3/3**: Adding magic audio...", expanded=True) as status:
             progress_bar.progress(100)
             audio_data = text2audio(story)