Update app.py

app.py

@@ -35,39 +35,40 @@ def img2text(image_bytes, _models):
 
 @st.cache_data
 def text2story(caption, _models):
-    """Generate a short story from image caption.
+    """Generate a short story from an image caption.
     
     Args:
-        caption: Caption describing the image
-        _models: Dictionary containing loaded models
+        caption (str): Caption describing the image.
+        _models (dict): Dictionary containing loaded models.
         
     Returns:
-        A generated story that expands on the image caption
+        str: A generated story that expands on the image caption.
     """
     story_generator = _models["story_generator"]
-    
-    # Format prompt to ensure the story expands on the image caption
-    prompt = f"""<|system|>
-You are a story generating assistant.
-<|user|>
-Please generate a story within 100 words based on the image caption: {caption}
-<|assistant|>"""
-    
-    # Generate story with parameters tuned for brevity and coherence
+
+    # Updated prompt includes explicit instructions to incorporate details from the caption
+    prompt = (
+        "<|system|>\n"
+        "You are a creative story generating assistant that always incorporates key details "
+        "from the provided image caption into your story.\n"
+        "<|user|>\n"
+        f"Please generate a short story within 100 words that clearly utilizes the following image caption and expands upon it thoroughly:\n\"{caption}\"\n"
+        "<|assistant|>"
+    )
+
+    # Generate story with parameters tuned for creativity and coherence
     response = story_generator(
         prompt,
-        max_new_tokens=120,  # Allow enough tokens for a complete story
+        max_new_tokens=120,  # Enough tokens for a complete story
         do_sample=True,
         temperature=0.7,     # Balanced creativity
-        top_p=0.9,           # Focus on more likely tokens
-        repetition_penalty=1.2,  # Avoid repetitive patterns
+        top_p=0.9,           # Focus on more likely tokens for coherence
+        repetition_penalty=1.2,  # Reduce repetitive patterns
         eos_token_id=story_generator.tokenizer.eos_token_id
     )
     
-    # Extract just the generated story text
+    # Extract the generated story text by separating the assistant's reply
     raw_story = response[0]['generated_text']
-    
-    # Parse out just the assistant's response from the conversation format
     story_text = raw_story.split("<|assistant|>")[-1].strip()
     
     return story_text