Update app.py

app.py

@@ -36,13 +36,21 @@ generator = pipeline(
 
 def generate_reply(prompt):
     full_prompt = (
-        spartan_intro.strip() + "\n\n"
-        "The Spartan Mentor awaits your question.\n\n"
-        "User: " + prompt.strip() + "\n"
+        spartan_intro.strip() +
+        "\n\n[Conversation begins]\n" +
+        "User: " + prompt.strip() + "\n" +
         "Spartan Mentor:"
     )
     result = generator(full_prompt, max_length=300, num_return_sequences=1)[0]["generated_text"]
-    return result.split("Spartan Mentor:")[-1].strip()
+    
+    # After "Spartan Mentor:", capture only what bot says
+    if "Spartan Mentor:" in result:
+        result = result.split("Spartan Mentor:")[-1]
+    
+    if "User:" in result:
+        result = result.split("User:")[0]
+    
+    return result.strip()
 
 iface = gr.Interface(
     fn=generate_reply,