Update app.py

app.py

@@ -141,35 +141,74 @@ def adr_predict(x):
     # ── NER ───────────────────────────────────────────────────────────────────
     try:
         res = ner_pipe(text_input)
-        entity_colors = {
-            "Severity":             "#ffcccb",
-            "Sign_symptom":         "#bcf5bc",
-            "Medication":           "#cfe2f3",
-            "Age":                  "#fff2cc",
-            "Sex":                  "#fff2cc",
-            "Diagnostic_procedure": "#eeeeee",
-            "Biological_structure": "#d9d9d9",
+
+        # 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"},
+            "Medication":           {"bg": "#d0e8ff", "border": "#4a90d9", "label": "Medication"},
+            "Age":                  {"bg": "#fff3cc", "border": "#d4a800", "label": "Age"},
+            "Sex":                  {"bg": "#ffe8fb", "border": "#c070b0", "label": "Sex"},
+            "Diagnostic_procedure": {"bg": "#e8e8e8", "border": "#888888", "label": "Diagnostic"},
+            "Biological_structure": {"bg": "#ddeeff", "border": "#6699cc", "label": "Body Part"},
         }
+        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"])))
+        legend_html = (
+            "<div style='display:flex; flex-wrap:wrap; gap:10px; "
+            "margin-bottom:16px; padding-bottom:12px; "
+            "border-bottom:1px solid #e0e0e0;'>"
+        )
+        for grp in seen_groups:
+            cfg = entity_config.get(grp, default_cfg)
+            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:{cfg[\"bg\"]}; "
+                f"border:2px solid {cfg[\"border\"]};'></span>"
+                f"{cfg['label']}</span>"
+            )
+        legend_html += "</div>"
 
-        htext = "<div style='line-height:2.0; font-size:1.1em; color:black;'>"
+        # 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;'>"
+        )
         prev_end = 0
-        res = sorted(res, key=lambda e: e["start"])
-        for entity in res:
+        res_sorted = sorted(res, key=lambda e: e["start"])
+        for entity in res_sorted:
             start, end = entity["start"], entity["end"]
-            word  = text_input[start:end]
-            color = entity_colors.get(entity["entity_group"], "#f3f3f3")
-            htext += text_input[prev_end:start]
-            htext += (
-                f"<mark style='background-color:{color}; color:black; "
-                f"padding:2px 4px; border-radius:4px; font-weight:500;'>"
-                f"{word} "
-                f"<small style='opacity:0.7;'>[{entity['entity_group']}]</small>"
-                f"</mark>"
+            word = text_input[start:end]
+            cfg  = entity_config.get(entity["entity_group"], default_cfg)
+
+            # 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 += (
+                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:{cfg[\"border\"]}; font-family:system-ui, sans-serif; "
+                f"line-height:1.1; margin-bottom:2px;'>{cfg['label']}</span>"
+                f"<span style='background:{cfg[\"bg\"]}; border:1.5px solid {cfg[\"border\"]}; "
+                f"color:#111; padding:3px 8px; border-radius:6px; "
+                f"font-weight:600; white-space:nowrap;'>{word}</span>"
+                f"</span>"
             )
             prev_end = end
-        htext += text_input[prev_end:] + "</div>"
-    except Exception:
-        htext = "<p style='color:black;'>NER processing error.</p>"
+
+        text_html += f"<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>"
 
     label_output = {
         "Severe Reaction":     float(scores[1]),
@@ -217,10 +256,14 @@ with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as
             gr.Markdown("### Examples")
             gr.Examples(
                 examples=[
-                    ["A 35 year-old male had severe headache after taking Aspirin. "
-                     "The lab results were normal."],
-                    ["A 35 year-old female had minor pain in upper abdomen after "
-                     "taking Acetaminophen."],
+                    ["A 42 year-old male developed a severe migraine and elevated blood pressure "
+                     "shortly after taking Aspirin. He was admitted for observation."],
+                    ["A 28 year-old female reported mild nausea and minor discomfort in the upper "
+                     "abdomen after taking Acetaminophen 500mg. Symptoms resolved within two hours."],
+                    ["A 67 year-old male with a history of renal impairment experienced acute kidney "
+                     "injury and oliguria following treatment with Ibuprofen for three days."],
+                    ["A 54 year-old female noted slight dizziness and dry mouth after her first dose "
+                     "of Metformin. No further symptoms were reported at the follow-up visit."],
                 ],
                 inputs=[prob1],
             )
@@ -248,4 +291,4 @@ with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as
         outputs=[label, shap_out, htext_out],
     )
 
-demo.launch()
+demo.launch()