| |
| """Axis 6 (temporal localization and trend): LLM extraction with deterministic computation. |
| |
| The LLM reads each candidate timeline and returns its numeric series. classify() computes the trend |
| shape from the series; extremum questions take the highest or lowest reading. Golds are formatted from |
| PROSE_TEMPLATE_* with the computed operands. synthesize() authors synthetic questions. |
| |
| python -m pipeline.questions.axis6 discovery -> _discover_axis6.json |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import random |
| from datetime import date, timedelta |
| from pathlib import Path |
|
|
| from pipeline.common import BASE, FIELDS, Progress, model, patient_context, patient_id, patient_line |
| from pipeline.llm import ChatClient |
|
|
| PATIENT = patient_id() |
| ENGINE = model("generation") |
|
|
| AXIS = 6 |
| SHAPE_VOCAB = ["steadily increased", "steadily decreased", "increased then decreased", |
| "decreased then increased", "remained essentially unchanged"] |
|
|
| PROSE_TEMPLATE_EXTREMUM_Q = ("Across the whole record, what was this patient's {extremum_kind} charted " |
| "{series_entity} value, and when was it measured?") |
| PROSE_TEMPLATE_EXTREMUM_A = ("The {extremum_kind} charted {series_entity} was {extremum_num:g}{unit_sp} " |
| "on {extremum_date}.") |
| PROSE_TEMPLATE_MC_Q = ("Between {range_from} and {range_to}, how did this patient's {series_entity} " |
| "change over the readings taken in that period?") |
| PROSE_TEMPLATE_MC_A = "Over {range_from} to {range_to}, {series_entity} {correct_option_text}." |
|
|
|
|
| def sign(x): |
| return 1 if x > 0 else (-1 if x < 0 else 0) |
|
|
|
|
| def classify(ns): |
| """Trend shape of a series, computed deterministically; None rejects the series.""" |
| lo, hi, med = min(ns), max(ns), sorted(ns)[len(ns) // 2] |
| if med != 0 and (hi - lo) / abs(med) <= 0.05: |
| return "remained essentially unchanged" |
| d = [b - a for a, b in zip(ns, ns[1:])] |
| nz = [sign(x) for x in d if x != 0] |
| if not nz: |
| return "remained essentially unchanged" |
| flips = sum(1 for a, b in zip(nz, nz[1:]) if a != b) |
| if flips == 0: |
| return "steadily increased" if nz[0] > 0 else "steadily decreased" |
| if flips == 1: |
| return "increased then decreased" if nz[0] > 0 else "decreased then increased" |
| return None |
|
|
|
|
| def best_window(points): |
| nums = [p["num"] for p in points] |
| for length in range(len(nums), 2, -1): |
| for i in range(0, len(nums) - length + 1): |
| sh = classify(nums[i:i + length]) |
| if sh is not None: |
| return (i, i + length, sh) |
| return None |
|
|
|
|
| def mid(d1, d2): |
| a, b = date.fromisoformat(d1), date.fromisoformat(d2) |
| return (a + (b - a) / 2).isoformat() |
|
|
|
|
| EXTRACT_SYS = ( |
| "You read one entity's timeline from a synthetic longitudinal medical chart and extract its\n" |
| "numeric readings. Charts are terse, and their format varies across patients and departments.\n" |
| "Return only JSON.") |
|
|
| EXTRACT_SCHEMA = """Return JSON in this form: |
| {"is_numeric_series": <true|false>, |
| "measurand": "<what is actually measured, normalized, e.g. 'Total Cholesterol' or 'Body weight'>", |
| "unit": "<unit, or '' if none>", |
| "points": [{"date": "YYYY-MM-DD", "visit": <int>, "num": <number>, "raw": "<the source text>"}], |
| "aliases": ["<every name or abbreviation this chart uses for this measurand>"], |
| "reject_reason": "<why, if is_numeric_series is false>"}""" |
|
|
|
|
| def extract_series(cli, t): |
| """LLM extraction of the numeric readings of one timeline. |
| |
| Trailing parentheticals are dropped; dates, censored values ("< 6") and readings that mix two measurands |
| are rejected. |
| """ |
| segs = [{"value": s.get("value"), "date": (s.get("valid_from") or {}).get("date"), |
| "visit": (s.get("valid_from") or {}).get("visit_number")} for s in t.get("segments") or []] |
| entity, department, aspect = t.get("entity"), t.get("department"), t.get("aspect") |
| segment_lines = "\n".join(f"{s['date']} v{s['visit']}: {s['value']!r}" for s in segs) |
| usr = ( |
| f"Entity: {entity}\nDepartment: {department}\nAspect: {aspect}\n\n" |
| "Segments, each one value and the date it was first recorded:\n" + segment_lines + "\n\n" |
| "Extract the numeric readings of one measurand.\n\n" |
| "Reject the series (is_numeric_series false) if any of these holds:\n" |
| "- The values are not numeric measurements but impressions, findings or narrative.\n" |
| "- More than one measurand is charted here, for example '<lab1>/<lab2> <value1>/<value2>', which\n" |
| " is ambiguous.\n" |
| "- The values are censored, such as '< 6' or '> 60', and cannot be ordered.\n" |
| "- Fewer than 3 readings carry a real number on distinct dates.\n" |
| "- A number is a date or a lab accession number, not a reading.\n\n" |
| "Take the number that is the reading, and ignore trailing text in parentheses, lab codes, embedded\n" |
| "dates, reference ranges and targets such as 'goal LDL < 100'. Read a value in scientific notation,\n" |
| "such as '1.2x10^3', as a number. Give one reading per date, and use the segment's own date and\n" |
| "visit.\n\n" |
| + EXTRACT_SCHEMA) |
| r = cli.complete_json(EXTRACT_SYS, usr) |
| if not r.get("is_numeric_series"): |
| return None, r.get("reject_reason", "not a numeric series") |
| pts = [p for p in (r.get("points") or []) if isinstance(p.get("num"), (int, float)) and p.get("date")] |
| seen, uniq = set(), [] |
| for p in sorted(pts, key=lambda x: x["date"]): |
| if p["date"] in seen: |
| continue |
| seen.add(p["date"]) |
| uniq.append(p) |
| if len(uniq) < 3: |
| return None, f"only {len(uniq)} usable readings" |
| return {"entity": r.get("measurand") or t.get("entity"), "department": t.get("department"), |
| "unit": (r.get("unit") or "").strip(), "points": uniq, |
| "aliases": [a.lower() for a in (r.get("aliases") or [])]}, None |
|
|
|
|
| UNIQ_SYS = ("You check whether a claimed extreme value is unique in a patient's chart. Return only JSON.") |
| LINES_PER_CHECK = 70 |
|
|
|
|
| def extremum_is_unique(cli, s, kind, ex_num, ex_date, hay_rows): |
| """Whether the extreme reading of a series is unique, checked against every chart line naming the |
| measurand, LINES_PER_CHECK lines per call; the first group that finds a higher (or lower) value or an |
| equal value on another date decides.""" |
| names = ", ".join(sorted({s["entity"], *s.get("aliases", [])})) |
| rel = [f" {d} v{v}: {tx}" for v, d, tx in hay_rows |
| if any(n and n.lower() in tx.lower() for n in {s["entity"], *s.get("aliases", [])})] |
| for i in range(0, len(rel), LINES_PER_CHECK): |
| unique, violations = _unique_in(cli, names, kind, ex_num, ex_date, rel[i:i + LINES_PER_CHECK]) |
| if not unique: |
| return False, violations |
| return True, [] |
|
|
|
|
| def _unique_in(cli, names, kind, ex_num, ex_date, rel): |
| """One uniqueness call over a group of chart lines.""" |
| chart_lines = "\n".join(rel) |
| beyond = "higher" if kind == "highest" else "lower" |
| r = cli.complete_json( |
| UNIQ_SYS, |
| f"Measurand: {names}\nClaim: the {kind} charted value is {ex_num:g} on {ex_date}.\n\n" |
| "Chart lines that mention it:\n" + chart_lines + "\n\n" |
| "Each line starts with the date of the visit that recorded it. Ignore reference ranges and targets.\n" |
| "Is the claim true, and is that date unique?\n" |
| f"- A line showing a {beyond} value anywhere makes the claim false.\n" |
| "- A separate measurement with an equal value on another date makes it not unique. A reading\n" |
| " copied forward into a later visit, with the same value and its original test date or no date\n" |
| " of its own, is the same reading and is fine, because charts copy readings forward.\n\n" |
| "Return JSON in this form:\n" |
| '{"unique": <true|false>,\n' |
| ' "violations": [{"date": "...", "value": <n>, "why": "beats|ties", "line": "..."}]}') |
| return bool(r.get("unique")), r.get("violations") or [] |
|
|
|
|
| def discover() -> None: |
| """Discover the axis-6 candidates of this patient -> _discover_axis6.json.""" |
| cli = ChatClient(model=ENGINE, reasoning_effort="medium", max_completion_tokens=4000) |
|
|
| ts = [json.loads(l) for l in open(f"{BASE}/evidence_chain.jsonl") if l.strip()] |
| chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl") if l.strip()] |
| LAST_V = max(r["visit_number"] for r in chain) |
| hay_rows = [(r["visit_number"], r["date"], str(a)) |
| for r in chain for f in FIELDS + ["icd_notes"] for a in (r.get(f) or [])] |
| hay_rows += [(r["visit_number"], r["date"], f"{k}={val}") |
| for r in chain for k, val in (r.get("vitals") or {}).items()] |
|
|
| cand = [t for t in ts if len(t.get("segments") or []) >= 3] |
| series, rejects = [], [] |
| with Progress("axis6 extract (LLM)", total=len(cand), out=f"{BASE}/progress_axis6_discover.json") as pg: |
| for t in cand: |
| try: |
| s, why = extract_series(cli, t) |
| except Exception as e: |
| s, why = None, f"error: {str(e)[:80]}" |
| if s: |
| series.append(s) |
| else: |
| rejects.append((t.get("entity"), why)) |
| pg.tick(ok=bool(s)) |
|
|
| candidates = [] |
| for s in series: |
| pts, nums = s["points"], [p["num"] for p in s["points"]] |
| ent, unit_sp = s["entity"], (f" {s['unit']}" if s["unit"] else "") |
| evidence = [{"visit": p.get("visit"), "date": p["date"], |
| "field": "evidence_chain", "atom": p.get("raw")} for p in pts] |
|
|
| for kind in ("highest", "lowest"): |
| ex = max(nums) if kind == "highest" else min(nums) |
| if nums.count(ex) != 1: |
| continue |
| k = nums.index(ex) |
| ex_date, ex_visit = pts[k]["date"], pts[k].get("visit") |
| try: |
| ok, _ = extremum_is_unique(cli, s, kind, ex, ex_date, hay_rows) |
| except Exception: |
| ok = False |
| if not ok: |
| continue |
| ops = {"extremum_kind": kind, "series_entity": ent, "extremum_num": ex, |
| "unit_sp": unit_sp, "extremum_date": ex_date, "n_points": len(pts), |
| "series_department": s["department"], "unit": s["unit"]} |
| answer = PROSE_TEMPLATE_EXTREMUM_A.format(**ops) |
| candidates.append({ |
| "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT, |
| "axis": f"{AXIS}_{kind}_value_and_date", |
| "question": PROSE_TEMPLATE_EXTREMUM_Q.format(**ops), "answer": answer, |
| "answer_type": "open_exact", "grounding": "record", |
| "evidence_visits": [p.get("visit") for p in pts], |
| "gt_visit": ex_visit, "gt_date": ex_date, |
| "question_appearance_window": {"from": ex_date, "to": max(r["date"] for r in chain)}, |
| "derivation": {"evidence": evidence, |
| "steps": [f"{ent} has {len(pts)} readings: {[(p['date'], p['num']) for p in pts]}.", |
| f"{kind} = {ex:g}{unit_sp}, at {ex_date}.", |
| f"Uniqueness confirmed against every chart line naming {ent}."], |
| "gold": answer, "rivals": [], "operands": ops, "derived": True}, |
| "provenance": f"LLM-extracted series ({ENGINE}); {kind} unique across the chart", |
| "axis_specific": {"template": f"{kind}_value_and_date", "n_points": len(pts)}, |
| }) |
|
|
| w = best_window(pts) |
| if w: |
| i, j, shape = w |
| wp = pts[i:j] |
| rf = mid(pts[i - 1]["date"], wp[0]["date"]) if i > 0 else \ |
| (date.fromisoformat(wp[0]["date"]) - timedelta(days=15)).isoformat() |
| rt = mid(wp[-1]["date"], pts[j]["date"]) if j < len(pts) else \ |
| (date.fromisoformat(wp[-1]["date"]) + timedelta(days=15)).isoformat() |
| rnd = random.Random(f"{PATIENT}|{ent}|trend") |
| opts = SHAPE_VOCAB[:] |
| rnd.shuffle(opts) |
| opts = opts[:4] |
| if shape not in opts: |
| opts[rnd.randrange(4)] = shape |
| ops = {"range_from": rf, "range_to": rt, "series_entity": ent, |
| "correct_option_text": shape, "shape_class": shape, |
| "series_department": s["department"], "n_window": len(wp), "n_points": len(pts)} |
| answer = PROSE_TEMPLATE_MC_A.format(**ops) |
| nums_w = [p["num"] for p in wp] |
| candidates.append({ |
| "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT, |
| "axis": f"{AXIS}_trend_direction_mc", |
| "question": PROSE_TEMPLATE_MC_Q.format(**ops), "answer": answer, |
| "answer_type": "multiple_choice", "grounding": "record", |
| "options": [{"label": "ABCD"[k], "content": o, "is_correct": o == shape} |
| for k, o in enumerate(opts)], |
| "evidence_visits": [p.get("visit") for p in wp], |
| "gt_visit": wp[-1].get("visit"), "gt_date": wp[-1]["date"], |
| "question_appearance_window": {"from": rt, "to": max(r["date"] for r in chain)}, |
| "derivation": {"evidence": evidence[i:j], |
| "steps": [f"{ent} in {rf}..{rt}: {nums_w}.", |
| f"Deltas: {[round(b-a,2) for a,b in zip(nums_w, nums_w[1:])]}.", |
| f"classify() -> '{shape}'."], |
| "gold": answer, "rivals": [], "operands": ops, "derived": True}, |
| "provenance": f"LLM-extracted series ({ENGINE}); shape COMPUTED by classify()", |
| "axis_specific": {"template": "trend_direction_mc", "shape_class": shape}, |
| }) |
|
|
| 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} LAST_V={LAST_V} engine={ENGINE}") |
| print(f" timelines with >=3 segments : {len(cand)}") |
| print(f" LLM-extracted numeric series: {len(series)}") |
| print(f" rejected : {len(rejects)}") |
| for e, why in rejects[:5]: |
| print(f" {str(e)[:38]:<38} {str(why)[:56]}") |
| print(f" CANDIDATES: {len(candidates)} -> {BASE}/_discover_axis{AXIS}.json") |
|
|
|
|
| def synthesize(cli, chain, by_v, forbidden, n, batch, existing_series): |
| """Author `n` numeric series whose values follow a fixed profile for the intended shape; one LLM call picks the |
| measurands and baselines.""" |
| out, vs = [], sorted(by_v) |
| SHAPES = ["increased then decreased", "decreased then increased", "steadily increased", "steadily decreased"] |
| FRAC = {"increased then decreased": [0.2, 0.5, 0.85, 0.5, 0.25], |
| "decreased then increased": [0.85, 0.5, 0.15, 0.5, 0.8], |
| "steadily increased": [0.15, 0.4, 0.65, 0.9], |
| "steadily decreased": [0.9, 0.65, 0.4, 0.15]} |
| patient, context = patient_line(), patient_context(chain) |
| tracked_list = "\n".join(f"- {e}" for e in existing_series[:50]) |
| forbidden_inline = ", ".join(forbidden[:25]) |
| r = cli.complete_json( |
| "You choose recurring numeric lab tests for a trend question of a medical-memory benchmark.\n" |
| "Return only JSON.", |
| f"{patient}\n{context}\n\n" |
| f"Choose {n} different recurring numeric lab tests that this patient's care might track over\n" |
| "years. Give a plausible baseline value and unit for each; the trajectory is computed from it.\n\n" |
| "Avoid lab tests that the chart already tracks, because a shared word or abbreviation makes both\n" |
| "unusable:\n" + tracked_list + |
| "\nAlso avoid these, which answer keys already track: " + forbidden_inline + "\n\n" |
| "Return JSON in this form:\n" |
| '{"labs": [{"label": "<chart name>", "unit": "<unit or \'\'>", "baseline": <number>,\n' |
| ' "decimals": <0|1|2>, "entity": "<measurand>", "topic": "<slug>",\n' |
| ' "department": "<one of the departments listed above>"}]}') |
| labs = r.get("labs") or [] |
| if not isinstance(labs, list): |
| print(" axis6: the reply's labs are not a list") |
| labs = [] |
| depts = {x["department"] for x in chain} |
| for i, lab in enumerate(labs[:n]): |
| if not (isinstance(lab, dict) and all(isinstance(lab.get(k), str) and lab[k].strip() |
| for k in ("label", "entity"))): |
| print(f" axis6[{i}] rejected: the lab lacks a label or an entity") |
| continue |
| try: |
| base_v, dec = float(lab["baseline"]), int(lab.get("decimals", 1)) |
| except (KeyError, TypeError, ValueError): |
| print(f" axis6[{i}] rejected: the lab lacks a numeric baseline") |
| continue |
| shape = SHAPES[i % len(SHAPES)] |
| frac = FRAC[shape] |
| dept = next((d for d in depts if lab.get("department", "") and lab["department"] in d), |
| by_v[vs[i % len(vs)]]["department"]) |
| hosts = [v for v in vs if by_v[v]["department"] == dept][:len(frac)] |
| if len(hosts) < 3: |
| hosts = vs[i::max(1, len(vs) // len(frac))][:len(frac)] |
| if len(hosts) < 3: |
| continue |
| unit = str(lab.get("unit") or "").strip() |
| pts, atoms = [], [] |
| for k, v in enumerate(hosts): |
| val = round(base_v * (0.5 + frac[k]), dec) |
| txt = f"{lab['label']} {val:g} {unit}".strip() |
| pts.append((v, by_v[v]["date"], val)) |
| atoms.append({"atom_id": f"{batch}/{i:02d}_{k:02d}", "batch_id": batch, "axis": "6", |
| "visit_number": v, "date": by_v[v]["date"], "department": by_v[v]["department"], |
| "field": "tests_and_medication", "text": txt, "entities": [lab["entity"]], |
| "topic": lab.get("topic"), "anchors_axis": "6", "anchors_on": []}) |
| nums = [p[2] for p in pts] |
| gold = f"Over {pts[0][1]} to {pts[-1][1]}, {lab['entity']} {shape}." |
| out.append({"atom": atoms[0], "atoms": atoms, "q": { |
| "qid": f"{PATIENT}/axis6/S{i:03d}", "patient_id": PATIENT, "axis": "6_trend_direction", |
| "question": (f"Between {pts[0][1]} and {pts[-1][1]}, how did this patient's {lab['entity']} " |
| f"change over the readings taken in that period?"), |
| "answer": gold, "answer_type": "open_exact", "grounding": f"synthetic:{batch}", |
| "evidence_visits": [p[0] for p in pts], "gt_visit": pts[-1][0], "gt_date": pts[-1][1], |
| "question_appearance_window": {"from": pts[-1][1], "to": by_v[vs[-1]]["date"]}, |
| "derivation": {"evidence": [{"visit": p[0], "date": p[1], "field": "tests_and_medication", |
| "atom": f"{lab['label']} {p[2]:g} {unit}".strip()} for p in pts], |
| "steps": [f"{lab['entity']} measured {len(pts)} times: {nums}.", |
| f"Deltas: {[round(b-a,2) for a,b in zip(nums,nums[1:])]}.", |
| f"Therefore the trajectory {shape}."], |
| "gold": gold, "rivals": [], |
| "operands": {"series_entity": lab["entity"], "series_department": dept, |
| "unit": unit, "points": [{"date": p[1], "num": p[2]} for p in pts], |
| "n_points": len(pts), "shape_class": shape}, |
| "derived": True}, |
| "provenance": f"synthetic trend ({batch}); value sequence computed so the shape is exact " |
| f"(absolute values immaterial for a memory test)", |
| "axis_specific": {"shape_class": shape, "depends_on_atoms": [x["atom_id"] for x in atoms]}, |
| }}) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| discover() |
|
|
