Update app.py

app.py

@@ -10,6 +10,8 @@ import matplotlib.pyplot as plt
 import matplotlib.patches as mpatches
 import io
 import base64
+import string
+import re
 import sys
 import csv
 import os
@@ -42,27 +44,92 @@ ner_pipe = pipeline(
 )
 
 
-# ── Custom SHAP bar-chart renderer ─────────────────────────────────────────────
-def render_shap_bar_chart(shap_values, class_idx: int = 1) -> str:
+# ── Severity rating renderer ──────────────────────────────────────────────────
+def build_severity_html(severe_prob: float) -> str:
     """
-    Builds a horizontal bar chart (red = pushes toward ADR, teal = pushes away)
-    that mirrors the style shown in the reference screenshot.
-    Returns an <img> HTML tag with an embedded base64 PNG.
+    Converts the Severe Reaction probability into a Low / Medium / High badge
+    with a colour-coded progress bar and plain-language description.
+
+    Thresholds (tunable):
+        < 0.35  → Low
+        0.35 – 0.65 → Medium
+        >= 0.65 → High
     """
-    # shap_values is a shap.Explanation object for a single sample
-    # .values shape: (n_tokens, n_classes)  or  (n_tokens,) when binary
-    values = shap_values.values          # (n_tokens, n_classes)
-    tokens = shap_values.data            # list/array of token strings
+    if severe_prob < 0.35:
+        level      = "Low"
+        bar_color  = "#2eaa5a"   # green
+        bg_color   = "#eafaf1"
+        border_col = "#2eaa5a"
+        description = (
+            "The clinical text shows limited indicators of a serious adverse drug reaction. "
+            "Routine monitoring is advised."
+        )
+    elif severe_prob < 0.65:
+        level      = "Medium"
+        bar_color  = "#e8a020"   # amber
+        bg_color   = "#fffbea"
+        border_col = "#e8a020"
+        description = (
+            "The clinical text contains some features associated with a significant adverse "
+            "drug reaction. Further clinical assessment is recommended."
+        )
+    else:
+        level      = "High"
+        bar_color  = "#cc1111"   # red
+        bg_color   = "#fff0f0"
+        border_col = "#cc1111"
+        description = (
+            "The clinical text shows strong indicators of a severe adverse drug reaction. "
+            "Prompt clinical review is strongly recommended."
+        )
+
+    pct = round(severe_prob * 100, 1)
+    bar_width = round(severe_prob * 100, 1)
+
+    html = (
+        "<div style='"
+        "background:" + bg_color + "; "
+        "border:1.5px solid " + border_col + "; "
+        "border-radius:10px; padding:18px 20px; font-family:system-ui, sans-serif;'>"
+
+        # Rating badge + percentage on same row
+        "<div style='display:flex; align-items:center; justify-content:space-between; "
+        "margin-bottom:12px;'>"
+        "<span style='font-size:1.5em; font-weight:800; color:" + bar_color + ";'>"
+        + level +
+        "</span>"
+        "<span style='font-size:1.0em; font-weight:600; color:#555;'>"
+        + str(pct) + "% severe probability"
+        "</span>"
+        "</div>"
+
+        # Progress bar track
+        "<div style='background:#e0e0e0; border-radius:6px; height:14px; "
+        "overflow:hidden; margin-bottom:14px;'>"
+        "<div style='background:" + bar_color + "; width:" + str(bar_width) + "%; "
+        "height:100%; border-radius:6px; "
+        "transition:width 0.4s ease;'></div>"
+        "</div>"
+
+        # Description
+        "<p style='margin:0; font-size:0.9em; color:#444; line-height:1.6;'>"
+        + description +
+        "</p>"
+        "</div>"
+    )
+    return html
+
+
+# ── Custom SHAP bar-chart renderer ─────────────────────────────────────────────
+def render_shap_bar_chart(shap_values, class_idx=1):
+    values = shap_values.values
+    tokens = shap_values.data
 
     if values.ndim == 2:
-        sv = values[:, class_idx]        # SHAP values for "Severe Reaction"
+        sv = values[:, class_idx]
     else:
         sv = values
 
