Update app.py

app.py

@@ -1,4 +1,4 @@
-# import part
+import re
 import streamlit as st
 from transformers import pipeline
 import textwrap
@@ -19,24 +19,52 @@ def load_pipelines():
 
 captioner, storyer, tts = load_pipelines()
 
-# Function part
-# Function to generate content from an image
+def clean_generated_story(raw_story: str) -> str:
+    """
+    Cleans the generated story by:
+    1. Removing digits.
+    2. Removing words that are likely random letter combinations based on having no vowels.
+    3. Removing single-letter words unless they are allowed (such as 'a' or 'I').
+    """
+    # Remove all digits using regex
+    story_without_numbers = re.sub(r'\d+', '', raw_story)
+    
+    vowels = set('aeiouAEIOU')
+    
+    def is_valid_word(word: str) -> bool:
+        # Allow "a" and "I" for single-letter words
+        if len(word) == 1 and word.lower() not in ['a', 'i']:
+            return False
+        # For words longer than one letter, filter out those that do not contain any vowels
+        if len(word) > 1 and not any(char in vowels for char in word):
+            return False
+        return True
+
+    # Split the story into words, apply filtering, and recombine into a clean story
+    words = story_without_numbers.split()
+    filtered_words = [word for word in words if is_valid_word(word)]
+    
+    # Optionally, you can trim the clean story to a certain word count
+    clean_story = " ".join(filtered_words[:100])
+    return clean_story
+
 def generate_content(image):
     pil_image = Image.open(image)
     
-    # Generate caption
+    # Generate caption from the image
     caption = captioner(pil_image)[0]["generated_text"]
     st.write("**🌟 What's in the picture: 🌟**")
     st.write(caption)
 
-    # Create prompt for story
+    # Create prompt for the story
+    # Notice there’s no need to include the extra cleaning instructions in this prompt,
+    # because our code handles them later.
     prompt = (
-        f"Write a funny, interesting children's story that is precisely centered on this scene {caption}\nStory:"
+        f"Write a funny, interesting story for young children precisely centered on this scene {caption}\nStory:"
         f" mention the exact place, location or venue within {caption}"
-        f" in third-person narrative, that describes this scene exactly: {caption} "
-        f" Avoid numbers, random letter combinations, and single-letter combinations."
     )
-    # Generate raw story
+    
+    # Generate raw story from the model
     raw = storyer(
         prompt,
         max_new_tokens=150,
@@ -45,38 +73,30 @@ def generate_content(image):
         no_repeat_ngram_size=2,
         return_full_text=False
     )[0]["generated_text"].strip()
-
-    # Define allowed characters to keep (removes digits and symbols like * and ~)
-    allowed_chars = string.ascii_letters + " .,!?\"'-"
-    
-    # Clean the raw story by keeping only allowed characters (this filters out any digits)
-    clean_raw = ''.join(c for c in raw if c in allowed_chars)
     
-    # Split into words and trim to 100 words
-    words = clean_raw.split()
-    story = " ".join(words[:100])
+    # Clean the raw story using our custom function
+    story = clean_generated_story(raw)
     
     st.write("**📖 Your funny story: 📖**")
     st.write(story)
 
-    # Generate audio from cleaned story
+    # Generate audio for the story
     chunks = textwrap.wrap(story, width=200)
     audio = np.concatenate([tts(chunk)["audio"].squeeze() for chunk in chunks])
 
-    # Save audio to temporary file
+    # Save audio to a temporary file
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
         sf.write(temp_file.name, audio, tts.model.config.sampling_rate)
         temp_file_path = temp_file.name
 
     return caption, story, temp_file_path
 
-# Streamlit UI
+# Streamlit UI section
 st.title("✨ Magic Story Maker ✨")
 st.markdown("Upload a picture to make a funny story and hear it too! 📸")
 
 uploaded_image = st.file_uploader("Choose your picture", type=["jpg", "jpeg", "png"])
 
-# Streamlit UI (modified image display section)
 if uploaded_image is None:
     st.image("https://example.com/placeholder_image.jpg", caption="Upload your picture here! 📷", use_container_width=True)
 else: