app.py

import re, json
import torch
import gradio as gr
from transformers import BertTokenizerFast, BertForTokenClassification

# === ตั้งค่าโมเดลจาก Hub ===
MODEL_ID = "Donlagon007/htn-ner-v1"

# โหลดโมเดล/โทเคนไนเซอร์ (CPU เป็นค่าเริ่มต้นใน Spaces)
tokenizer = BertTokenizerFast.from_pretrained(MODEL_ID)
model = BertForTokenClassification.from_pretrained(MODEL_ID)
model.eval()
id2label = model.config.id2label

# ---------- Utils ----------
def decode_bio_to_spans(labels):
    spans, cur_type, s = [], None, None
    for i, lab in enumerate(labels):
        if lab == "O" or lab is None:
            if cur_type is not None:
                spans.append((cur_type, s, i-1))
                cur_type, s = None, None
            continue
        tag, typ = lab.split("-", 1)
        if tag == "B":
            if cur_type is not None:
                spans.append((cur_type, s, i-1))
            cur_type, s = typ, i
        elif tag == "I":
            if cur_type != typ:
                if cur_type is not None:
                    spans.append((cur_type, s, i-1))
                cur_type, s = typ, i
    if cur_type is not None:
        spans.append((cur_type, s, len(labels)-1))
    return spans

def ner_predict_with_tokens(text, max_length=256):
    enc = tokenizer(
        text,
        return_offsets_mapping=True,
        return_tensors="pt",
        truncation=True, max_length=max_length
    )
    with torch.no_grad():
        out = model(
            input_ids=enc["input_ids"],
            attention_mask=enc["attention_mask"]
        )
        pred_ids = out.logits.argmax(-1).squeeze(0).tolist()

    offsets = enc["offset_mapping"].squeeze(0).tolist()
    input_ids = enc["input_ids"].squeeze(0).tolist()

    tokens_info, kept_labels, kept_offsets = [], [], []
    # ตัด [CLS]/[SEP]/padding โดยเช็ค offset (0,0)
    for lid, (st, ed), tid in zip(pred_ids, offsets, input_ids):
        if st == ed == 0:
            continue
        tok = tokenizer.convert_ids_to_tokens([tid])[0]
        lab = id2label[lid]
        tokens_info.append({"token": tok, "label": lab, "start": st, "end": ed})
        kept_labels.append(lab)
        kept_offsets.append((st, ed))

    spans_tok = decode_bio_to_spans(kept_labels)
    entities = []
    for typ, s_tok, e_tok in spans_tok:
        cs, ce = kept_offsets[s_tok][0], kept_offsets[e_tok][1]
        entities.append({"type": typ, "text": text[cs:ce], "start": cs, "end": ce})
    return tokens_info, entities

def cn_num(s):
    m = re.search(r"(\d+(?:\.\d+)?)", s.replace(" ", ""))
    return float(m.group(1)) if m else None

def parse_bp(value_text):
    m = re.search(r"(\d{2,3})\s*/\s*(\d{2,3})", value_text.replace(" ", ""))
    if m:
        return int(m.group(1)), int(m.group(2))
    return None, None

THRESHOLDS = {
    "空腹血糖": {"unit":"mg/dL", "abnormal": lambda v: v is not None and v >= 126},
    "HbA1c":   {"unit":"%",     "abnormal": lambda v: v is not None and v >= 6.5},
    "LDL":     {"unit":"mg/dL", "abnormal": lambda v: v is not None and v >= 160,
                "borderline": lambda v: v is not None and 130 <= v < 160},
}

def status_for(test, val):
    if test in THRESHOLDS:
        th = THRESHOLDS[test]
        v = cn_num(val)
        if "borderline" in th and th["borderline"](v):
            return "偏高"
        return "異常" if th["abnormal"](v) else "正常"
    return None

def pair_tests_values(entities):
    ents = sorted(entities, key=lambda x: x["start"])
    pairs, lone = [], []
    last_test = None
    for e in ents:
        if e["type"] == "TEST":
            if last_test: lone.append(last_test)
            last_test = {"test": e["text"], "start": e["start"], "value": None}
        elif e["type"] == "VALUE" and last_test and (e["start"] - last_test["start"]) < 40:
            last_test["value"] = e["text"]
            pairs.append({"test": last_test["test"], "value": e["text"]})
            last_test = None
    if last_test: lone.append(last_test)
    return pairs, lone