-    # ── Filter out noise tokens ──────────────────────────────────────────────
-    import string
-    import re
-
     STOPWORDS = {
         "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for",
         "of", "with", "by", "from", "is", "was", "are", "were", "be", "been",
@@ -78,19 +145,14 @@ def render_shap_bar_chart(shap_values, class_idx: int = 1) -> str:
         t = tok.strip().lower()
         if not t:
             return False
-        # Pure punctuation or whitespace
         if all(c in string.punctuation + " \t\n" for c in t):
             return False
-        # Single characters (letters or digits alone)
         if len(t) <= 1:
             return False
-        # Common stopwords
         if t in STOPWORDS:
             return False
-        # Subword fragments starting with ## (BERT-style)
         if t.startswith("##"):
             return False
-        # Tokens that are only digits
         if re.fullmatch(r"\d+", t):
             return False
         return True
@@ -100,30 +162,26 @@ def render_shap_bar_chart(shap_values, class_idx: int = 1) -> str:
     sv_f  = sv[mask]
     tok_f = tok_arr[mask]
 
-    # ── Deduplicate: for repeated tokens keep the one with the largest |SHAP| ─
+    # Deduplicate: keep highest |SHAP| per unique token
     seen = {}
     for i, tok in enumerate(tok_f):
         key = tok.strip().lower()
         if key not in seen or abs(sv_f[i]) > abs(sv_f[seen[key]]):
             seen[key] = i
-    keep = sorted(seen.values())
+    keep  = sorted(seen.values())
     sv_f  = sv_f[keep]
     tok_f = tok_f[keep]
 
-    # ── Top-N by absolute magnitude, then sort for chart display ─────────────
-    TOP_N = 20
-    order = np.argsort(np.abs(sv_f))[::-1][:TOP_N]
-    sv_top  = sv_f[order]
-    tok_top = tok_f[order]
-
-    # Re-sort so chart reads largest positive at top, largest negative at bottom
+    TOP_N      = 20
+    order      = np.argsort(np.abs(sv_f))[::-1][:TOP_N]
+    sv_top     = sv_f[order]
+    tok_top    = tok_f[order]
     plot_order = np.argsort(sv_top)
-    sv_plot  = sv_top[plot_order]
-    tok_plot = tok_top[plot_order]
-
-    COLOR_POSITIVE = "#cc1111"   # bold red  — increases severe ADR probability
-    COLOR_NEGATIVE = "#1a6fcc"   # strong blue — decreases severe ADR probability
+    sv_plot    = sv_top[plot_order]
+    tok_plot   = tok_top[plot_order]
 
+    COLOR_POSITIVE = "#cc1111"
+    COLOR_NEGATIVE = "#1a6fcc"
     colors = [COLOR_POSITIVE if v > 0 else COLOR_NEGATIVE for v in sv_plot]
 
     fig_height = max(4, len(sv_plot) * 0.38)
@@ -131,11 +189,8 @@ def render_shap_bar_chart(shap_values, class_idx: int = 1) -> str:
     ax.set_facecolor("white")
 
     y_pos = np.arange(len(sv_plot))
-    bars = ax.barh(y_pos, sv_plot, color=colors, height=0.6, edgecolor="none")
-
-    # Zero line
+    ax.barh(y_pos, sv_plot, color=colors, height=0.6, edgecolor="none")
     ax.axvline(0, color="#333333", linewidth=0.9, zorder=3)
-
     ax.set_yticks(y_pos)
     ax.set_yticklabels(tok_plot, fontsize=10, color="#222222")
     ax.set_xlabel("SHAP Value — impact on ADR prediction", fontsize=10, color="#444444")
@@ -144,34 +199,83 @@ def render_shap_bar_chart(shap_values, class_idx: int = 1) -> str:
         fontsize=12, fontweight="bold", color="#222222", pad=12
     )
 
-    # Legend patches — colors match the bars exactly
-    red_patch  = mpatches.Patch(color=COLOR_POSITIVE,
-                                label="Increases severe ADR probability")
-    blue_patch = mpatches.Patch(color=COLOR_NEGATIVE,
-                                label="Decreases severe ADR probability")
-    ax.legend(handles=[red_patch, blue_patch], fontsize=9,
-              loc="lower right", framealpha=0.7)
+    red_patch  = mpatches.Patch(color=COLOR_POSITIVE, label="Increases severe ADR probability")
+    blue_patch = mpatches.Patch(color=COLOR_NEGATIVE, label="Decreases severe ADR probability")
+    ax.legend(handles=[red_patch, blue_patch], fontsize=9, loc="lower right", framealpha=0.7)
 
     ax.spines["top"].set_visible(False)
     ax.spines["right"].set_visible(False)
     ax.spines["left"].set_visible(False)
     ax.tick_params(axis="y", length=0)
     ax.tick_params(axis="x", colors="#555555")
-    ax.xaxis.label.set_color("#555555")
 
     plt.tight_layout()
-
     buf = io.BytesIO()
