#!/usr/bin/env python3 """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()