Update app.py

app.py

@@ -59,22 +59,36 @@ def img2text(image):
     text = image_to_text(image)[0]["generated_text"]
     return text
 
-# Improved text-to-story function with natural ending
+# Improved text-to-story function with longer stories (approaching 100 words)
 def text2story(text):
     generator = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-    prompt = f"Write a short children's story based on this: {text}. The story should have a clear beginning, middle, and end. Keep it under 150 words. Once upon a time, "
     
-    # Generate a longer text to ensure we get a complete story
+    # Specifically ask for a longer story (150-200 words) to ensure we get at least 100
+    prompt = f"""Write a children's story based on this: {text}. 
+    The story should have a clear beginning, middle, and end.
+    Make the story approximately 150-200 words long with descriptive language.
+    Start with "Once upon a time, "
+    """
+    
+    # Generate a longer text with higher max_length to ensure we get a complete story
     story_result = generator(
         prompt,
-        max_length=300,
+        max_length=500,  # Increased max length
         num_return_sequences=1,
         temperature=0.7,
         do_sample=True
     )
    
     story_text = story_result[0]['generated_text']
-    story_text = story_text.replace(prompt, "Once upon a time, ")
+    
+    # Extract just the story part (after the prompt)
+    if "Once upon a time, " in story_text:
+        # Find the index of "Once upon a time" and extract from there
+        start_idx = story_text.find("Once upon a time, ")
+        story_text = story_text[start_idx:]
+    else:
+        # If we can't find the exact phrase, try to find the story after the prompt
+        story_text = story_text.replace(prompt, "Once upon a time, ")
     
     # Find natural ending points (end of sentences)
     periods = [i for i, char in enumerate(story_text) if char == '.']
@@ -84,22 +98,38 @@ def text2story(text):
     # Combine all ending punctuation and sort
     all_endings = sorted(periods + question_marks + exclamation_marks)
     
+    # Count words in the story so far
+    def count_words(text):
+        return len(text.split())
+    
     # If we have any sentence endings
     if all_endings:
-        # Get the index where the story should reasonably end (after at least 100 characters)
-        min_story_length = 100
-        suitable_endings = [i for i in all_endings if i >= min_story_length]
+        # Find endings that give us stories of approximately 100 words or more
+        target_word_count = 100
+        min_word_count = 80  # Allow slightly shorter stories that end naturally
+        
+        suitable_endings = []
+        for ending_idx in all_endings:
+            candidate_text = story_text[:ending_idx+1]
+            word_count = count_words(candidate_text)
+            if word_count >= min_word_count:
+                suitable_endings.append((ending_idx, word_count))
         
         if suitable_endings:
-            # Find an ending that completes a thought (not just the first sentence)
-            if len(suitable_endings) > 2:
-                # Use the third sentence ending or later for a more complete story
-                return story_text[:suitable_endings[2]+1]
-            else:
-                # If we don't have many sentences, use the last one we found
-                return story_text[:suitable_endings[-1]+1]
+            # Find the ending that gets us closest to our target word count
+            suitable_endings.sort(key=lambda x: abs(x[1] - target_word_count))
+            best_ending_idx = suitable_endings[0][0]
+            return story_text[:best_ending_idx+1]
+    
+    # If we couldn't find a good ending point, try to get at least 80-100 words
+    if len(all_endings) > 0:
+        for i in range(len(all_endings)-1, -1, -1):
+            ending_idx = all_endings[i]
+            candidate_text = story_text[:ending_idx+1]
+            if count_words(candidate_text) >= 80:
+                return candidate_text
     
-    # If no good ending is found, return as is
+    # If all else fails, return the story as is
     return story_text
 
 # Basic Streamlit interface
@@ -121,7 +151,10 @@ if uploaded_file is not None:
     # Text to Story
     with st.spinner("Creating story..."):
         story = text2story(caption)
-    st.write(f"Story: {story}")
+        # Display word count for transparency
+        word_count = len(story.split())
+        st.write(f"Story ({word_count} words):")
+        st.write(story)
     
     # Text to Audio
     with st.spinner("Generating audio..."):