Update app.py

app.py

@@ -50,8 +50,17 @@ def ask():
 )
 
 
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = text.split("Assistant:")[-1].strip()
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+reply = decoded.split("Response:", 1)[-1]
+
+STOP_WORDS = ["User:", "AI:", "Assistant:", "Question:", "Answer:"]
+for s in STOP_WORDS:
+    if s in reply:
+        reply = reply.split(s)[0]
+
+reply = reply.strip()
+
 
     return jsonify({"reply": reply})