def extract_structured(text):
    tokens, entities = ner_predict_with_tokens(text)

    # basic fields
    name, ages, sex = None, [], None
    for e in entities:
        if e["type"] == "PER" and (name is None or len(e["text"]) > len(name)):
            name = e["text"]
        elif e["type"] == "AGE":
            v = cn_num(e["text"])
            if v is not None: ages.append(int(v))
        elif e["type"] == "SEX":
            sex = e["text"]

    # fallback sex detect
    if not sex:
        m_sex = re.search(r"(男|女)", text)
        if m_sex: sex = m_sex.group(1)

    pairs, _ = pair_tests_values(entities)
    key_findings = []
    for p in pairs:
        st = status_for(p["test"], p["value"])
        row = {"test": p["test"], "value": p["value"]}
        if st: row["status"] = st
        key_findings.append(row)

    risks = set()
    # 4.1 Diabetes
    fpg = next((cn_num(p["value"]) for p in pairs if p["test"] == "空腹血糖"), None)
    a1c = next((cn_num(p["value"]) for p in pairs if p["test"] == "HbA1c"), None)
    if (fpg is not None and fpg >= 126) or (a1c is not None and a1c >= 6.5):
        risks.add("糖尿病")
    # 4.2 Hyperlipidemia via LDL
    ldl = next((cn_num(p["value"]) for p in pairs if p["test"] == "LDL"), None)
    if ldl is not None and ldl >= 160:
        risks.add("高血脂")
    elif ldl is not None and ldl >= 130:
        risks.add("高血脂(輕度)")
    # 4.3 Hypertension via BP or DISEASE mention
    bp_val = next((p["value"] for p in pairs if p["test"] in ["診間血壓","家庭血壓","24小時動態血壓"]), None)
    if bp_val:
        sys, dia = parse_bp(bp_val)
        if sys and dia and (sys >= 140 or dia >= 90):
            risks.add("高血壓")
    if any(e["type"] == "DISEASE" and "高血壓" in e["text"] for e in entities):
        risks.add("高血壓")

    # recommendations
    recs = []
    for e in entities:
        if e["type"] == "DRUG":
            recs.append(f"開始服用 {e['text']}")
        elif e["type"] == "DRUG_CLASS":
            recs.append(f"考慮 {e['text']} 類藥物")
        elif e["type"] == "TREATMENT":
            t = e["text"]
            if "飲食" in t and "低鹽" not in t:
                t = "控制飲食"
            recs.append(f"建議{t}")

    age = max(ages) if ages else None
    name_disp = name if name else "病人"
    age_disp = f"{age}歲" if age is not None else ""
    abns = [f"{k['test']} {k['value']}" for k in key_findings if k.get("status") in ("異常","偏高")]
    parts = [f"{name_disp}（{age_disp}）"] if age_disp else [name_disp]
    if abns: parts.append(f"檢查顯示 " + "、".join(abns[:3]))
    if "糖尿病" in risks: parts.append("符合糖尿病診斷")
    if "高血脂" in risks or "高血脂(輕度)" in risks: parts.append("另見 LDL 偏高")
    if recs: parts.append("建議：" + "、".join(recs[:3]))
    summary = "，".join(parts) + "。"

    structured = {
        "name": name or None,
        "age": age if age is not None else None,
        "sex": sex,
        "key_findings": key_findings,
        "disease_risk": sorted(list(risks)),
        "recommendations": recs,
        "summary": summary
    }
    return tokens, entities, structured

# ---------- Human-readable report ----------
def make_readable_report(structured: dict) -> str:
    name = structured.get("name") or "病人"
    age  = structured.get("age")
    sex  = structured.get("sex")
    head = f"【健檢摘要】{name}"
    if age is not None: head += f"（{age}歲"
    else: head += "（年齡不詳"
    if sex: head += f"，{sex}）"
    else: head += "）"

    # key findings
    kfs = structured.get("key_findings", [])
    abn_lines, nor_lines = [], []
    for k in kfs:
        t, v = k.get("test"), k.get("value")
        st = k.get("status")
        if st in ("異常","偏高"):
            abn_lines.append(f"．{t}: {v}（{st}）")
        elif st == "正常":
            nor_lines.append(f"．{t}: {v}（正常）")
        else:
            nor_lines.append(f"．{t}: {v}")

    risks = structured.get("disease_risk", [])
    recs  = structured.get("recommendations", [])
    summary = structured.get("summary") or ""

    sections = [head, ""]
    if abn_lines:
        sections += ["【異常/偏高項目】"] + abn_lines + [""]
    if nor_lines:
        sections += ["【其他檢測】"] + nor_lines + [""]
    if risks:
        sections += ["【疾病風險/診斷】", "．" + "、".join(risks), ""]
    if recs:
        sections += ["【建議】", "．" + "、".join(recs), ""]
    if summary:
        sections += ["【摘要敘述】", summary, ""]
    return "\n".join(sections).strip()

# ---------- Gradio UI ----------
EXAMPLE = "李偉（65歲，男），有高血壓與糖尿病。\n診間血壓152/94mmHg，空腹血糖138mg/dL，HbA1c 7.1%。\n建議使用ARB類藥物並低鹽飲食。"

def run(text):
    tokens, entities, structured = extract_structured(text)
    human_report = make_readable_report(structured)
    return (
        human_report,
        json.dumps(structured, ensure_ascii=False, indent=2),
        json.dumps(entities, ensure_ascii=False, indent=2),
        json.dumps(tokens, ensure_ascii=False, indent=2),
    )

with gr.Blocks(title="HTN NER (Chinese)") as demo:
    gr.Markdown("## Hypertension NER → Human Report / JSON / Entities / Tokens")
    inp = gr.Textbox(label="輸入文字 (中文)", lines=6, value=EXAMPLE)
    btn = gr.Button("Analyze")
    out_report   = gr.Textbox(label="Doctor Report", lines=12)
    out_struct   = gr.Code(label="Structured JSON")
    out_entities = gr.Code(label="Entities (spans)")
    out_tokens   = gr.Code(label="Token-level (B/I/O)")
    btn.click(run, inputs=inp, outputs=[out_report, out_struct, out_entities, out_tokens])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")