| |
| """Axis 4 (entity scoping): LLM discovery, then question generation. |
| |
| Stage 1: the LLM reads candidate atoms and resolves who a condition belongs to (the patient, a named |
| relative or a household contact) and whether a model without memory could answer. Stage 2: the axis-4 |
| rules (pipeline.prompts) write the question and gold. |
| |
| python -m pipeline.questions.axis4 discovery -> _discover_axis4.json |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import re |
| from collections import Counter |
| from pathlib import Path |
|
|
| from pipeline import prompts |
| from pipeline.common import BASE, Progress, model, patient_id |
| from pipeline.llm import ChatClient, subst |
|
|
| PATIENT = patient_id() |
| ENGINE = model("generation") |
|
|
| AXIS = 4 |
| SCAN_FIELDS = ["diagnosis", "patient_self_description", "doctor_notes", "doctor_extra_notes"] |
|
|
| DISCOVER_SYS = ( |
| "You read one patient's chart and find facts whose owner is ambiguous or not obvious: a condition or\n" |
| "event that belongs to a relative, such as a mother, father, son, daughter, wife, husband or\n" |
| "sibling; a household-contact or contagion event; or the patient's own condition that a reader could\n" |
| "take for someone else's. These facts test whether a memory method attributes a fact to the right\n" |
| "person. Return only JSON.") |
|
|
|
|
| RELATIVE = re.compile(r"(?i)\b(?:mother|father|mom|dad|parents?|sisters?|brothers?|siblings?|sons?|daughters?" |
| r"|child(?:ren)?|wife|husband|spouse|family|fhx|aunt|uncle|grand\w*|niece|nephew|cousin" |
| r"|household|contact)\b") |
|
|
|
|
| HISTORY = re.compile(r"(?i)^\s*(?:history of|hx of|h/o|known)\b") |
| ABBREVIATIONS = { |
| "htn": "hypertension", "dm": "diabetes", "t2dm": "type 2 diabetes", "t1dm": "type 1 diabetes", |
| "cad": "coronary artery disease", "chf": "heart failure", "hf": "heart failure", "mi": "myocardial infarction", |
| "af": "atrial fibrillation", "afib": "atrial fibrillation", "cva": "stroke", "ckd": "chronic kidney disease", |
| "esrd": "end-stage renal disease", "copd": "chronic obstructive pulmonary disease", "tb": "tuberculosis", |
| "hbv": "hepatitis b", "hcv": "hepatitis c", "gerd": "reflux", "oa": "osteoarthritis", "ra": "rheumatoid arthritis", |
| "bph": "benign prostatic hyperplasia", "uti": "urinary tract infection", "pud": "peptic ulcer", |
| } |
|
|
|
|
| def expand(text): |
| """The lower-cased text followed by the full names of the abbreviations it contains.""" |
| t = str(text or "").lower() |
| extra = [full for ab, full in ABBREVIATIONS.items() if re.search(rf"(?<![a-z0-9]){ab}(?![a-z0-9])", t)] |
| return " ".join([t] + extra) |
|
|
|
|
| def own_diagnoses(chain): |
| """The patient's own conditions as (text, atom, visit, field), text lower-cased with abbreviations expanded: |
| diagnosis atoms and the patient_self_description atoms that state a history ("history of X"); atoms that |
| mention a relative or a household contact are left out.""" |
| atoms = [(a, r.get("visit_number"), "diagnosis") for r in chain for a in (r.get("diagnosis") or [])] |
| atoms += [(a, r.get("visit_number"), "patient_self_description") for r in chain |
| for a in (r.get("patient_self_description") or []) if HISTORY.search(str(a))] |
| return [(expand(a), a, v, f) for a, v, f in atoms if not RELATIVE.search(str(a))] |
|
|
|
|
| def patient_has(entity, own): |
| """(atom, visit, field) of the first of the patient's own conditions that names the entity (abbreviations |
| expanded), or None.""" |
| e = str(entity or "").strip().lower() |
| names = [e] if len(e) >= 4 else [] |
| names += [full for ab, full in ABBREVIATIONS.items() if re.search(rf"(?<![a-z0-9]){ab}(?![a-z0-9])", e)] |
| return next(((a, v, f) for text, a, v, f in own if any(n in text for n in names)), None) |
|
|
|
|
| def field_of(rec, text): |
| """The scanned field of a visit whose atom is the text, or else contains it (case and spacing ignored), or |
| None.""" |
| def norm(t): |
| return re.sub(r"\s+", " ", str(t or "")).strip().lower() |
| t = norm(text) |
| atoms = [(f, norm(a)) for f in SCAN_FIELDS for a in rec.get(f) or []] |
| return (next((f for f, a in atoms if t and a == t), None) |
| or next((f for f, a in atoms if t and t in a), None)) |
|
|
|
|
| def patient_owned(scoping): |
| """Whether discovery resolved the fact to the patient.""" |
| owner = str(scoping.get("owner") or "").strip().lower() |
| return scoping.get("kind") == "patient_owned" or bool( |
| re.fullmatch(r"(?:the )?patient(?:'s)?(?: (?:him|her)self)?(?: own)?", owner)) |
|
|
|
|
| def discover_scopings(cli, by_v): |
| rows = [] |
| for v in sorted(by_v): |
| for f in SCAN_FIELDS: |
| for a in by_v[v].get(f) or []: |
| rows.append(f"v{v} {by_v[v]['date']} {f}: {str(a)[:120]}") |
| chunks = ["\n".join(rows[i:i + 140]) for i in range(0, len(rows), 140)] or [""] |
| found = [] |
| for chart_facts in chunks: |
| usr = ( |
| "Chart facts, one per line (visit, date, field, text):\n" + chart_facts + "\n\n" |
| "Find facts that belong to someone other than the patient, or whose owner a reader could get\n" |
| "wrong. Classify each one:\n" |
| "- relative_owned: a condition or event of a named relative, such as family history,\n" |
| " 'mother's ...' or 'father had ...'.\n" |
| "- contagion: a household contact or exposure event, such as a family member's infection.\n" |
| "- patient_owned: the patient's own condition or event that a reader could take for someone\n" |
| " else's, for example because it is mentioned together with a relative's.\n" |
| "Skip the patient's own diagnoses that nothing links to another person. Skip negated family\n" |
| "history such as 'denied family history of X': there the entity does not exist to be owned.\n\n" |
| "For each fact give the owner (the relative, 'household contact' or 'the patient'), the entity\n" |
| "(the condition or event), the visit, and a one-line note on why a model without memory would\n" |
| "attribute it to the wrong person.\n\n" |
| "Return JSON in this form:\n" |
| '{"scopings": [{\n' |
| ' "visit": <int>,\n' |
| ' "kind": "<relative_owned|contagion|patient_owned>",\n' |
| ' "owner": "<who it belongs to, e.g. \'the patient\'s mother\' or \'the patient\'>",\n' |
| ' "entity": "<the condition or event>",\n' |
| ' "atom_verbatim": "<the source text>",\n' |
| ' "why_memory_required": "<one line>"\n' |
| "}]}") |
| try: |
| found += cli.complete_json(DISCOVER_SYS, usr).get("scopings") or [] |
| except Exception as e: |
| print(f" discover chunk FAILED: {str(e)[:100]}") |
| return found |
|
|
|
|
| def discover() -> None: |
| """Discover the axis-4 candidates of this patient -> _discover_axis4.json.""" |
| cli = ChatClient(model=ENGINE, reasoning_effort="medium", max_completion_tokens=6000) |
|
|
| chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl") if l.strip()] |
| by_v = {r["visit_number"]: r for r in chain} |
| plast = max(r["date"] for r in chain) |
|
|
| with Progress("axis4 discover (LLM)", total=1, out=f"{BASE}/progress_axis4_discover.json") as pg: |
| scopings = discover_scopings(cli, by_v) |
| pg.tick(ok=bool(scopings)) |
|
|
| AX4 = subst(prompts.AXIS4_RULES) |
| COMMON = subst(prompts.AXIS4_COMMON) |
| system = subst(f"{COMMON}\n\n{AX4}") |
|
|
| candidates = [] |
| seen = set() |
| own = own_diagnoses(chain) |
| for s in scopings: |
| rv = s.get("visit") |
| if rv not in by_v: |
| continue |
| k = (str(s.get("owner")).lower().strip(), str(s.get("entity")).lower().strip()[:40]) |
| if k in seen: |
| continue |
| seen.add(k) |
| own_case = patient_owned(s) |
| also = None if own_case else patient_has(s.get("entity"), own) |
| seed = {"kind": s.get("kind"), "owner": s.get("owner"), "entity": s.get("entity"), |
| "atom_verbatim": s.get("atom_verbatim"), "why_memory_required": s.get("why_memory_required")} |
| if also: |
| seed["patient_also_has"] = also[0] |
| fact_json = json.dumps(seed, ensure_ascii=False) |
| usr = ( |
| "Write one entity-attribution question for this seed, following the axis-4 rules above. The\n" |
| "question must not reveal the owner; the gold answer states it. Return only JSON:\n" |
| '{"question": "<question>", "answer": "<gold naming the correct owner>"}\n\n' |
| "Seed:\n" + fact_json) |
| try: |
| q = cli.complete_json(system, usr) |
| except Exception as e: |
| print(f" gen FAILED: {str(e)[:100]}") |
| continue |
| if not q.get("question") or not q.get("answer"): |
| continue |
| if str(q["answer"]).lstrip().startswith("Yes") != (own_case or bool(also)): |
| print(f" reject {s.get('entity')!r}: the gold's yes/no does not match the resolved owner") |
| continue |
| evidence = [{"visit": rv, "date": by_v[rv]["date"], "field": field_of(by_v[rv], s.get("atom_verbatim")), |
| "atom": s.get("atom_verbatim")}] |
| owner_step = f"Owner resolved: {s.get('owner')} (kind={s.get('kind')})." |
| ev_visits = [rv] |
| if also: |
| patom, pv, pfield = also |
| if pv != rv: |
| ev_visits.append(pv) |
| evidence.append({"visit": pv, "date": by_v[pv]["date"] if pv in by_v else None, "field": pfield, |
| "atom": patom}) |
| owner_step = (f"The fact at v{rv} is {s.get('owner')}'s; the patient's own {s.get('entity')} is " |
| f"documented at v{pv}.") |
| gv = max(v for v in ev_visits if v in by_v) |
| candidates.append({ |
| "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT, |
| "axis": "4_entity_scoping", "question": q["question"], "answer": q["answer"], |
| "answer_type": "open_exact", "grounding": "record", |
| "evidence_visits": ev_visits, "gt_visit": gv, "gt_date": by_v[gv]["date"], |
| "question_appearance_window": {"from": by_v[gv]["date"], "to": plast}, |
| "derivation": { |
| "evidence": evidence, |
| "steps": [f"Atom at v{rv}: {s.get('atom_verbatim')!r}.", owner_step, |
| f"{s.get('why_memory_required')}"], |
| "gold": q["answer"], "rivals": [], |
| "operands": {"owner": s.get("owner"), "entity": s.get("entity"), "kind": s.get("kind"), |
| "patient_also_has": also[0] if also else None}, |
| "derived": False}, |
| "provenance": f"LLM-resolved scoping ({ENGINE}); owner={s.get('owner')}", |
| "axis_specific": {"kind": s.get("kind"), "owner": s.get("owner"), "entity": s.get("entity")}, |
| }) |
|
|
| out = {"axis": AXIS, "patient_id": PATIENT, "n_record": len(candidates), "candidates": candidates} |
| Path(f"{BASE}/_discover_axis{AXIS}.json").write_text(json.dumps(out, ensure_ascii=False, indent=1)) |
| print(f"\npatient {PATIENT} engine={ENGINE}") |
| print(f" scopings discovered (LLM): {len(scopings)} by kind: {dict(Counter(s.get('kind') for s in scopings))}") |
| print(f" CANDIDATES: {len(candidates)} -> {BASE}/_discover_axis{AXIS}.json") |
| for c in candidates[:3]: |
| print(f" Q: {c['question'][:70]} A: {str(c['answer'])[:40]}") |
|
|
|
|
| if __name__ == "__main__": |
| discover() |
|
|
