fixed output formatting for model

app.py

@@ -65,7 +65,6 @@ print(
 # === Role Agent with instruction/input/output format ===
 class RoleAgent:
     def __init__(self, role_instruction, tokenizer, model):
-
         self.tokenizer = tokenizer
         self.model = model
         self.role_instruction = role_instruction
@@ -81,36 +80,97 @@ class RoleAgent:
 
         outputs = self.model.generate(
             **inputs,
-            max_new_tokens=256,
+            max_new_tokens=128,
             do_sample=True,
             temperature=0.7,
             pad_token_id=self.tokenizer.eos_token_id,
         )
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        thinking = ""
-        answer = response
-        if "Output:" in response:
-            # Split on the last occurrence of 'Output:' in case it's repeated
-            answer = response.rsplit("Output:", 1)[-1].strip()
+        # Extract only the new generated content after the prompt
+        prompt_length = len(prompt)
+        if len(response) > prompt_length:
+            generated_text = response[prompt_length:].strip()
+        else:
+            generated_text = response.strip()
+
+        # Clean up the response - remove any repeated instruction/input/output patterns
+        lines = generated_text.split("\n")
+        clean_lines = []
+
+        for line in lines:
+            line = line.strip()
+            # Skip lines that look like instruction formatting
+            if (
+                line.startswith("instruction:")
+                or line.startswith("input:")
+                or line.startswith("output:")
+                or line == ""
+            ):
+                continue
+            clean_lines.append(line)
+
+        # Join the clean lines and take the first substantial response
+        if clean_lines:
+            answer = clean_lines[0]
+            # If there are multiple clean lines, take the first one that's substantial
+            for line in clean_lines:
+                if len(line) > 20:  # Arbitrary threshold for substantial content
+                    answer = line
+                    break
         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()
+            # Fallback: try to extract after "output:" if present
+            if "output:" in generated_text.lower():
+                parts = generated_text.lower().split("output:")
+                if len(parts) > 1:
+                    answer = parts[-1].strip()
+                else:
+                    answer = generated_text
+            else:
+                answer = generated_text
+
+        # Additional cleanup - remove any remaining instruction artifacts
+        answer = (
+            answer.replace("instruction:", "")
+            .replace("input:", "")
+            .replace("output:", "")
+            .strip()
+        )
+
+        # If answer is still messy, try to extract the actual medical content
+        if "patient" in answer.lower() and len(answer) > 100:
+            # Look for sentences that contain medical information
+            sentences = answer.split(".")
+            medical_sentences = []
+            for sentence in sentences:
+                sentence = sentence.strip()
+                if len(sentence) > 10 and any(
+                    word in sentence.lower()
+                    for word in [
+                        "patient",
+                        "pain",
+                        "symptom",
+                        "diagnosis",
+                        "treatment",
+                        "knee",
+                        "reports",
+                        "experiencing",
+                    ]
+                ):
+                    medical_sentences.append(sentence)
+
+            if medical_sentences:
+                answer = ". ".join(
+                    medical_sentences[:2]
+                )  # Take first 2 medical sentences
+                if not answer.endswith("."):
+                    answer += "."
+
         print(
-            "[THINKING ANSWER SPLIT] thinking/answer split:",
-            response,
-            "→",
-            "[THINKING] thinking:",
-            thinking,
-            "[ANSWER] answer:",
-            answer,
+            f"[CLEANED RESPONSE] Original length: {len(response)}, Cleaned: '{answer}'"
         )
 
+        thinking = ""  # For now, we'll focus on getting clean answers
         return {"thinking": thinking, "output": answer}