Update app.py

app.py

@@ -32,7 +32,6 @@ If the question is unrelated to Python, say:
 ### ASSISTANT:
 """
 
-# Chat function
 def chat(instruction):
     prompt = format_prompt(instruction)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -47,16 +46,15 @@ def chat(instruction):
             repetition_penalty=1.1
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Debugging optional:
-    # print("🧠 Full output:", repr(response))
-
-    # Split and return the assistant's answer
-    if "### Answer:" in response:
-        return response.split("### Answer:")[-1].strip()
+    # Extract only the final answer
+    if "### Answer:" in full_output:
+        return full_output.split("### Answer:")[-1].strip()
+    elif "### ASSISTANT:" in full_output:
+        return full_output.split("### ASSISTANT:")[-1].strip()
     else:
-        return response.strip()
+        return full_output.strip()
 
 # Streamlit UI
 st.set_page_config(page_title="🐍 Python Tutor Chatbot")