fixed treatment input code

app.py

@@ -65,15 +65,22 @@ print(
 # === Role Agent with instruction/input/output format ===
 class RoleAgent:
     def __init__(self, role_instruction, tokenizer, model):
+        format_instr = (
+            "When you respond, ALWAYS use exactly two sections:\n"
+            "THINKING:  <your detailed reasoning>\n"
+            "ANSWER:    <your concise final output>\n"
+            "Do NOT output anything else.\n\n"
+        )
+
         self.tokenizer = tokenizer
         self.model = model
-        self.role_instruction = role_instruction
+        self.role_instruction = role_instruction + format_instr
 
     def act(self, input_text):
         prompt = (
-            f"Instruction: {self.role_instruction}\n"
-            f"Input: {input_text}\n"
-            f"Output:"
+            f"instruction: {self.role_instruction}\n"
+            f"input: {input_text}\n"
+            f"output:"
         )
         encoding = self.tokenizer(prompt, return_tensors="pt")
         inputs = {k: v.to(self.model.device) for k, v in encoding.items()}
@@ -172,20 +179,14 @@ def simulate_interaction(user_input, iterations=1):
         diagnosis = diag_out["output"]
 
         # Question
-        q_in = (
-            f"Vignette: {summary}\n"
-            f"Current Estimated Diagnosis:\n"
-            f"ANSWER: {diagnosis}"
-        )
+        q_in = f"Vignette: {summary}\n" f"Current Estimated Diagnosis:\n" f"{diagnosis}"
         q_out = questioner.act(q_in)
         history.append(f"Doctor: {q_out['output']}")
         # Append actual user input to keep the conversation realistic
         history.append(f"Patient: {user_input}")
 
         # Treatment
-        treatment_out = treatment_agent.act(
-            f"Diagnosis: {diagnosis}\nVignette: {summary}"
-        )
+        treatment_out = treatment_agent.act(f"Diagnosis: {diagnosis}\n")
 
     return {
         "summary": sum_out,