-    fig.savefig(buf, format="png", dpi=130, bbox_inches="tight",
-                facecolor="white")
+    fig.savefig(buf, format="png", dpi=130, bbox_inches="tight", facecolor="white")
     plt.close(fig)
     buf.seek(0)
     b64 = base64.b64encode(buf.read()).decode("utf-8")
     return (
-        f"<div style='background:white; padding:12px; border-radius:8px;'>"
-        f"<img src='data:image/png;base64,{b64}' "
-        f"style='width:100%; max-width:760px; display:block; margin:auto;' />"
-        f"</div>"
+        "<div style='background:white; padding:12px; border-radius:8px;'>"
+        "<img src='data:image/png;base64," + b64 + "' "
+        "style='width:100%; max-width:760px; display:block; margin:auto;' /></div>"
+    )
+
+
+# ── Severity rating widget ─────────────────────────────────────────────────────
+def build_severity_html(severe_prob):
+    if severe_prob >= 0.70:
+        rating        = "HIGH"
+        rating_color  = "#cc1111"
+        rating_bg     = "#fff0f0"
+        rating_border = "#e07070"
+        rating_desc   = "Strong indicators of a serious adverse drug reaction are present."
+        rating_icon   = "\u26a0\ufe0f"
+    elif severe_prob >= 0.40:
+        rating        = "MEDIUM"
+        rating_color  = "#b07000"
+        rating_bg     = "#fffbe6"
+        rating_border = "#d4a800"
+        rating_desc   = "Some indicators of an adverse drug reaction detected. Clinical review recommended."
+        rating_icon   = "\U0001f536"
+    else:
+        rating        = "LOW"
+        rating_color  = "#2a7a2a"
+        rating_bg     = "#f0fff0"
+        rating_border = "#5aaa5a"
+        rating_desc   = "Few or no indicators of a severe adverse drug reaction detected."
+        rating_icon   = "\u2705"
+
+    bar_pct = int(severe_prob * 100)
+
+    return (
+        "<div style='background:" + rating_bg + "; border:2px solid " + rating_border + "; "
+        "border-radius:10px; padding:16px 20px; font-family:system-ui, sans-serif;'>"
+
+        "<div style='display:flex; align-items:center; gap:12px; margin-bottom:10px;'>"
+        "<span style='font-size:1.8em;'>" + rating_icon + "</span>"
+
+        "<div>"
+        "<div style='font-size:0.75em; font-weight:700; letter-spacing:0.1em; "
+        "text-transform:uppercase; color:#666;'>ADR Severity Rating</div>"
+        "<div style='font-size:1.6em; font-weight:900; color:" + rating_color + "; "
+        "letter-spacing:0.04em;'>" + rating + "</div>"
+        "</div>"
+
+        "<div style='margin-left:auto; text-align:right;'>"
+        "<div style='font-size:0.72em; color:#888; font-weight:600;'>Severe reaction probability</div>"
+        "<div style='font-size:1.4em; font-weight:800; color:" + rating_color + ";'>"
+        + str(bar_pct) + "%</div>"
+        "</div>"
+        "</div>"
+
+        "<div style='background:#e0e0e0; border-radius:999px; height:10px; margin-bottom:10px;'>"
+        "<div style='background:" + rating_color + "; width:" + str(bar_pct) + "%; "
+        "height:10px; border-radius:999px;'></div>"
+        "</div>"
+
+        "<div style='font-size:0.88em; color:#555; margin-top:4px;'>" + rating_desc + "</div>"
+        "</div>"
     )
 
 
@@ -182,18 +286,16 @@ def adr_predict(x):
     output = model(**encoded_input)
     scores = torch.softmax(output.logits, dim=-1)[0].detach().cpu().numpy()
 
-    # ── SHAP (bar chart) ──────────────────────────────────────────────────────
+    # SHAP
     try:
         shap_values = explainer([text_input])
         shap_html = render_shap_bar_chart(shap_values[0], class_idx=1)
     except Exception as e:
-        shap_html = f"<p style='color:red;'>SHAP explanation error: {e}</p>"
+        shap_html = "<p style='color:red;'>SHAP explanation error: " + str(e) + "</p>"
 
-    # ── NER ───────────────────────────────────────────────────────────────────
+    # NER
     try:
         res = ner_pipe(text_input)
-
-        # Richer config: background, border accent, and human-readable display name
         entity_config = {
             "Severity":             {"bg": "#ffe0de", "border": "#e07070", "label": "Severity"},
             "Sign_symptom":         {"bg": "#d4f5d4", "border": "#5aaa5a", "label": "Symptom"},
@@ -205,30 +307,27 @@ def adr_predict(x):
         }
         default_cfg = {"bg": "#f0f0f0", "border": "#aaaaaa", "label": "Other"}
 
-        # Legend — only show entity types actually found in this result
-        seen_groups = list(dict.fromkeys(e["entity_group"] for e in sorted(res, key=lambda e: e["start"])))
+        seen_groups = list(dict.fromkeys(
+            e["entity_group"] for e in sorted(res, key=lambda e: e["start"])
+        ))
         legend_html = (
             "<div style='display:flex; flex-wrap:wrap; gap:10px; "
-            "margin-bottom:16px; padding-bottom:12px; "
-            "border-bottom:1px solid #e0e0e0;'>"
+            "margin-bottom:16px; padding-bottom:12px; border-bottom:1px solid #e0e0e0;'>"
         )
         for grp in seen_groups:
