File size: 2,110 Bytes
633ba95 0f5d296 eb1f007 633ba95 0f5d296 a392df0 633ba95 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from streamlit_agraph import Node, Edge, Config
import re, itertools
from typing import List, Tuple
_GREEN = "#00b894"
_ORANGE = "#d35400"
_BLUE = "#0984e3"
def _safe(node_like):
"""Return empty dict if node_like is an Exception."""
return node_like if isinstance(node_like, dict) else {}
def build_agraph(
papers: List[dict],
umls: List[dict],
drug_safety: List[dict],
) -> Tuple[List[Node], List[Edge], Config]:
nodes, edges = [], []
# UMLS concept nodes
for c in filter(bool, map(_safe, umls)):
cui, name = c.get("cui"), c.get("name", "")
if cui and name:
nid = f"cui:{cui}"
nodes.append(Node(id=nid, label=name, color=_GREEN, size=25))
# Drug nodes
def _drug_name(d: dict) -> str:
return (
d.get("drug_name")
or d.get("patient", {}).get("drug", "")
or d.get("medicinalproduct", "")
)
for idx, rec in enumerate(itertools.chain.from_iterable(
[r if isinstance(r, list) else [r] for r in drug_safety])):
dn = _drug_name(rec) or f"drug_{idx}"
did = f"drug_{idx}"
nodes.append(Node(id=did, label=dn, color=_ORANGE, size=25))
# Paper nodes
for i, p in enumerate(papers, 1):
pid = f"paper_{i}"
nodes.append(Node(id=pid, label=f"P{i}", tooltip=p["title"], color=_BLUE, size=15))
txt = f"{p['title']} {p['summary']}".lower()
# link ↔ concepts
for c in filter(bool, map(_safe, umls)):
if (name := c.get("name", "")).lower() in txt and c.get("cui"):
edges.append(Edge(source=pid, target=f"cui:{c['cui']}", label="mentions"))
# link ↔ drugs
for n in nodes:
if n.id.startswith("drug_") and n.label.lower() in txt:
edges.append(Edge(source=pid, target=n.id, label="mentions"))
cfg = Config(
width="100%", height="600px", directed=False,
node={"labelProperty": "label"},
nodeHighlightBehavior=True, highlightColor="#f1c40f",
collapsible=True,
)
return nodes, edges, cfg
|