Update app.py

app.py

@@ -34,6 +34,49 @@ def text2story(text):
     story_text = story_result[0]['generated_text']
     story_text = story_text.replace(prompt, "Once upon a time, ")
     
+    # Find a natural ending point (end of sentence) before 100 words
+    words = story_text.split()
+    if len(words) > 100:
+        # Join the first 100 words
+        shortened_text = " ".join(words[:100])
+        
+        # Find the last complete sentence
+        last_period = shortened_text.rfind('.')
+        last_question = shortened_text.rfind('?')
+        last_exclamation = shortened_text.rfind('!')
+        
+        # Find the last sentence ending punctuation
+        last_end = max(last_period, last_question, last_exclamation)
+        
+        if last_end > 0:
+            # Truncate at the end of the last complete sentence
+            story_text = shortened_text[:last_end + 1]
+        else:
+            # If no sentence ending found, just use the shortened text
+            story_text = shortened_text
+    
+    return story_textdef text2story(text):
+    # Using a smaller text generation model
+    generator = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+    
+    # Create a prompt for the story generation
+    prompt = f"Write a fun children's story based on this: {text}. The story should be short and end naturally with a conclusion. Once upon a time, "
+    
+    # Generate the story
+    story_result = generator(
+        prompt,
+        max_length=250,  # Increased to allow for a complete story
+        num_return_sequences=1,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95,
+        do_sample=True
+    )
+   
+    # Extract the generated text
+    story_text = story_result[0]['generated_text']
+    story_text = story_text.replace(prompt, "Once upon a time, ")
+    
     # Find a natural ending point (end of sentence) before 100 words
     words = story_text.split()
     if len(words) > 100: