| |
| """Axis 2 (state change): seeds from the evidence chain, LLM generation (pipeline.llm step axis2), post-processing. |
| |
| seeds categorical state transitions of a condition, treatment or diagnosis (numeric trends belong to |
| axis 6): 'evolved' (the current value supersedes a different one) and 'resolved' (the latest |
| segment ended) -> _seed_axis2.json |
| post GEN recomputes the appearance window, deduplicates, assigns qid and grounding, and counts the questions |
| that contain a date or a distinctive token of the old value or of their own answer |
| -> questions_axis2.jsonl |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| from collections import Counter, defaultdict |
|
|
| from pipeline.common import BASE, has_text, norm_value, other_department_segments, patient_id |
|
|
| PATIENT = patient_id() |
|
|
|
|
| def seeds() -> None: |
| """Seed builder; writes _seed_axis2.json.""" |
| ALLOWED = {"disease_state": "diagnosis", "symptom": "symptom", "medication": "medication", |
| "procedure_or_followup": "treatment", "care_plan": "treatment", |
| "differential_impression": "differential"} |
| CAP_PER_KEY = 4 |
|
|
| UNIT = re.compile(r"\b(mg|kg|lb|lbs|cgy|gy|mmhg|ml|cm|mm|fx|iu|meq)\b", re.I) |
|
|
| def numeric_dominated(vals): |
| hits = 0 |
| for v in vals: |
| vs = v.lower() |
| if UNIT.search(vs) or "%" in vs or re.search(r"\d+\s*(?:/|->|~|-)\s*\d+", vs) \ |
| or re.search(r"(?:flex|ext|iop|grade|loss)\s*[:=]?\s*\d", vs): |
| hits += 1 |
| return hits >= max(1, (len(vals) + 1) // 2) |
|
|
| def norm_val(v): |
| return re.sub(r"\s+", " ", re.sub(r"\s*\([^)]*\)", "", v).strip().lower()) |
|
|
| ADMIN = re.compile(r"opd|follow.?up|f\s*/\s*u|\bfu\b|refill|schedul|appointment", re.I) |
| ONGOING = re.compile(r"still|persist|ongoing|complain|continu|recurrent|worse|flare", re.I) |
|
|
| tl = [json.loads(l) for l in open(f"{BASE}/evidence_chain.jsonl")] |
|
|
| seeds = [] |
| for t in tl: |
| asp = t["aspect"] |
| if asp not in ALLOWED: |
| continue |
| if ADMIN.search(t["entity"]): |
| continue |
| segs = t["segments"] |
| if len(segs) < 2: |
| continue |
| latest = segs[-1] |
| old_vals = [s["value"] for s in segs[:-1]] |
| all_vals = [s["value"] for s in segs] |
| if numeric_dominated(all_vals): |
| continue |
| st = latest.get("status") |
| if st == "current": |
| if norm_val(latest["value"]) == norm_val(segs[0]["value"]): |
| continue |
| change_type, current_value = "evolved", latest["value"] |
| elif st == "resolved": |
| if ONGOING.search(latest["value"]): |
| continue |
| change_type, current_value = "resolved", latest["value"] |
| else: |
| continue |
| vf = latest["valid_from"] |
| if any(((o.get("valid_from") or {}).get("date") or "") > vf["date"] |
| and norm_value(o.get("value")) != norm_value(current_value) |
| for o in other_department_segments(tl, t["entity"], t["department"])): |
| continue |
| seeds.append({ |
| "change_type": change_type, |
| "entity": t["entity"], "aspect": asp, "qtype": ALLOWED[asp], "department": t["department"], |
| "current_value": current_value, |
| "old_values_DO_NOT_MENTION": old_vals, |
| "gt_date": vf["date"], "gt_visit": vf["visit_number"], |
| "evidence_visits": sorted({v for s in segs for v in s.get("supporting_visits", [])})[-8:], |
| "n_segments": len(segs), "current_status": t.get("current_status", ""), |
| }) |
|
|
| dedup = {} |
| for s in seeds: |
| k = norm_val(s["entity"]) |
| if k not in dedup or s["n_segments"] > dedup[k]["n_segments"]: |
| dedup[k] = s |
| seeds = list(dedup.values()) |
|
|
| by = defaultdict(list) |
| for s in seeds: |
| by[(s["qtype"], s["change_type"])].append(s) |
| out = [] |
| for k, lst in by.items(): |
| lst.sort(key=lambda x: -x["n_segments"]) |
| out += lst[:CAP_PER_KEY] |
|
|
| json.dump({"axis": "2_state_change", "seeds": out}, open(f"{BASE}/_seed_axis2.json", "w"), |
| ensure_ascii=False, indent=1) |
|
|
| print(f"axis-2 state-change seeds: {len(out)} (from {len(seeds)} candidates)") |
| print("by change_type:", dict(Counter(s["change_type"] for s in out))) |
| print("by qtype :", dict(Counter(s["qtype"] for s in out))) |
| print() |
| for s in out: |
| print(f" [{s['change_type']:8}/{s['qtype']:10}] {s['entity'][:38]:38} CURR={s['current_value'][:32]!r}") |
| print(f" old(forbidden)={[v[:22] for v in s['old_values_DO_NOT_MENTION']][:3]}") |
| print("\nwrote _seed_axis2.json") |
|
|
|
|
| def post(gen_path: str) -> None: |
| """Post-processing of the generated axis-2 questions; writes questions_axis2.jsonl.""" |
| LAST_V = max(json.loads(_l)["visit_number"] for _l in open(f"{BASE}/event_chain.jsonl")) |
| MIN_HORIZON_VISITS = 24 |
|
|
| recs = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl")] |
| vdate = {r["visit_number"]: r["date"] for r in recs} |
| last_date = max(vdate.values()) |
| seeds = {s["entity"]: s for s in json.load(open(f"{BASE}/_seed_axis2.json"))["seeds"]} |
|
|
| def gt_visit(d): |
| later = [r["visit_number"] for r in recs if r["date"] >= d] |
| return min(later) if later else 1 |
|
|
| WORD = re.compile(r"[a-z0-9]+") |
| STOP = {"the", "patient", "current", "currently", "still", "have", "has", "does", "did", "any", "what", "is", |
| "was", "are", "for", "and", "with", "left", "right", "bilateral", "bil", "his", "her", "of", "to", "on", |
| "disease", "status", "diagnosis", "diagnoses", "condition", "evidence", "stage", "chronic", "disorder", |
| "finding", "findings", "documented", "workup", "middle", "overall", "functional", "receive", "part", |
| "treatment", "symptom", "symptoms", "problem", "area", "recent", "no", "not"} |
|
|
| def toks(s): |
| return {w for w in WORD.findall(str(s).lower()) if len(w) >= 4 and w not in STOP} |
|
|
| DATE = re.compile(r"\b\d{4}-\d{2}-\d{2}\b|\b\d{1,2}\s*/\s*\d{1,2}\b|\b(19|20)\d{2}\b") |
|
|
| raw = json.load(open(gen_path)) |
| qs = raw["questions"] if isinstance(raw, dict) else raw |
|
|
| out, seen, n_leak = [], set(), 0 |
| no_text = sum(1 for q in qs if not has_text(q)) |
| if no_text: |
| print(f"skipped {no_text} generated question(s) without question text") |
| qs = [q for q in qs if has_text(q)] |
| bad = [q for q in qs if not (isinstance(q.get("gt_date"), str) and q["gt_date"].strip())] |
| for q in bad: |
| print(f"skipped a generated question without a gt_date: {q['question'][:60]!r}") |
| qs = [q for q in qs if q not in bad] |
| for q in qs: |
| key = q["question"].strip().lower() |
| if key in seen: |
| continue |
| seen.add(key) |
| gv = gt_visit(q["gt_date"]) |
| fv = min(gv + MIN_HORIZON_VISITS, LAST_V) |
| q["gt_visit"] = gv |
| q["question_appearance_window"] = {"from": vdate[fv], "to": last_date} |
| q["axis"] = "2_state_change" |
| q["qid"] = f"{PATIENT}/axis2/{len(out):03d}" |
| q["grounding"] = "record" |
|
|
| qtext = q["question"] |
| qtok = toks(qtext) |
| sd = seeds.get(q.get("entity", ""), {}) |
| old_vals = sd.get("old_values_DO_NOT_MENTION", []) |
| cur = sd.get("current_value", q.get("answer", "")) |
| stable = toks(sd.get("entity", "")) |
| old_distinct = (set().union(*[toks(v) for v in old_vals]) if old_vals else set()) - stable |
| ans_distinct = toks(cur) - stable |
| if DATE.search(qtext) or qtok & old_distinct or qtok & ans_distinct: |
| n_leak += 1 |
| out.append(q) |
|
|
| with open(f"{BASE}/questions_axis2.jsonl", "w", encoding="utf-8") as fh: |
| for q in out: |
| fh.write(json.dumps(q, ensure_ascii=False) + "\n") |
|
|
| print(f"axis-2 questions: {len(out)} ({len(qs) - len(out)} dropped as dup)") |
| print("by change_type:", dict(Counter(q["change_type"] for q in out))) |
| print("by qtype :", dict(Counter(q["qtype"] for q in out))) |
| print(f"leak-lint violations: {n_leak}") |
| print("\nwrote questions_axis2.jsonl") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| sp = ap.add_subparsers(dest="cmd", required=True) |
| sp.add_parser("seeds") |
| sp.add_parser("post").add_argument("generated") |
| a = ap.parse_args() |
| seeds() if a.cmd == "seeds" else post(a.generated) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
