Update app.py

app.py

@@ -8,10 +8,34 @@ def img2text(url):
 
 def text2story(text):
     story_generator = pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator", device_map="auto", return_full_text=False)
-    prompt = f"A short story under 100 English words about: {text}."
-    generated = story_generator(text, max_new_tokens=140, do_sample=True, temperature=0.7)
+    prompt = f"<BOS> {text}"
+    generated = story_generator(prompt, max_new_tokens=140, do_sample=True, temperature=0.7)
     story_text = generated[0]['generated_text']
-    return story_text
+
+    # Split into sentences
+    sentences = re.split(r'(?<=[.!?])\s+', story_text.strip())
+    
+    # Initialize variables
+    current_word_count = 0
+    final_sentences = []
+    
+    # Iterate through each sentence and accumulate until the word count is within 100
+    for sentence in sentences:
+        words = sentence.split()
+        word_count = len(words)
+        if current_word_count + word_count > 100:
+            break
+        final_sentences.append(sentence)
+        current_word_count += word_count
+    
+    # Join the final sentences to form the story
+    final_story = ' '.join(final_sentences)
+    
+    # Ensure it ends with a punctuation mark
+    if not final_story.endswith(('.', '!', '?')):
+        final_story += '.'
+    
+    return final_story
 
 def text2audio(story_text):
     audio_generator = pipeline("text-to-speech", model="ylacombe/vits_vctk_scottish_female")