Update app.py

app.py

@@ -31,6 +31,21 @@ generator = pipeline(
     tokenizer=tokenizer,
     return_full_text=False
 )
+def chat_prompt(user_text: str) -> str:
+    messages = [
+        {"role": "system", "content": "You are a helpful storyteller that writes engaging prose."},
+        {"role": "user", "content": (user_text or "").strip()}
+    ]
+    # Use chat template if available, otherwise fall back to a simple format
+    if hasattr(tokenizer, "apply_chat_template"):
+        return tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    # Fallback prompt format
+    return "System: You are a helpful storyteller that writes engaging prose.\nUser: " + (user_text or "").strip() + "\nAssistant:"
+
 
 # Function to generate stories
 def generate_story(prompt, max_tokens=300, temperature=0.8):