added more debugging code and fixed up summarizer's input

app.py

@@ -75,7 +75,6 @@ class RoleAgent:
             f"Input: {input_text}\n"
             f"Output:"
         )
-        print("[DEBUG] prompt:", prompt)
         encoding = self.tokenizer(prompt, return_tensors="pt")
         inputs = {k: v.to(self.model.device) for k, v in encoding.items()}
 
@@ -89,13 +88,19 @@ class RoleAgent:
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
         thinking = ""
-        print(response)
+        print(f"[RESPONSE]: {response}")
         answer = response
-        if all(tag in response for tag in ("THINKING:", "ANSWER:", "END")):
-            print("[FIX] response:", response)
-            block = response.split("THINKING:")[1].split("END")[0]
-            thinking = block.split("ANSWER:")[0].strip()
-            answer = block.split("ANSWER:")[1].strip()
+        if "Output:" in response:
+            # Split on the last occurrence of 'Output:' in case it's repeated
+            answer = response.rsplit("Output:", 1)[-1].strip()
+        else:
+            # Fallback: if thinking/answer/end tags exist, use previous logic
+            tags = ("THINKING:", "ANSWER:", "END")
+            if all(tag in response for tag in tags):
+                print("[FIX] tagged response detected:", response)
+                block = response.split("THINKING:", 1)[1].split("END", 1)[0]
+                thinking = block.split("ANSWER:", 1)[0].strip()
+                answer = block.split("ANSWER:", 1)[1].strip()
 
         return {"thinking": thinking, "output": answer}
 
@@ -122,6 +127,19 @@ treatment_agent = RoleAgent(
     model=model,
 )
 
+"""[DEBUG] prompt: Instruction: You are a clinical summarizer trained to extract structured vignettes from doctor–patient dialogues.
+Input: Doctor: What brings you in today?
+Patient: I am a male. I am 15. My knee hurts. What may be the issue with my knee?
+
+Previous Vignette: 
+Output:
+Instruction: You are a clinical summarizer trained to extract structured vignettes from doctor–patient dialogues.
+Input: Doctor: What brings you in today?
+Patient: I am a male. I am 15. My knee hurts. What may be the issue with my knee?
+
+Previous Vignette: 
+Output: The patient is a 15-year-old male presenting with knee pain."""
+
 
 # === Inference State ===
 conversation_history = []