Update explainer.py

explainer.py

@@ -1,135 +1,96 @@
 # explainer.py
 # ECG Condition Classifier - Clinical Explanation Generator
 
-# ── Keyword-based label matching ──────────────────────────────────────────────
-# Maps keywords found in model output labels to clinical explanations.
-# This works regardless of whether the model returns "NORM", "Normal",
-# "LABEL_0", or any other format.
-
-CONDITION_MAP = [
-    {
-        "keywords": ["norm", "normal", "sinus"],
-        "code": "NORM",
+CONDITION_MAP = {
+    "NORM": {
         "summary": "Normal Sinus Rhythm",
-        "explanation": (
-            "The ECG pattern is consistent with a normal sinus rhythm. "
-            "Regular P waves, QRS complexes, and T waves are present within "
-            "expected parameters. No significant cardiac abnormalities identified."
-        ),
         "urgency": "Routine",
-        "recommendation": "No immediate action indicated. Routine monitoring as appropriate."
+        "description": "The ECG shows a normal sinus rhythm with regular rate and rhythm, normal P waves, QRS complexes, and T waves.",
+        "recommendation": "No immediate cardiac intervention required. Continue routine follow-up as clinically indicated.",
     },
-    {
-        "keywords": ["afib", "atrial fib", "fibrillation", "af"],
-        "code": "AFIB",
+    "AFIB": {
         "summary": "Atrial Fibrillation",
-        "explanation": (
-            "The ECG pattern is consistent with Atrial Fibrillation. "
-            "Characteristic findings include an irregularly irregular rhythm, "
-            "absence of distinct P waves, and variable R-R intervals. "
-            "AFib is associated with increased risk of stroke and heart failure."
-        ),
         "urgency": "Elevated",
-        "recommendation": "Clinical evaluation recommended. Rate/rhythm management and stroke risk assessment advised."
+        "description": "The ECG pattern is consistent with Atrial Fibrillation. Characteristic findings include an irregularly irregular rhythm, absence of distinct P waves, and variable R-R intervals. AFib is associated with increased risk of stroke and heart failure.",
+        "recommendation": "Clinical evaluation recommended. Rate/rhythm management and stroke risk assessment advised.",
     },
-    {
-        "keywords": ["mi", "infarct", "stemi", "nstemi", "myocardial"],
-        "code": "MI",
+    "MI": {
         "summary": "Myocardial Infarction",
-        "explanation": (
-            "The ECG pattern is consistent with Myocardial Infarction. "
-            "Findings may include ST-segment elevation or depression and "
-            "pathological Q waves suggesting myocardial injury or ischemia."
-        ),
         "urgency": "URGENT",
-        "recommendation": "⚠️ IMMEDIATE clinical review strongly recommended. Activate emergency cardiac protocols if STEMI suspected."
+        "description": "The ECG pattern is consistent with Myocardial Infarction. Findings may include ST-segment elevation or depression and pathological Q waves suggesting myocardial injury or ischemia.",
+        "recommendation": "⚠️ IMMEDIATE clinical review strongly recommended. Activate emergency cardiac protocols if STEMI suspected.",
     },
-    {
-        "keywords": ["sttc", "st ", "t-wave", "t wave", "st/t", "ischemi"],
-        "code": "STTC",
+    "STTC": {
         "summary": "ST/T-Wave Change",
-        "explanation": (
-            "Non-specific ST-segment and/or T-wave changes detected. "
-            "These findings can represent myocardial ischemia, electrolyte "
-            "abnormalities, medication effects, or normal variants. "
-            "Clinical context is essential for interpretation."
-        ),
         "urgency": "Moderate",
-        "recommendation": "Clinical correlation advised. Serial ECGs and lab results may be warranted."
+        "description": "Non-specific ST-segment and/or T-wave changes detected. These findings can represent myocardial ischemia, electrolyte abnormalities, medication effects, or normal variants. Clinical context is essential for interpretation.",
+        "recommendation": "Clinical correlation advised. Serial ECGs and lab results may be warranted.",
     },
-    {
-        "keywords": ["cd", "conduction", "bundle", "block", "bbb", "lbbb", "rbbb"],
-        "code": "CD",
+    "CD": {
         "summary": "Conduction Disturbance",
-        "explanation": (
-            "The ECG shows signs of a conduction disturbance. "
-            "Findings may include widened QRS complex or abnormal morphology "
-            "consistent with bundle branch block or intraventricular conduction delay."
-        ),
         "urgency": "Moderate",
-        "recommendation": "Cardiology follow-up recommended. Echocardiogram may help determine significance."
+        "description": "The ECG shows features consistent with a conduction disturbance, such as bundle branch block or AV conduction delay. These findings alter the normal electrical pathway through the ventricles.",
+        "recommendation": "Cardiology consultation recommended to evaluate hemodynamic significance and underlying etiology.",
     },
-]
-
-DEFAULT = {
-    "code": "?",
-    "summary": "Unrecognized Pattern",
-    "explanation": (
-        "The model returned a label that could not be mapped to a known "
-        "cardiac condition. This may indicate an unusual ECG pattern or "
-        "a label category outside the standard NORM/AFIB/MI/STTC/CD set."
-    ),
-    "urgency": "Unknown",
-    "recommendation": "Manual review by a qualified cardiologist is recommended."
 }
 
-DISCLAIMER = (
-    "\n\n⚠️ DISCLAIMER: For educational and research purposes only. "
-    "Does not constitute medical advice. Always consult a qualified healthcare provider."
-)
+DISCLAIMER = "⚠️ DISCLAIMER: For educational and research purposes only. Does not constitute medical advice. Always consult a qualified healthcare provider."
 
 
 def match_condition(label):
-    """Matches a raw model label string to a known condition using keyword search."""
-    label_lower = label.lower()
-    for condition in CONDITION_MAP:
-        if any(kw in label_lower for kw in condition["keywords"]):
-            return condition
-    return DEFAULT
+    """Return condition info dict for a given label string."""
+    label_upper = (label or "").upper()
+    for key, val in CONDITION_MAP.items():
+        if key in label_upper:
+            return {**val, "key": key}
+    return {
+        "key": "UNKNOWN",
+        "summary": "Unrecognized Pattern",
+        "urgency": "Unknown",
+        "description": "The model returned a label that could not be mapped to a known cardiac condition.",
+        "recommendation": "Manual review by a qualified cardiologist is recommended.",
+    }
 
 
-def generate_explanation(label, confidence):
+def generate_explanation(label, confidence, full_result=None):
     """
-    Generates a clinical explanation for any label the model returns.
-
-    Args:
-        label: str — raw label from model (e.g. "Normal", "AFIB", "LABEL_2")
-        confidence: float — confidence percentage
-
-    Returns:
-        (explanation_text: str, urgency: str)
+    Generate the clinical explanation text block.
+    If full_result (from Claude API) is available, use its rich findings.
+    Otherwise fall back to template text.
     """
-    data = match_condition(label)
+    condition = match_condition(label)
+
+    if full_result:
+        # Use Claude's own clinical findings text
+        findings = full_result.get("findings", condition["description"])
+        recommendation = full_result.get("recommendation", condition["recommendation"])
+        urgency = full_result.get("urgency", condition["urgency"])
+        full_name = full_result.get("full_name", condition["summary"])
+    else:
+        findings = condition["description"]
+        recommendation = condition["recommendation"]
+        urgency = condition["urgency"]
+        full_name = condition["summary"]
 
-    explanation_text = (
-        f"Condition: {data['summary']}\n"
+    text = (
+        f"Condition: {full_name}\n"
         f"Raw Model Label: {label}\n"
         f"Confidence: {confidence:.1f}%\n"
-        f"Urgency: {data['urgency']}\n\n"
-        f"{data['explanation']}\n\n"
-        f"Recommendation: {data['recommendation']}"
+        f"Urgency: {urgency}\n\n"
+        f"{findings}\n\n"
+        f"Recommendation: {recommendation}\n\n"
         f"{DISCLAIMER}"
     )
 
-    return explanation_text, data["urgency"]
+    return text, urgency
 
 
 def format_top3(results):
-    """Formats top predictions into a readable breakdown."""
+    """Format the confidence breakdown for display."""
     lines = ["Top Predictions:"]
     for i, r in enumerate(results[:3], 1):
-        lbl = r["label"]
-        pct = r["score"] * 100
-        bar = "█" * int(pct / 5)
-        lines.append(f"  {i}. {lbl}: {pct:.1f}%  {bar}")
+        label = r["label"]
+        score = r["score"] * 100
+        bar = "█" * int(score / 5)
+        lines.append(f"  {i}. {label}: {score:.1f}% {bar}")
     return "\n".join(lines)