-            cfg = entity_config.get(grp, default_cfg)
+            cfg    = entity_config.get(grp, default_cfg)
             bg     = cfg["bg"]
             border = cfg["border"]
             lbl    = cfg["label"]
             legend_html += (
-                f"<span style='display:inline-flex; align-items:center; gap:6px; "
-                f"font-size:0.8em; font-weight:600; color:#444; "
-                f"font-family:system-ui, sans-serif;'>"
-                f"<span style='display:inline-block; width:13px; height:13px; "
-                f"border-radius:3px; background:{bg}; "
-                f"border:2px solid {border};'></span>"
-                f"{lbl}</span>"
+                "<span style='display:inline-flex; align-items:center; gap:6px; "
+                "font-size:0.8em; font-weight:600; color:#444; font-family:system-ui, sans-serif;'>"
+                "<span style='display:inline-block; width:13px; height:13px; border-radius:3px; "
+                "background:" + bg + "; border:2px solid " + border + ";'></span>"
+                + lbl + "</span>"
             )
         legend_html += "</div>"
 
-        # Annotated text — each entity as a pill with a small label badge above it
         text_html = (
             "<div style='line-height:3.4; font-size:1.05em; color:#111; "
             "font-family:Georgia, serif; letter-spacing:0.01em;'>"
@@ -243,36 +342,35 @@ def adr_predict(x):
             border = cfg["border"]
             lbl    = cfg["label"]
 
-            # Plain text before this entity
-            text_html += f"<span style='color:#111;'>{text_input[prev_end:start]}</span>"
-
-            # Entity pill: small uppercase label above, coloured word below
+            text_html += "<span style='color:#111;'>" + text_input[prev_end:start] + "</span>"
             text_html += (
-                f"<span style='display:inline-block; position:relative; "
-                f"vertical-align:middle; margin:0 2px; text-align:center;'>"
-                f"<span style='display:block; font-size:0.6em; font-weight:800; "
-                f"letter-spacing:0.08em; text-transform:uppercase; "
-                f"color:{border}; font-family:system-ui, sans-serif; "
-                f"line-height:1.1; margin-bottom:2px;'>{lbl}</span>"
-                f"<span style='background:{bg}; border:1.5px solid {border}; "
-                f"color:#111; padding:3px 8px; border-radius:6px; "
-                f"font-weight:600; white-space:nowrap;'>{word}</span>"
-                f"</span>"
+                "<span style='display:inline-block; position:relative; "
+                "vertical-align:middle; margin:0 2px; text-align:center;'>"
+                "<span style='display:block; font-size:0.6em; font-weight:800; "
+                "letter-spacing:0.08em; text-transform:uppercase; color:" + border + "; "
+                "font-family:system-ui, sans-serif; line-height:1.1; margin-bottom:2px;'>"
+                + lbl + "</span>"
+                "<span style='background:" + bg + "; border:1.5px solid " + border + "; "
+                "color:#111; padding:3px 8px; border-radius:6px; "
+                "font-weight:600; white-space:nowrap;'>" + word + "</span>"
+                "</span>"
             )
             prev_end = end
 
-        text_html += f"<span style='color:#111;'>{text_input[prev_end:]}</span></div>"
+        text_html += "<span style='color:#111;'>" + text_input[prev_end:] + "</span></div>"
         htext = legend_html + text_html
 
     except Exception as ex:
-        htext = f"<p style='color:#c00;'>NER processing error: {ex}</p>"
+        htext = "<p style='color:#c00;'>NER processing error: " + str(ex) + "</p>"
 
     label_output = {
         "Severe Reaction":     float(scores[1]),
         "Non-severe Reaction": float(scores[0]),
     }
 
-    return label_output, shap_html, htext
+    severity_html = build_severity_html(float(scores[1]))
+
+    return label_output, severity_html, shap_html, htext
 
 
 # ── UI ─────────────────────────────────────────────────────────────────────────
@@ -330,6 +428,9 @@ with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as
             gr.Markdown("### Classification")
             label = gr.Label(label="Severity Probability")
 
+            gr.Markdown("### Severity Rating")
+            severity_out = gr.HTML(label="Severity Rating", elem_classes="output-box")
+
             gr.Markdown("### Medical Entities")
             htext_out = gr.HTML(label="NER Mapping", elem_classes="output-box")
 
@@ -345,7 +446,7 @@ with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as
     submit_btn.click(
         fn=adr_predict,
         inputs=[prob1],
-        outputs=[label, shap_out, htext_out],
+        outputs=[label, severity_out, shap_out, htext_out],
     )
 
 demo.launch()