| |
| """Export one finished patient build into the per-patient files of the benchmark. |
| |
| event_chain.jsonl <- event_chain.jsonl (patient_id set to the exported identifier) |
| evidence_chain.jsonl <- evidence_chain.jsonl (entity null for the facts of synthetic questions) |
| dialogues.jsonl <- dialogues.jsonl (meta keeps only atom_index) |
| questions.jsonl <- questions_axis{1..8}.jsonl + checkpoints.jsonl + validation.jsonl |
| (sound questions in axis-file order with their checkpoint; axis-3 questions with the |
| same options and answer kept once; axis 8 keeps at most AXIS8_PER_MECHANISM questions per |
| mechanism, evenly spaced by evidence-to-checkpoint distance, not-yet-due deadlines first) |
| |
| python -m pipeline.export --work build/<patient_id> --out <dir> --patient-id P<nn> |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import re |
| from collections import Counter |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| import tiktoken |
|
|
| COPIES = [("event_chain.jsonl", "event_chain.jsonl"), ("evidence_chain.jsonl", "evidence_chain.jsonl"), |
| ("dialogues.jsonl", "dialogues.jsonl")] |
| AXIS8_PER_MECHANISM = {"deadline": 10, "event_bounded": 5, "fixed_duration": 5} |
| AXIS8_STATE_ORDER = {"deadline": ["not_yet_due", "overdue"]} |
|
|
|
|
| @lru_cache(maxsize=None) |
| def enc(): |
| """The tiktoken encoding named by MEMEBENCH_TOKENIZER, loaded on first use.""" |
| name = os.environ.get("MEMEBENCH_TOKENIZER", "") |
| if not name or name.startswith("<"): |
| raise SystemExit("set MEMEBENCH_TOKENIZER (the tiktoken encoding of the exported size fields)") |
| return tiktoken.get_encoding(name) |
|
|
|
|
| def jl(path): |
| return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()] |
|
|
|
|
| def check_build(w): |
| need = [f"questions_axis{a}.jsonl" for a in range(1, 9)] + [ |
| "checkpoints.jsonl", "validation.jsonl", "MANIFEST.json", "visits.jsonl"] + [s for s, _ in COPIES] |
| missing = [f for f in need if not (w / f).exists()] |
| if missing: |
| raise SystemExit(f"{w}: build incomplete, missing {missing}") |
| empty = [r["visit_number"] for r in jl(w / "dialogues.jsonl") if not r.get("turns")] |
| if empty: |
| raise SystemExit(f"{w}: {len(empty)} dialogues have no turns (re-run step F): {empty[:20]}") |
| n_visits = len(jl(w / "visits.jsonl")) |
| n_dia = len(jl(w / "dialogues.jsonl")) |
| if n_dia != n_visits: |
| raise SystemExit(f"{w}: {n_dia} dialogues for {n_visits} visits") |
|
|
|
|
| def build_questions(patient, w): |
| """The sound questions of the build in axis-file order, and the (mechanism, state) of each axis-8 question.""" |
| dia = sorted(jl(w / "dialogues.jsonl"), key=lambda r: r["visit_number"]) |
| tok = {r["visit_number"]: len(enc().encode(" ".join(f"{t['speaker']}: {t['text']}" for t in r["turns"]), |
| disallowed_special=())) for r in dia} |
| date_of = {r["visit_number"]: r["date"] for r in dia} |
| ck = {c["qid"]: c for c in jl(w / "checkpoints.jsonl")} |
| va = {v["qid"]: v for v in jl(w / "validation.jsonl")} |
|
|
| out, axis8, dropped, seen_ax3 = [], {}, [], set() |
| for ax in range(1, 9): |
| for q in jl(w / f"questions_axis{ax}.jsonl"): |
| qid = q["qid"] |
| if qid not in va or qid not in ck: |
| raise SystemExit(f"{qid}: no validation or checkpoint row") |
| if not va[qid]["sound"]: |
| continue |
| if ax == 3: |
| key = json.dumps([q.get("options"), q.get("answer")], ensure_ascii=False, sort_keys=True) |
| if key in seen_ax3: |
| dropped.append((qid, "axis3_duplicate")) |
| continue |
| seen_ax3.add(key) |
| ev, ckv = q["evidence_visits"], ck[qid]["ckpt_visit"] |
| if ckv not in tok: |
| raise SystemExit(f"{qid}: checkpoint v{ckv} not in dialogues") |
| if ck[qid]["ckpt_date"] and str(ck[qid]["ckpt_date"]) != date_of[ckv]: |
| raise SystemExit(f"{qid}: checkpoint date {ck[qid]['ckpt_date']} differs from the dialogue date") |
| out_qid = f"{patient}/{qid.split('/', 1)[1]}" |
| if ax == 8: |
| spec = q["axis_specific"] |
| axis8[out_qid] = (spec["mechanism"], spec["state"]) |
| out.append({ |
| "qid": out_qid, "patient_id": patient, "axis": q["axis"], "question": q["question"], |
| "options": q.get("options"), "answer": q["answer"], "answer_type": q["answer_type"], |
| "explanation": q["provenance"] if ax in (1, 2) else " ".join(q["derivation"]["steps"]), |
| "checkpoint_visit": ckv, "checkpoint_date": ck[qid]["ckpt_date"], |
| "evidence_visits": ev, "grounding": q["grounding"], "n_evidence_visits": len(ev), |
| "evidence_context_tokens": sum(tok.get(v, 0) for v in ev), |
| "context_visits": sum(1 for v in tok if v <= ckv), |
| "context_tokens": sum(t for v, t in tok.items() if v <= ckv), |
| }) |
| return out, axis8, dropped |
|
|
|
|
| def spread(qs, k): |
| """k questions evenly spaced over the list sorted by the evidence-to-checkpoint gap.""" |
| if k <= 0: |
| return [] |
| if len(qs) <= k: |
| return qs |
| s = sorted(qs, key=lambda q: q["checkpoint_visit"] - q["evidence_visits"][0]) |
| idx = sorted(set(round(i * (len(s) - 1) / (k - 1)) for i in range(k))) if k > 1 else [0] |
| return [s[i] for i in idx] |
|
|
|
|
| def select_axis8(rows, axis8): |
| """At most AXIS8_PER_MECHANISM questions per mechanism, evenly spaced by evidence-to-checkpoint distance; for |
| deadlines, not-yet-due questions first, then overdue ones.""" |
| keep = set() |
| for mech, cap in AXIS8_PER_MECHANISM.items(): |
| taken = [] |
| for state in AXIS8_STATE_ORDER.get(mech, [None]): |
| pool = [r for r in rows if r["qid"] in axis8 and axis8[r["qid"]][0] == mech |
| and state in (None, axis8[r["qid"]][1])] |
| taken += spread(pool, cap - len(taken)) |
| keep.update(r["qid"] for r in taken) |
| return [r for r in rows if r["qid"] not in axis8 or r["qid"] in keep] |
|
|
|
|
| def export(w, dest, patient): |
| """Write the per-patient files of the build in w into the new directory dest, as patient `patient` (P<nn>).""" |
| if not re.fullmatch(r"P\d{2}", str(patient or ""), re.ASCII): |
| raise SystemExit(f"patient id {patient!r} is not of the form P<nn>") |
| check_build(w) |
| rows, axis8, dropped = build_questions(patient, w) |
| unknown = sorted({m for m, _ in axis8.values()} - set(AXIS8_PER_MECHANISM)) |
| if unknown: |
| raise SystemExit(f"axis-8 mechanisms without a cap: {unknown}") |
| kept = select_axis8(rows, axis8) |
| dest.mkdir(parents=True) |
| for src, dst in COPIES: |
| records = jl(w / src) |
| for r in records: |
| if src == "dialogues.jsonl": |
| r["meta"] = {"atom_index": (r.get("meta") or {}).get("atom_index", {})} |
| if src == "event_chain.jsonl" and "patient_id" in r: |
| r["patient_id"] = patient |
| if src == "evidence_chain.jsonl" and str(r.get("batch_id", "")).startswith("synthetic-"): |
| r["entity"] = None |
| with open(dest / dst, "w", encoding="utf-8") as fh: |
| for r in records: |
| fh.write(json.dumps(r, ensure_ascii=False) + "\n") |
| with open(dest / "questions.jsonl", "w", encoding="utf-8") as fh: |
| for r in kept: |
| fh.write(json.dumps(r, ensure_ascii=False) + "\n") |
| return rows, kept, axis8, dropped |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Export one finished patient build.") |
| ap.add_argument("--work", required=True, help="the patient's build directory (build/<patient_id>)") |
| ap.add_argument("--out", required=True, help="output directory (must not exist)") |
| ap.add_argument("--patient-id", required=True, |
| help="identifier written as patient_id and as the qid prefix, of the form P<nn>") |
| a = ap.parse_args() |
| dest = Path(a.out) |
| if dest.exists(): |
| raise SystemExit(f"{dest} already exists; refusing to overwrite") |
| rows, kept, axis8, dropped = export(Path(a.work), dest, a.patient_id) |
| per_axis = Counter(r["axis"][0] for r in kept) |
| print(f"-> {dest}") |
| print(f"sound questions: {len(rows)} axis-3 duplicates dropped: {len(dropped)}") |
| print(f"axis 8: {len(axis8)} sound, kept per mechanism " |
| f"{dict(Counter(axis8[r['qid']][0] for r in kept if r['qid'] in axis8))}") |
| print(f"questions.jsonl: {len(kept)} per axis {dict(sorted(per_axis.items()))}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
