File size: 8,699 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | #!/usr/bin/env python3
"""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()
|