| |
| """Stage F: two-agent role-play dialogue generation over the event chain. |
| |
| Per visit, the doctor agent sees the visit's atoms (flagged NEW or CARRIED), up to two prior visits in the |
| same department and a checklist of atoms to convey; the patient agent sees its own complaints, prior |
| same-department visits and an optional cross-department recap. The agents alternate as separate LLM |
| calls sharing the transcript until every atom is conveyed or the turn limit is reached; short visits get a |
| minimum length and every dialogue ends with a closing turn. The prompts are DIALOGUE_RULES, DOCTOR_SYSTEM and |
| PATIENT_SYSTEM in pipeline.prompts. |
| |
| generate --all --workers N --out dialogues.jsonl |
| count-empty number of visits whose dialogue has no turns |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import re |
| import threading |
| import unicodedata |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| from pipeline import prompts |
| from pipeline.common import BASE, FIELDS, model |
| from pipeline.llm import ChatClient |
|
|
| DATA = BASE |
| MODEL = model("dialogue") |
| SHARED = prompts.DIALOGUE_RULES |
| DOC_SYS = SHARED + "\n\n" + prompts.DOCTOR_SYSTEM |
| PAT_SYS = SHARED + "\n\n" + prompts.PATIENT_SYSTEM |
|
|
| _STOP = {"the", "of", "a", "an", "and", "or", "to", "with", "for", "no", "on", "at", "in", "by", |
| "is", "was", "are", "s", "p", "his", "her", "you", "your"} |
| _ALIASES = { |
| "lbp": "low back pain", "oa": "osteoarthritis", "mmt": "muscle strength", |
| "pt": "physical therapy", "ot": "occupational therapy", "st": "speech therapy", "nsaid": "painkiller", |
| "vas": "pain score", "r/o": "rule out", "bil": "bilateral", "gerd": "reflux", |
| "hivd": "disc herniation", "ar": "allergic rhinitis", "u/e": "upper limb", "ldl": "cholesterol", |
| "ck": "muscle enzyme", "rom": "range of motion", "bw": "weight", "bp": "blood pressure", |
| } |
| SMALL_TALK_TOPICS = ("traffic", "waiting time", "the hospital being busy", "feeling mildly nervous", |
| "coming from work", "the weather", "being worried while waiting for results") |
| PATIENT_SMALL_TALK_TOPICS = SMALL_TALK_TOPICS + ("getting to the clinic",) |
| _NUMBER = re.compile(r"\d+(?:\.\d+)?") |
| _VITAL_NAMES = {"BP": ("bp", "blood pressure"), "pulse": ("pulse", "hr", "heart rate", "pr"), |
| "weight_kg": ("bw", "weight", "wt"), "height_cm": ("height", "ht", "bh")} |
|
|
|
|
| def norm(s: str) -> str: |
| s = str(s).lower() |
| for k, v in _ALIASES.items(): |
| s = re.sub(rf"(?<![a-z]){re.escape(k)}(?![a-z])", f"{k} {v}", s) |
| return re.sub(r"\s+", " ", s) |
|
|
|
|
| def toks(s: str) -> set: |
| return {t for t in re.findall(r"[a-z0-9]+", norm(s)) if len(t) > 2 and t not in _STOP} |
|
|
|
|
| def write_jsonl(path, records): |
| with open(path, "w", encoding="utf-8") as fh: |
| for r in records: |
| fh.write(json.dumps(r, ensure_ascii=False) + "\n") |
| return len(records) |
|
|
|
|
| def drop_caseless_letters(s): |
| """The text without letters that have no case (Unicode category Lo), with empty parentheses and repeated spaces |
| removed.""" |
| s = "".join(ch for ch in str(s) if unicodedata.category(ch) != "Lo") |
| s = re.sub(r"\(\s*\)", "", s) |
| return re.sub(r"\s{2,}", " ", re.sub(r"\s+\)", ")", re.sub(r"\(\s+", "(", s))).strip() |
|
|
|
|
| recs, tags = {}, {} |
|
|
|
|
| def load_chain(): |
| """Read the patient's event chain and change tags.""" |
| recs.update({r["visit_number"]: r for r in (json.loads(l) for l in open(DATA / "event_chain.jsonl"))}) |
| tags.update({r["visit_number"]: r for r in (json.loads(l) for l in open(DATA / "tags.jsonl"))}) |
|
|
|
|
| def atoms_of(rec): |
| out = [] |
| for f in FIELDS: |
| for i, a in enumerate(rec.get(f) or []): |
| out.append({"atom_id": f"v{rec['visit_number']}:{f}:{i}", "field": f, "text": drop_caseless_letters(a)}) |
| for k, val in (rec.get("vitals") or {}).items(): |
| if in_doctor_notes(rec, k, val): |
| continue |
| out.append({"atom_id": f"v{rec['visit_number']}:vitals:{k}", "field": "vitals", "text": f"{k}={val}"}) |
| return [a for a in out if a["text"]] |
|
|
|
|
| def in_doctor_notes(rec, key, val): |
| """Whether a doctor_notes atom of the visit already records this vital sign: it names the sign and gives the |
| same reading.""" |
| names = _VITAL_NAMES.get(key, (key.lower(),)) |
| value = re.compile(rf"(?<![\d.]){re.escape(str(val))}(?![\d.])") |
| for a in rec.get("doctor_notes") or []: |
| t = str(a).lower() |
| if value.search(t) and any(re.search(rf"(?<![a-z]){re.escape(n)}(?![a-z])", t) for n in names): |
| return True |
| return False |
|
|
|
|
| def flag_novelty(vn, atoms): |
| """NEW vs CARRIED: an atom is CARRIED when an atom of the same field in a prior same-department visit shares |
| at least 70% of its words and has every number the atom states; a number that is new makes it NEW.""" |
| prev = tags.get(vn, {}).get("prev_visits", []) or [] |
| prior = {} |
| for pv in prev: |
| for pa in atoms_of(recs[pv]) if pv in recs else []: |
| prior.setdefault(pa["field"], []).append((toks(pa["text"]), set(_NUMBER.findall(pa["text"])))) |
| for a in atoms: |
| at, an = toks(a["text"]), set(_NUMBER.findall(a["text"])) |
| a["novelty"] = "CARRIED" if at and any(len(at & pt) / len(at) >= 0.7 and an <= pn |
| for pt, pn in prior.get(a["field"], [])) else "NEW" |
|
|
|
|
| def window_txt(vn): |
| prev = tags.get(vn, {}).get("prev_visits", []) or [] |
| if not prev: |
| return " none: this is the patient's first visit to this department." |
| out = [] |
| for tag, pv in zip(("most recent", "one before"), reversed(prev[-2:])): |
| if pv in recs: |
| out.append(f" [{tag}: v{pv} {recs[pv]['date']}] " + " | ".join(a["text"] for a in atoms_of(recs[pv]))) |
| return "\n".join(out) |
|
|
|
|
| def doctor_blob(rec, atoms): |
| department, date = drop_caseless_letters(rec["department"]), rec["date"] |
| p = [f"You are a {department} physician. Today's date is {date}.", |
| "\n<current_visit_facts>: today's full record, one fact per line with its novelty flag. Put the atom " |
| "id in covers."] |
| for a in atoms: |
| fl = "NEW " if a["novelty"] == "NEW" else "CARRIED" |
| atom_id, field, text = a["atom_id"], a["field"], a["text"] |
| p.append(f" {fl} | {atom_id} [{field}] {text}") |
| p.append("\n<recent_visits_here>: your previous visits in this department, the source for recaps.") |
| p.append(window_txt(rec["visit_number"])) |
| p.append("\n<must_convey_this_visit>: every atom id above must be conveyed by the end of the visit.") |
| return "\n".join(p) |
|
|
|
|
| def patient_blob(rec, atoms): |
| p = ["You are the patient. You do not know today's exam findings, results, diagnosis or plan until the doctor " |
| "tells you.", "\n<your_reasons_today>: your own complaints for today, in your own words. You may bring " |
| "these up."] |
| reasons = [(a["atom_id"], a["text"]) for a in atoms if a["field"] == "patient_self_description"] |
| for atom_id, text in reasons: |
| p.append(f" {atom_id} {text}") |
| if not reasons: |
| p.append(" none: a routine follow-up with no new complaint of your own") |
| p.append("\n<your_recent_visits_here>: your previous visits in this department. You lived through these.") |
| p.append(window_txt(rec["visit_number"])) |
| return "\n".join(p) |
|
|
|
|
| def cross_dept_recap(rec): |
| """The patient's most recent visit to another department, as the <cross_dept_recap> line of the patient's |
| opening <turn_task>, or ''.""" |
| cd = tags.get(rec["visit_number"], {}).get("prev_cross_dept_visit") |
| if not (cd and cd in recs): |
| return "" |
| other_department, other_date = drop_caseless_letters(recs[cd]["department"]), recs[cd]["date"] |
| return (f"\n<cross_dept_recap>: you recently saw a different department, {other_department}, on {other_date}. " |
| "You may briefly mention that you went there; you know none of its findings or results.") |
|
|
|
|
| def agent_turn(client, system, blob, transcript, steering): |
| conversation = "\n".join(f"{x['speaker']}: {x['text']}" for x in transcript) or "(the visit is just beginning)" |
| user = (f"{blob}\n\n=== Conversation so far ===\n{conversation}\n\n=== Your turn: one short spoken turn ===\n" |
| f"{steering}") |
| o = client.complete_json(system, user) |
| if not isinstance(o, dict) or "text" not in o: |
| return {"text": str(o)[:200], "covers": [], "type": "clinical", "end_ok": False} |
| cov = o.get("covers") or [] |
| return {"text": drop_caseless_letters(o.get("text", "")), |
| "covers": [str(c) for c in cov] if isinstance(cov, list) else [], |
| "type": str(o.get("type", "clinical")), "end_ok": bool(o.get("end_ok", False))} |
|
|
|
|
| WRAP = " <steering>: this is the final wrap-up turn; you may set end_ok=true." |
| NOWRAP = " <steering>: this is not the wrap-up turn; set end_ok=false." |
|
|
|
|
| def generate_dialogue(vn, client): |
| rec = recs[vn] |
| atoms = atoms_of(rec) |
| flag_novelty(vn, atoms) |
| ids = {a["atom_id"] for a in atoms} |
| a_by_id = {a["atom_id"]: a for a in atoms} |
| n_new = sum(1 for a in atoms if a["novelty"] == "NEW") |
| small_talk_budget = int(tags.get(vn, {}).get("noise_score", 0) or 0) |
| turns_target = max(5, min(48, 2 + 2 * min(n_new, 6) + (len(atoms) - n_new + 3) // 4 + small_talk_budget)) |
| d_blob, p_blob = doctor_blob(rec, atoms), patient_blob(rec, atoms) |
| transcript, covered = [], set() |
| small_talk_left = [small_talk_budget] |
| used_topics = set() |
|
|
| def small_talk_hint(patient=False): |
| """The small-talk sentence of a turn: the speaker's next topic not yet used in the visit, starting at a |
| position set by the visit number, while the visit's budget lasts.""" |
| if small_talk_left[0] > 0: |
| small_talk_left[0] -= 1 |
| topics = PATIENT_SMALL_TALK_TOPICS if patient else SMALL_TALK_TOPICS |
| order = [topics[(vn + k) % len(topics)] for k in range(len(topics))] |
| t = next((x for x in order if x not in used_topics), order[0]) |
| used_topics.add(t) |
| return (f" You may add one brief bit of small talk about {t}, an allowed topic. Keep it natural, and do " |
| "not repeat a small-talk topic already used.") |
| return " No small talk this turn." |
|
|
| def add(speaker, turn): |
| covered.update(c for c in turn["covers"] if c in ids) |
| transcript.append({"speaker": speaker, "type": turn.get("type", "clinical"), |
| "text": turn["text"], "covers": [c for c in turn["covers"] if c in ids]}) |
|
|
| def uncovered(): |
| u = [a for a in atoms if a["atom_id"] not in covered] |
| return sorted(u, key=lambda a: (a["novelty"] != "NEW", a["field"] != "doctor_extra_notes")) |
|
|
| def reasons(): |
| return "; ".join(f"[{a['atom_id']}] {a['text']}" for a in uncovered() |
| if a["field"] == "patient_self_description") |
|
|
| def react(): |
| """The patient's turn after a doctor's clinical turn.""" |
| t = ("<turn_task>: phase=react. React to what the doctor just said, with a lay reaction or a short " |
| "question.") |
| own = reasons() |
| if own: |
| t += f" If you have not said them yet, you may voice these reasons and list their ids in covers: {own}." |
| add("Patient", agent_turn(client, PAT_SYS, p_blob, transcript, t + small_talk_hint(patient=True) + "\n" |
| + NOWRAP)) |
|
|
| def task(batch, phase, wrap=False): |
| facts = "; ".join(f"[{a['atom_id']}] {a['text']}" for a in batch) |
| t = f"<turn_task>: phase={phase}." |
| if facts: |
| t += f" Convey only these facts in this turn, and list their ids in covers: {facts}." |
| return t + small_talk_hint() + ("\n" + (WRAP if wrap else NOWRAP)) |
|
|
| own = reasons() |
| add("Patient", agent_turn(client, PAT_SYS, p_blob, transcript, |
| "<turn_task>: phase=intake. Open the visit: say in your own everyday words why you are here today." |
| + (f" Voice your reasons for today and list their ids in covers: {own}." if own else "") |
| + small_talk_hint(patient=True) + cross_dept_recap(rec) + "\n" + NOWRAP)) |
| add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript, |
| task([a for a in uncovered() if a["field"] != "patient_self_description"][:1], "intake"))) |
|
|
| guard = 0 |
| while uncovered() and len(transcript) < turns_target + 8 and guard < 40: |
| guard += 1 |
| u = uncovered() |
| batch = u[:1] if u[0]["novelty"] == "NEW" else u[:3] |
| add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript, task(batch, "clinical"))) |
| if uncovered(): |
| react() |
|
|
| n_recap = 0 |
| while uncovered() and n_recap < 6: |
| n_recap += 1 |
| add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript, |
| task(uncovered()[:5], "recap"))) |
|
|
| while len(transcript) < 4: |
| add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript, |
| "<turn_task>: phase=plan. Briefly restate today's plan or next step in plain words, using only plans " |
| "already in <current_visit_facts>. Do not add new instructions, tests or follow-up intervals." |
| + small_talk_hint() |
| + "\n" + NOWRAP)) |
| add("Doctor", agent_turn(client, DOC_SYS, d_blob, transcript, |
| "<turn_task>: phase=closing. Wrap the visit up warmly in one short line, without new instructions, " |
| "tests or follow-up intervals." + WRAP)) |
|
|
| return {"visit_number": vn, "date": rec["date"], "department": rec["department"], "session_id": f"v{vn}", |
| "turns": transcript, |
| "meta": {"tag": tags.get(vn, {}).get("tag"), "noise_score": small_talk_budget, |
| "turns_target": turns_target, "n_atoms": len(atoms), "n_new": n_new, "engine": MODEL, |
| "atom_index": {a["atom_id"]: {"text": a["text"], "field": a["field"], "novelty": a["novelty"]} |
| for a in atoms}, |
| "coverage": {"ledger": len(ids), "covered": len(covered), |
| "uncovered": [a_by_id[i]["atom_id"] for i in ids - covered]}}} |
|
|
|
|
| def generate(args) -> None: |
| """Generate the dialogues of the selected visits, resuming from any existing output.""" |
| load_chain() |
| vns = sorted(recs) if args.all else [int(x) for x in args.visits.split(",") if x.strip()] |
| outpath = Path(args.out) if os.path.isabs(args.out) else DATA / args.out |
| done = {} |
| if outpath.exists() and not args.fresh: |
| done = {r["visit_number"]: r for r in (json.loads(l) for l in open(outpath) if l.strip()) if r.get("turns")} |
| todo = [v for v in vns if v not in done] |
| print(f"visits: {len(vns)} total | already done {len(done)} | to generate {len(todo)}") |
| client = ChatClient(model=MODEL, reasoning_effort="low", max_completion_tokens=800) |
| lock = threading.Lock() |
| n_ok = [0] |
| n_err = [0] |
| fh = open(outpath, "a", encoding="utf-8") |
|
|
| def one(vn): |
| try: |
| r = generate_dialogue(vn, client) |
| except Exception as e: |
| r = {"visit_number": vn, "error": str(e)[:150], "turns": []} |
| with lock: |
| fh.write(json.dumps(r, ensure_ascii=False) + "\n") |
| fh.flush() |
| (n_ok if r.get("turns") else n_err)[0] += 1 |
| if (n_ok[0] + n_err[0]) % 25 == 0: |
| print(f" progress {n_ok[0]+n_err[0]}/{len(todo)} (ok {n_ok[0]}, err {n_err[0]})", flush=True) |
| return r |
|
|
| with ThreadPoolExecutor(max_workers=args.workers) as ex: |
| list(ex.map(one, todo)) |
| fh.close() |
| allr = {r["visit_number"]: r for r in (json.loads(l) for l in open(outpath) if l.strip())} |
| out = sorted(allr.values(), key=lambda r: r["visit_number"]) |
| write_jsonl(outpath, out) |
| ok = [r for r in out if r.get("turns")] |
| err = [r for r in out if not r.get("turns")] |
| tc = [len(r["turns"]) for r in ok] |
| cov = [r["meta"]["coverage"] for r in ok if "meta" in r] |
| print(f"\ndialogues: {len(ok)} ok, {len(err)} error/empty turns min/med/max = " |
| f"{min(tc)}/{sorted(tc)[len(tc)//2]}/{max(tc)}") |
| print(f"visits <5 turns: {sum(1 for t in tc if t < 5)}") |
| print(f"atom coverage: {sum(c['covered'] for c in cov)}/{sum(c['ledger'] for c in cov)} " |
| f"({100*sum(c['covered'] for c in cov)/max(1,sum(c['ledger'] for c in cov)):.1f}%)") |
| if err: |
| print("visits with errors (retried on the next run):", [r["visit_number"] for r in err][:30]) |
| print(f"wrote {args.out}") |
|
|
|
|
| def count_empty() -> None: |
| """Print the number of visits whose dialogue has no turns.""" |
| print(sum(1 for l in open(BASE / "dialogues.jsonl") if l.strip() and not json.loads(l).get("turns"))) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Stage F: dialogue generation.") |
| sp = ap.add_subparsers(dest="cmd", required=True) |
| g = sp.add_parser("generate") |
| g.add_argument("--all", action="store_true") |
| g.add_argument("--visits", default="") |
| g.add_argument("--workers", type=int, default=12) |
| g.add_argument("--out", default="dialogues.jsonl") |
| g.add_argument("--fresh", action="store_true", help="ignore any existing output and regenerate all") |
| sp.add_parser("count-empty") |
| a = ap.parse_args() |
| generate(a) if a.cmd == "generate" else count_empty() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
