File size: 6,498 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 | """Dataset loading: patient registry, visit sessions, questions and checkpoints.
Each patient directory holds `dialogues.jsonl` (one record per visit) and `questions.jsonl`
(one record per question, carrying its checkpoint visit). Sessions are sorted by
(date, visit_number) and records without turns are dropped. A visit is rendered as
`=== Visit v{n} 路 {date} 路 {department} ===` followed by `Speaker: text` lines; annotation
fields (`meta`, per-turn `covers`) never enter memory text. A question's checkpoint must
name an existing visit, and its checkpoint date must equal that visit's date.
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
def read_jsonl(path) -> list[dict]:
return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]
def load_sessions(path) -> list[dict]:
"""Return non-empty sessions sorted by (date, visit_number), keeping only the fields memory may see."""
rows = read_jsonl(path)
rows = [r for r in rows if r.get("turns")]
rows.sort(key=lambda r: (r["date"], r["visit_number"]))
return [{"visit_number": r["visit_number"], "date": r["date"],
"department": r["department"], "session_id": r.get("session_id"),
"turns": [{"speaker": t["speaker"], "text": t["text"]} for t in r["turns"]]}
for r in rows]
def render_dialogue(rec: dict) -> str:
return "\n".join(f"{t['speaker']}: {t['text']}" for t in rec["turns"])
def render_visit(rec: dict) -> str:
return f"=== Visit v{rec['visit_number']} 路 {rec['date']} 路 {rec['department']} ===\n" + render_dialogue(rec)
@dataclass
class Question:
qid: int
raw: dict
ckpt_visit: int = -1
ckpt_date: str = ""
@property
def axis(self) -> str:
return str(self.raw.get("axis", ""))
@property
def answer_type(self) -> str:
return str(self.raw.get("answer_type", ""))
@property
def question(self) -> str:
return self.raw["question"]
@property
def grounding(self) -> str:
return str(self.raw.get("grounding", "record"))
@property
def evidence_visits(self) -> list[int]:
return list(self.raw.get("evidence_visits", []))
def load_questions(questions_file, sessions: list[dict]) -> list[Question]:
"""Load a patient's questions, each with its checkpoint visit and date."""
records = read_jsonl(questions_file)
date_by_visit = {s["visit_number"]: s["date"] for s in sessions}
questions, problems = [], []
for i, r in enumerate(records):
ck = r.get("checkpoint_visit")
if ck is None:
problems.append(f"qid {r.get('qid')}: no checkpoint_visit")
continue
ck = int(ck)
if ck not in date_by_visit:
problems.append(f"qid {r.get('qid')}: checkpoint_visit v{ck} not in dialogues")
continue
stated = r.get("checkpoint_date")
if stated and str(stated) != date_by_visit[ck]:
problems.append(f"qid {r.get('qid')}: checkpoint_date {stated} != v{ck} date {date_by_visit[ck]}")
continue
questions.append(Question(qid=i, raw=r, ckpt_visit=ck,
ckpt_date=str(stated or date_by_visit[ck])))
if problems:
raise SystemExit("questions with checkpoints failed validation:\n " + "\n ".join(problems))
return questions
QUESTIONS_FILE = "questions.jsonl"
@dataclass
class Patient:
"""One patient: give a `dir` and let `resolve()` find the files inside it, or set the paths explicitly."""
name: str
patient_id: str
dialogues: str = ""
questions_dir: str = ""
dir: str = ""
questions_file: str = ""
def resolve(self) -> bool:
"""Fill unset paths from `dir`; return whether the questions file exists."""
d = Path(self.dir) if self.dir else None
if d and not self.dialogues and (d / "dialogues.jsonl").exists():
self.dialogues = str(d / "dialogues.jsonl")
if d and not self.questions_dir:
self.questions_dir = str(d)
qfile = Path(self.questions_dir) / QUESTIONS_FILE if self.questions_dir else None
if qfile is None or not qfile.exists():
return False
self.questions_file = str(qfile)
return True
def load(self):
"""Return (sessions, questions) for this patient."""
found = self.resolve()
if not self.dialogues:
raise SystemExit(f"patient '{self.name}': no dialogues file found (dir={self.dir})")
if not found:
raise SystemExit(f"patient '{self.name}': no questions file in {self.questions_dir} "
f"(looked for {QUESTIONS_FILE})")
sessions = load_sessions(self.dialogues)
questions = load_questions(self.questions_file, sessions)
if not questions:
raise SystemExit(f"patient '{self.name}': no questions found in {self.questions_file}")
return sessions, questions
def load_patients(config_path, skipped: list | None = None) -> dict[str, Patient]:
"""Build the patient registry from a YAML file.
`patients` lists explicit entries; `patient_globs` lists globs of patient directories, each
registered under its directory name. Directories without a questions file are left out and,
when `skipped` is given, appended to it as (name, reason).
"""
import glob as _glob
import yaml
cfg = yaml.safe_load(open(config_path, encoding="utf-8")) or {}
out: dict[str, Patient] = {}
for p in cfg.get("patients") or []:
out[p["name"]] = Patient(
name=p["name"], patient_id=str(p["patient_id"]), dialogues=p.get("dialogues", ""),
questions_dir=p.get("questions_dir", ""), dir=p.get("dir", ""))
for entry in cfg.get("patient_globs") or []:
pattern, prefix = (entry, "") if isinstance(entry, str) else (entry["path"], entry.get("prefix", ""))
for path in sorted(_glob.glob(str(pattern))):
d = Path(path)
if not d.is_dir():
continue
pid = d.name
name = f"{prefix}{pid}"
if name in out:
continue
pat = Patient(name=name, patient_id=pid, dir=str(d))
if not pat.resolve():
if skipped is not None:
skipped.append((name, f"no {QUESTIONS_FILE} in {d}"))
continue
out[name] = pat
return out
|