Update app.py

app.py

@@ -7,6 +7,7 @@ st.set_page_config(
 )
 
 # Other imports AFTER Streamlit config
+import re
 import time
 import tempfile
 from PIL import Image
@@ -33,11 +34,11 @@ def load_models():
         device_map="auto",
         trust_remote_code=True,
         torch_dtype="auto",
-        max_new_tokens=150,
+        max_new_tokens=250,
         temperature=0.7,
         top_p=0.85,
         repetition_penalty=1.15,
-        eos_token_id=151645  # Qwen3's specific EOS token
+        eos_token_id=151645
     )
     
     return captioner, storyteller
@@ -69,10 +70,10 @@ if uploaded_image:
     # Create story prompt
     story_prompt = (
         f"<|im_start|>system\n"
-        f"You are a children's book author. Create a 50-100 word story based on this image description: {image_caption}\n"
+        f"You are a children's book author. Create a 100-150 word story based on: {image_caption}\n"
         "Use simple language, friendly characters, and a positive lesson.<|im_end|>\n"
         f"<|im_start|>user\n"
-        f"Write a short, child-friendly story with a clear beginning, middle, and end.<|im_end|>\n"
+        f"Write a child-friendly story with a clear beginning, middle, and end.<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
 
@@ -82,25 +83,38 @@ if uploaded_image:
         story_result = story_pipe(
             story_prompt,
             do_sample=True,
-            num_return_sequences=1
+            num_return_sequences=1,
+            pad_token_id=151645
         )
         generation_time = time.time() - start_time
-        st.text(f"⏱ Generation time: {generation_time:.1f}s")
 
     # Process output
     raw_story = story_result[0]['generated_text']
+    
+    # Clean up story text
     clean_story = raw_story.split("<|im_start|>assistant\n")[-1]
+    clean_story = clean_story.split("<|im_start|>")[0]  # Remove any new turns
     clean_story = clean_story.replace("<|im_end|>", "").strip()
+    
+    # Remove assistant mentions using regex
+    clean_story = re.sub(
+        r'^(assistant[:>]?\s*)+', 
+        '', 
+        clean_story, 
+        flags=re.IGNORECASE
+    ).strip()
 
-    # Ensure proper story formatting
+    # Format story punctuation
     final_story = []
     for sentence in clean_story.split(". "):
-        if not sentence: continue
+        sentence = sentence.strip()
+        if not sentence:
+            continue
         if not sentence.endswith('.'):
             sentence += '.'
         final_story.append(sentence[0].upper() + sentence[1:])
     
-    final_story = " ".join(final_story).replace("..", ".")[:600]  # Character limit safeguard
+    final_story = " ".join(final_story).replace("..", ".")[:800]
 
     # Display story
     st.subheader("✨ Your Magical Story")