Update app.py

app.py

@@ -59,36 +59,22 @@ def img2text(image):
     text = image_to_text(image)[0]["generated_text"]
     return text
 
-# Improved text-to-story function with longer stories (approaching 100 words)
+# Improved text-to-story function with natural ending
 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, "
     
-    # 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
+    # Generate a longer text to ensure we get a complete story
     story_result = generator(
         prompt,
-        max_length=500,  # Increased max length
+        max_length=300,
         num_return_sequences=1,
         temperature=0.7,
         do_sample=True
     )
    
     story_text = story_result[0]['generated_text']
-    
-    # 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, ")
+    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 == '.']
@@ -98,38 +84,22 @@ 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:
-        # 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))
+        # 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]
         
         if suitable_endings:
-            # 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
+            # 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]
     
-    # If all else fails, return the story as is
+    # If no good ending is found, return as is
     return story_text
 
 # Basic Streamlit interface
@@ -151,10 +121,7 @@ if uploaded_file is not None:
     # Text to Story
     with st.spinner("Creating story..."):
         story = text2story(caption)
-        # Display word count for transparency
-        word_count = len(story.split())
-        st.write(f"Story ({word_count} words):")
-        st.write(story)
+    st.write(f"Story: {story}")
     
     # Text to Audio
     with st.spinner("Generating audio..."):