Update app.py

app.py

@@ -3,6 +3,7 @@ import streamlit as st
 from transformers import pipeline
 from gtts import gTTS  # Using gTTS for text-to-speech
 import os
+import re  # For removing unwanted words
 
 # function part
 # img2text
@@ -15,13 +16,15 @@ def img2text(url):
 
 # text2story
 def text2story(text):
-    # Use a text-generation model to create a fun and realistic story
-    story_generator = pipeline("text-generation", model="gpt2")
-    # Add a prompt to guide the story generation (avoid magical elements)
-    prompt = f"Write a fun and realistic story for kids based on this: {text}. Keep it simple and under 95 words."
+    # Use facebook/bart-large-cnn for text generation
+    story_generator = pipeline("text-generation", model="facebook/bart-large-cnn")
+    # Add a more explicit prompt to guide the story generation
+    prompt = f"Write a fun and realistic story for kids based on this: {text}. The story should be under 95 words and suitable for children aged 3-10. Avoid using the word 'illustration'."
     story = story_generator(prompt, max_length=95, num_return_sequences=1)[0]["generated_text"]
     # Remove the prompt from the generated story
-    story = story.replace(prompt, "").strip()  # Clean up the output
+    story = story.replace(prompt, "").strip()
+    # Remove the word "illustration" (case-insensitive) using regex
+    story = re.sub(r"\billustration\b", "", story, flags=re.IGNORECASE)
     return story[:95]  # Limit to 95 words
 
 # text2audio