| """Locations, identifiers, models, question counts, progress reporting and the term glossary shared by the |
| stages. |
| |
| build_patient.sh sets MEMEBENCH_WORK (the patient's working directory), MEMEBENCH_PATIENT (the patient's |
| number in the visit table), the per-patient facts and the model names. The visit table (the synthetic |
| timeline) and the term glossary are not distributed (MEMEBENCH_VISIT_TABLE, MEMEBENCH_GLOSSARY). |
| """ |
| from __future__ import annotations |
|
|
| import io |
| import json |
| import os |
| import re |
| import time |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| VISIT_TABLE = Path(os.environ.get("MEMEBENCH_VISIT_TABLE", "<VISIT_TABLE>")) |
| GLOSSARY = Path(os.environ.get("MEMEBENCH_GLOSSARY", "<TERM_GLOSSARY>")) |
| SEX_MAP = {"0": "female", "1": "male"} |
| DEFAULT_PDESC = "longitudinal multi-department history" |
| FIELDS = ["diagnosis", "patient_self_description", "tests_and_medication", "doctor_notes", |
| "doctor_plans", "health_advice", "doctor_extra_notes"] |
|
|
| |
| |
| PER_AXIS_MINIMUM = {3: 3, 5: 10, 6: 10, 7: 10} |
| AXIS1_TARGET = "8-10" |
| AXIS1_FORM_MIX = "5 form=range + 3 form=never_again" |
| AXIS2_TARGET = "4" |
|
|
|
|
| def _required(name: str) -> str: |
| value = os.environ.get(name, "") |
| if not value or value.startswith("<"): |
| raise SystemExit(f"{name} is not set; run the stages through build_patient.sh") |
| return value |
|
|
|
|
| BASE = Path(_required("MEMEBENCH_WORK")) |
|
|
|
|
| def patient_id() -> int: |
| """The patient's number in the visit table.""" |
| return int(_required("MEMEBENCH_PATIENT")) |
|
|
|
|
| def model(role: str) -> str: |
| """The model configured for a role: extraction, generation or dialogue.""" |
| return _required(f"MEMEBENCH_{role.upper()}_MODEL") |
|
|
|
|
| def no_memory_exempt_axes() -> set: |
| """Axis numbers whose questions are kept regardless of the no-memory probe (MEMEBENCH_NO_MEMORY_EXEMPT_AXES).""" |
| if "MEMEBENCH_NO_MEMORY_EXEMPT_AXES" not in os.environ: |
| raise SystemExit("MEMEBENCH_NO_MEMORY_EXEMPT_AXES is not set; run the stages through build_patient.sh") |
| raw = os.environ["MEMEBENCH_NO_MEMORY_EXEMPT_AXES"] |
| return {int(x) for x in raw.split(",") if x.strip()} |
|
|
|
|
| def _sex_str() -> str: |
| """The patient's sex from the visit table (0 = female, 1 = male), for the synthetic-atom prompts.""" |
| patient = patient_id() |
| try: |
| import pandas as pd |
| df = pd.read_parquet(VISIT_TABLE, columns=["record_id", "sex"]) |
| v = str(df[df["record_id"] == patient]["sex"].dropna().iloc[0]).strip() |
| except Exception as e: |
| raise RuntimeError(f"cannot read the patient's sex from {VISIT_TABLE}: {e}") from e |
| if v not in SEX_MAP: |
| raise RuntimeError(f"sex value {v!r} of the patient in {VISIT_TABLE} is not one of {sorted(SEX_MAP)}") |
| return SEX_MAP[v] |
|
|
|
|
| @lru_cache(maxsize=None) |
| def patient_line() -> str: |
| """The one-line patient description given to the synthetic-atom prompts.""" |
| pdesc = os.environ.get("MEMEBENCH_PDESC", DEFAULT_PDESC) |
| return f"Patient ({_sex_str()}): {pdesc}." |
|
|
|
|
| def norm_value(v) -> str: |
| """A timeline value for comparison: lower case, parentheses and a leading 'history of' removed.""" |
| v = re.sub(r"\s*\([^)]*\)", "", str(v or "")).strip().lower() |
| v = re.sub(r"^(?:history of|hx of|h/o|known)\s+", "", v) |
| return re.sub(r"\s+", " ", v).strip(" .;,") |
|
|
|
|
| def other_department_segments(timelines, entity, department): |
| """Segments of the timelines that other departments keep for the same entity (name compared ignoring case and |
| spacing).""" |
| key = re.sub(r"\s+", " ", str(entity or "").strip().lower()) |
| return [s for t in timelines if t.get("department") != department |
| and re.sub(r"\s+", " ", str(t.get("entity") or "").strip().lower()) == key |
| for s in t.get("segments") or []] |
|
|
|
|
| def has_text(q, key="question") -> bool: |
| """Whether a model reply carries a non-empty string under `key`.""" |
| return isinstance(q, dict) and isinstance(q.get(key), str) and bool(q[key].strip()) |
|
|
|
|
| def lists_options(question, options) -> bool: |
| """Whether a question text already lists its options: every option label appears as an item marker, '<label>.' |
| or '<label>)', at the start of a line or after a space, '/' or '('.""" |
| labels = [str(o.get("label")) for o in options or [] if isinstance(o, dict) and o.get("label")] |
| return bool(labels) and all(re.search(rf"(?:^|[\s/(]){re.escape(x)}[.)]\s", str(question or "")) for x in labels) |
|
|
|
|
| def patient_context(chain) -> str: |
| """The record's departments, 15 most frequent diagnoses and 20 most frequent tests and medications, for the |
| synthetic-atom prompts.""" |
| from collections import Counter |
| depts = sorted({r["department"] for r in chain}) |
| dx = Counter(str(a)[:60] for r in chain for a in r.get("diagnosis") or []) |
| tm = Counter(str(a)[:60] for r in chain for a in r.get("tests_and_medication") or []) |
| return ("Departments: " + "; ".join(depts) + "\n" |
| "Most frequent diagnoses: " + "; ".join(d for d, _ in dx.most_common(15)) + "\n" |
| "Most frequent tests and medications: " + "; ".join(t for t, _ in tm.most_common(20))) |
|
|
|
|
| |
| def _hms(s: float) -> str: |
| s = int(max(0, s)) |
| if s < 60: |
| return f"{s}s" |
| if s < 3600: |
| return f"{s//60}m{s%60:02d}s" |
| return f"{s//3600}h{(s%3600)//60:02d}m" |
|
|
|
|
| class Progress: |
| def __init__(self, label: str, total: int, out: str | None = None, |
| every: int = 10, secs: float = 15.0): |
| self.label, self.total, self.out = label, total, Path(out) if out else None |
| self.every, self.secs = every, secs |
| self.n = self.ok = self.fail = 0 |
| self.t0 = self.last_t = time.time() |
| self.last_n = 0 |
| self.done = False |
|
|
| def __enter__(self): |
| self._say(force=True) |
| return self |
|
|
| def __exit__(self, *exc): |
| self.done = True |
| self._say(force=True, final=True) |
| return False |
|
|
| def tick(self, ok: bool = True, n: int = 1): |
| self.n += n |
| if ok: |
| self.ok += n |
| else: |
| self.fail += n |
| now = time.time() |
| if (self.n - self.last_n) >= self.every or (now - self.last_t) >= self.secs: |
| self._say() |
|
|
| def _state(self, final=False): |
| el = time.time() - self.t0 |
| rate = self.n / el if el > 0 else 0 |
| left = (self.total - self.n) / rate if rate > 0 and not final else 0 |
| return { |
| "label": self.label, "done": self.done, |
| "n": self.n, "total": self.total, |
| "pct": round(100 * self.n / self.total, 1) if self.total else 0, |
| "ok": self.ok, "fail": self.fail, |
| "elapsed_s": round(el, 1), "rate_per_s": round(rate, 2), |
| "eta_s": round(left, 1) if not final else 0, |
| "updated": time.strftime("%H:%M:%S"), |
| } |
|
|
| def _say(self, force: bool = False, final: bool = False): |
| st = self._state(final) |
| self.last_t, self.last_n = time.time(), self.n |
| bar_w = 24 |
| filled = int(bar_w * st["pct"] / 100) |
| bar = "#" * filled + "." * (bar_w - filled) |
| tail = f", {st['fail']} FAILED" if st["fail"] else "" |
| eta = "done" if final else f"eta {_hms(st['eta_s'])}" |
| print(f" [{bar}] {st['pct']:>5.1f}% {st['n']:>4}/{st['total']} " |
| f"{st['rate_per_s']:.2f}/s {_hms(st['elapsed_s'])} elapsed, {eta}{tail}", |
| flush=True) |
| if self.out: |
| self.out.parent.mkdir(parents=True, exist_ok=True) |
| tmp = self.out.with_suffix(".tmp") |
| tmp.write_text(json.dumps(st, indent=1)) |
| os.replace(tmp, self.out) |
|
|
|
|
| |
| |
| def load_glossary(): |
| g = {} |
| for ln in io.open(GLOSSARY, encoding="utf-8"): |
| ln = ln.rstrip("\n") |
| if ln.strip(): |
| src, norm = ln.split("\t") |
| g[src] = norm |
| return g |
|
|
|
|
| def glossary_re(g): |
| keys = sorted(g, key=len, reverse=True) |
| return re.compile("|".join(re.escape(k) for k in keys)) |
|
|
|
|
| def apply_glossary(text, g, pat): |
| return pat.sub(lambda m: f"{g[m.group(0)]} ({m.group(0)})", text) |
|
|
