diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index 8d9a59d21..742e69231 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -229,6 +229,7 @@ Available Datasets datasets/pyhealth.datasets.eICUDataset datasets/pyhealth.datasets.ISRUCDataset datasets/pyhealth.datasets.MIMICExtractDataset + datasets/pyhealth.datasets.MIMIC4NoteExtDIBHCDataset datasets/pyhealth.datasets.OMOPDataset datasets/pyhealth.datasets.DREAMTDataset datasets/pyhealth.datasets.SHHSDataset diff --git a/docs/api/datasets/pyhealth.datasets.MIMIC4NoteExtDIBHCDataset.rst b/docs/api/datasets/pyhealth.datasets.MIMIC4NoteExtDIBHCDataset.rst new file mode 100644 index 000000000..c079aabcb --- /dev/null +++ b/docs/api/datasets/pyhealth.datasets.MIMIC4NoteExtDIBHCDataset.rst @@ -0,0 +1,14 @@ +pyhealth.datasets.MIMIC4NoteExtDIBHCDataset +=================================== + +The Mimic4Note Extension DischargeInstructions BriefHospitalCourse dataset is used in this paper: https://arxiv.org/pdf/2402.15422 to create AI generated medical note summaries. + +.. autoclass:: pyhealth.datasets.MIMIC4NoteExtDIBHCDataset + :members: + :undoc-members: + :show-inheritance: + + + + + diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 23a4e06e5..78fa1499c 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -230,3 +230,4 @@ Available Tasks Mutation Pathogenicity (COSMIC) Cancer Survival Prediction (TCGA) Cancer Mutation Burden (TCGA) + MIMIC-IV Note Tasks diff --git a/docs/api/tasks/pyhealth.tasks.mimic4_note_ext_dibhc_tasks.rst b/docs/api/tasks/pyhealth.tasks.mimic4_note_ext_dibhc_tasks.rst new file mode 100644 index 000000000..2b2b878c6 --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.mimic4_note_ext_dibhc_tasks.rst @@ -0,0 +1,7 @@ +pyhealth.tasks.mimic4_note_ext_dibhc_tasks +=================================== + +.. automodule:: pyhealth.tasks.mimic4_note_ext_dibhc_tasks + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/examples/mimic4_note_ext_dibhc_hallucination_detection_logreg.py b/examples/mimic4_note_ext_dibhc_hallucination_detection_logreg.py new file mode 100644 index 000000000..ffd46dcdb --- /dev/null +++ b/examples/mimic4_note_ext_dibhc_hallucination_detection_logreg.py @@ -0,0 +1,1070 @@ +# -*- coding: utf-8 -*- +"""Hallucination detection ablation study on MIMIC-IV-Note summaries. + +Reproduces and extends Section 4.7 (Automatic Hallucination Detection) from: + + Hegselmann et al. "A Data-Centric Approach To Generate Faithful and + High Quality Patient Summaries with Large Language Models." CHIL 2024. + https://arxiv.org/abs/2402.15422 + +Part 1 Task demonstration: + Shows BHCSummarizationTask and HallucinationDetectionTask processing + synthetic patient records, printing input/output schemas and sample + outputs to illustrate how both tasks integrate with the PyHealth pipeline. + +Part 2 Ablation study (novel contribution): + Binary classification ablation comparing three feature configurations + for document-level hallucination detection using logistic regression. + This is a novel extension not in the original paper, which only + evaluates span-level detection using MedCat and GPT-4. + + Feature configurations: + + - Config A: TF-IDF only (bag-of-words over summary text) + - Config B: TF-IDF + lexical overlap features + - Config C: TF-IDF + overlap + structural (number mismatch) + +Experimental results (5-fold stratified CV, 100 synthetic samples): + + +-------------------------------------+-------+-------+-------+ + | Config | Prec | Rec | F1 | + +=====================================+=======+=======+=======+ + | Config A: TF-IDF only | 40.0% | 33.3% | 36.0% | + +-------------------------------------+-------+-------+-------+ + | Config B: TF-IDF + Overlap | 56.7% | 60.0% | 57.3% | + +-------------------------------------+-------+-------+-------+ + | Config C: TF-IDF + Overlap + Struct | 73.3% | 53.3% | 57.3% | + +-------------------------------------+-------+-------+-------+ + +Key findings: + + - Feature engineering matters more than model complexity on small + imbalanced datasets. With only ~12% positive samples, TF-IDF + alone achieves just 36% F1. + + - Adding lexical overlap features (Config B) jumps F1 from 36% + to 57.3% by flagging summary words absent from the BHC context + (the same grounding signal used by the paper's MedCat baseline). + + - Structural features (Config C) improve precision from 56.7% to + 73.3% without hurting F1, showing that number mismatch (wrong + dosages, wrong dates) reliably identifies unsupported facts. + + - These results mirror the paper's broader conclusion: hallucination + detection is fundamentally difficult. Even GPT-4 achieves only + 19-20% F1 at span level. Document-level binary classification with + grounding-aware features achieves higher F1 but on a coarser task. + +Usage: + + python examples/mimic4_note_ext_dibhc_hallucination_detection_logreg.py + +Requirements: + + pip install scikit-learn +""" + +import os +import sys +from unittest.mock import MagicMock + +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from pyhealth.tasks.mimic4_note_tasks import ( + BHCSummarizationTask, + HallucinationDetectionTask, +) + +try: + from sklearn.linear_model import LogisticRegression + from sklearn.feature_extraction.text import TfidfVectorizer + from sklearn.metrics import precision_recall_fscore_support + from sklearn.model_selection import StratifiedKFold + from sklearn.base import BaseEstimator, TransformerMixin + from sklearn.pipeline import Pipeline + import scipy.sparse as sp +except ImportError: + print("Install scikit-learn: pip install scikit-learn") + sys.exit(1) + + +# ----------------------------------------------------------------- +# Synthetic demo data +# Mirrors the structure of PhysioNet ann-pt-summ annotation files. +# All BHC text, summaries, and labels are fully fabricated. +# ---------------------------------------------------------------- + +_SYNTHETIC_SAMPLES = [ + ( + "Brief Hospital Course: Patient presented with chest pain. Started on IV" + "vancomycin and ceftriaxone. Infection improved and patient discharged.", + "You were admitted for chest pain. You received blood thinners.", + 1, + ), + ( + "Brief Hospital Course: Patient with hypertension admitted for pneumonia." + "Started on azithromycin. Afebrile at discharge.", + "You were treated with antibiotics for a lung infection.", + 0, + ), + ( + "Brief Hospital Course: Post-op day 2 after appendectomy. Pain controlled." + "Vital signs stable. Tolerating diet.", + "You had your appendix removed and recovered well.", + 0, + ), + ( + "Brief Hospital Course: Patient with atrial fibrillation. Rate controlled" + "with metoprolol. Anticoagulation continued.", + "You were admitted for an irregular heartbeat and given Tylenol 500mg twice" + "daily.", + 1, + ), + ( + "Brief Hospital Course: Diabetic patient, HbA1c 9.2. Insulin regimen" + "adjusted. Glucose controlled prior to discharge.", + "Your blood sugar was high and we started insulin therapy.", + 0, + ), + ( + "Brief Hospital Course: Patient with stroke admitted for left hemisphere" + "infarct. MRI confirmed.", + "You were admitted for a mild fracture of the left clavicle.", + 1, + ), + ( + "Brief Hospital Course: Patient with COPD exacerbation. Started on steroids" + "and bronchodilators. Improved.", + "You were treated for a flare-up of your lung condition.", + 0, + ), + ( + "Brief Hospital Course: Post-cardiac catheterization. Stent placed in LAD. No" + "complications.", + "You had a procedure to open a blocked artery in your heart.", + 0, + ), + ( + "Brief Hospital Course: Patient with sepsis from UTI. Started broad-spectrum" + "antibiotics. Blood cultures negative.", + "You had a serious infection and received strong antibiotics for 14 days.", + 1, + ), + ( + "Brief Hospital Course: Patient with AKI. Hydrated with IV fluids. Creatinine" + "improved to baseline.", + "Your kidneys were not working well and we gave you fluids to help.", + 0, + ), + ( + "Brief Hospital Course: Patient with DVT in left leg. Started on heparin" + "bridge to warfarin. No PE on imaging.", + "You were found to have a blood clot and were started on blood thinners.", + 0, + ), + ( + "Brief Hospital Course: GI bleed from gastric ulcer. EGD performed. No active" + "bleeding. PPI started.", + "You were admitted for bleeding in your stomach. You were given a stress test.", + 1, + ), + ( + "Brief Hospital Course: Heart failure exacerbation. Diuresed with IV" + "furosemide. BNP trending down.", + "You were given diuretics to remove excess fluid from your lungs.", + 0, + ), + ( + "Brief Hospital Course: Patient with meningitis. LP performed. Started on" + "ceftriaxone and vancomycin.", + "You were treated for an infection of the fluid surrounding your brain.", + 0, + ), + ( + "Brief Hospital Course: Acute pancreatitis. NPO and IV fluids. Lipase" + "trending down. Diet advanced.", + "You were admitted for inflammation of your pancreas and given IV fluids.", + 0, + ), + ( + "Brief Hospital Course: NSTEMI. Troponin peaked at 2.4. Cardiac cath showed" + "80% LAD stenosis. Stent placed.", + "You had a mild heart attack. You received physical therapy for two weeks.", + 1, + ), + ( + "Brief Hospital Course: Cellulitis of right leg. Treated with IV nafcillin." + "Erythema resolved.", + "You were treated for a skin infection of your right leg with antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Hypertensive urgency. BP 210/110. IV labetalol given." + "Oral meds resumed.", + "Your blood pressure was dangerously high and was treated with medications.", + 0, + ), + ( + "Brief Hospital Course: Pulmonary embolism found on CTA. Anticoagulation" + "initiated with heparin.", + "You were found to have a blood clot in your lungs.", + 0, + ), + ( + "Brief Hospital Course: Alcohol withdrawal. CIWA protocol initiated. Thiamine" + "and folate given.", + "You were admitted for alcohol withdrawal and given medications to help.", + 0, + ), + ( + "Brief Hospital Course: Community acquired pneumonia. Treated with" + "levofloxacin. O2 requirements resolved.", + "You were treated for pneumonia. You were also found to have a urinary tract" + "infection.", + 1, + ), + ( + "Brief Hospital Course: Hypoglycemia. Glucose 38 on arrival. D50 given." + "Insulin dose adjusted.", + "Your blood sugar dropped very low and we gave you sugar through your IV.", + 0, + ), + ( + "Brief Hospital Course: Hip fracture after fall. ORIF performed. Physical" + "therapy initiated.", + "You broke your hip and had surgery to repair it.", + 0, + ), + ( + "Brief Hospital Course: Acute cholecystitis. Laparoscopic cholecystectomy" + "performed without complication.", + "You had your gallbladder removed due to an infection.", + 0, + ), + ( + "Brief Hospital Course: Seizure. EEG performed. Keppra initiated. No further" + "events.", + "You were admitted for a seizure and started on anti-seizure medication.", + 0, + ), + ( + "Brief Hospital Course: Dehydration from vomiting. IV fluids given." + "Electrolytes normalized.", + "You were dehydrated and given IV fluids. You were also started on" + "chemotherapy.", + 1, + ), + ( + "Brief Hospital Course: Chest pain, ruled out ACS. Stress test negative." + "Discharged on aspirin.", + "You were monitored for chest pain and found not to have a heart attack.", + 0, + ), + ( + "Brief Hospital Course: Type 2 DM poorly controlled. Metformin increased." + "HbA1c 10.2.", + "Your diabetes was poorly controlled and your medications were adjusted.", + 0, + ), + ( + "Brief Hospital Course: Ischemic stroke. tPA administered. MRI showed small" + "left MCA infarct.", + "You had a stroke affecting the right side of your brain.", + 1, + ), + ( + "Brief Hospital Course: Asthma exacerbation. Nebulizers and steroids given." + "Peak flow improved.", + "You were treated for an asthma attack with steroids and breathing treatments.", + 0, + ), + ( + "Brief Hospital Course: Pyelonephritis. IV ceftriaxone started. Urine culture" + "grew E coli.", + "You were treated for a kidney infection with antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Syncope. Holter monitor placed. Echo normal." + "Orthostatics negative.", + "You fainted and we monitored your heart for abnormal rhythms.", + 0, + ), + ( + "Brief Hospital Course: Bowel obstruction. NGT placed. Conservative" + "management successful.", + "You had a blockage in your intestine that improved with conservative" + "treatment.", + 0, + ), + ( + "Brief Hospital Course: Anemia requiring transfusion. Two units pRBC given." + "Hgb improved to 9.", + "You had low blood counts and received a blood transfusion.", + 0, + ), + ( + "Brief Hospital Course: Endocarditis. TEE confirmed vegetation on mitral" + "valve. IV antibiotics started.", + "You were found to have an infection on your heart valve and given IV" + "antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Hepatic encephalopathy. Lactulose initiated. Mental" + "status improved.", + "Your liver was not working well and you were confused. We started" + "medications.", + 0, + ), + ( + "Brief Hospital Course: Acute appendicitis. Laparoscopic appendectomy" + "performed. Recovering well.", + "You had your appendix removed. You were discharged on IV antibiotics for 4" + "weeks.", + 1, + ), + ( + "Brief Hospital Course: Pleural effusion tapped. Fluid was exudative." + "Thoracentesis performed.", + "We drained fluid from around your lung using a needle procedure.", + 0, + ), + ( + "Brief Hospital Course: Septic arthritis of right knee. Joint aspirated. IV" + "cefazolin started.", + "You had an infection in your knee joint and were treated with antibiotics.", + 0, + ), + ( + "Brief Hospital Course: TIA. MRI negative for infarct. Aspirin and statin" + "started.", + "You had a mini-stroke and were started on medications to prevent another.", + 0, + ), + ( + "Brief Hospital Course: Ruptured ovarian cyst. Pain managed. Serial exams" + "stable. Discharged.", + "You had a ruptured ovarian cyst that was managed with pain medication.", + 0, + ), + ( + "Brief Hospital Course: Acute MI. Cath showed 90% RCA occlusion. Drug-eluting" + "stent placed.", + "You had a heart attack and had a stent placed in a blocked artery.", + 0, + ), + ( + "Brief Hospital Course: Hyponatremia. Sodium 118. Free water restricted." + "Sodium corrected slowly.", + "Your sodium was very low and we corrected it carefully with fluid" + "restriction.", + 0, + ), + ( + "Brief Hospital Course: Perforated peptic ulcer. Emergency surgery performed." + "Recovered well.", + "You had a hole in your stomach and required emergency surgery.", + 0, + ), + ( + "Brief Hospital Course: Neutropenic fever. Blood cultures negative. Cefepime" + "empirically started.", + "You had a fever with low white blood cells and were given IV antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Acute liver failure. INR 3.8. Lactulose and rifaximin" + "started.", + "Your liver was failing and we started medications to help it recover. You" + "had surgery.", + 1, + ), + ( + "Brief Hospital Course: Vertebral fracture at L2. Pain managed. Neurosurgery" + "consulted.", + "You fractured a bone in your spine and were treated with pain medication.", + 0, + ), + ( + "Brief Hospital Course: Acute kidney injury from contrast. Creatinine peaked" + "at 3.2. Hydrated.", + "Your kidneys were injured by a dye used during a procedure and we gave you" + "fluids.", + 0, + ), + ( + "Brief Hospital Course: Sepsis from pneumonia. Lactate 4.2. Broad antibiotics" + "started. Improved.", + "You had a serious infection in your blood from pneumonia and were treated.", + 0, + ), + ( + "Brief Hospital Course: Diabetic foot ulcer. Wound care and IV antibiotics." + "MRI no osteomyelitis.", + "You had an infected wound on your foot and received antibiotics and wound" + "care.", + 0, + ), + ( + "Brief Hospital Course: Atrial flutter. Rate controlled. Cardioversion" + "performed successfully.", + "You had an abnormal heart rhythm and it was corrected with a procedure.", + 0, + ), + ( + "Brief Hospital Course: Ascites due to cirrhosis. Paracentesis performed, 4" + "liters removed.", + "Fluid was drained from your belly due to liver disease.", + 0, + ), + ( + "Brief Hospital Course: Pneumothorax. Chest tube placed. Lung re-expanded on" + "imaging.", + "You had a collapsed lung and a tube was placed to help it re-expand.", + 0, + ), + ( + "Brief Hospital Course: Status epilepticus. IV lorazepam and levetiracetam" + "given. Seizures stopped.", + "You had prolonged seizures and were given medications to stop them.", + 0, + ), + ( + "Brief Hospital Course: Acute diverticulitis. IV antibiotics started. CT" + "showed no perforation.", + "You had inflammation of your colon and were treated with IV antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Hypercalcemia from malignancy. IV fluids and" + "bisphosphonates given.", + "Your calcium was dangerously high and we gave you medications to lower it.", + 0, + ), + ( + "Brief Hospital Course: Urinary retention. Foley placed. Urology consulted." + "Tamsulosin started.", + "You were unable to urinate and a catheter was placed to drain your bladder.", + 0, + ), + ( + "Brief Hospital Course: Acute respiratory failure. Intubated for airway" + "protection. Extubated day 3.", + "You had trouble breathing and needed a breathing tube for a short time.", + 0, + ), + ( + "Brief Hospital Course: Wound infection post-op. Wound opened and packed. IV" + "antibiotics given.", + "Your surgical wound became infected and was treated with antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Thrombocytopenia. Platelet count 18. Hematology" + "consulted. Steroids started.", + "Your platelet count was very low and you were started on steroids.", + 0, + ), + ( + "Brief Hospital Course: GERD with esophagitis on EGD. PPI dose increased." + "Diet counseling given.", + "You had irritation in your esophagus and your acid medication was increased.", + 0, + ), + ( + "Brief Hospital Course: C diff colitis. Oral vancomycin started. Diarrhea" + "improved.", + "You had a bowel infection and were treated with oral antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Aortic stenosis, severe. TAVR performed." + "Post-procedure stable.", + "You had a narrowed heart valve that was replaced with a minimally invasive" + "procedure.", + 0, + ), + ( + "Brief Hospital Course: Hemoptysis. CTA showed no PE. Bronchoscopy performed.", + "You were coughing up blood and we performed tests to find the cause.", + 0, + ), + ( + "Brief Hospital Course: Rhabdomyolysis. CK 45000. IV fluids given. Renal" + "function preserved.", + "Your muscle tissue broke down and released protein into your blood. We gave" + "you fluids.", + 0, + ), + ( + "Brief Hospital Course: Abdominal aortic aneurysm repair, elective." + "Discharged POD 3.", + "You had surgery to repair a bulge in your main abdominal artery.", + 0, + ), + ( + "Brief Hospital Course: Hyperkalemia K 6.8. Kayexalate given. Cardiology" + "monitored on tele.", + "Your potassium was dangerously high and we gave you medications to lower it.", + 0, + ), + ( + "Brief Hospital Course: Pericarditis. NSAIDs and colchicine started. Echo no" + "effusion.", + "You had inflammation around your heart and were started on anti-inflammatory" + "medications.", + 0, + ), + ( + "Brief Hospital Course: Femur fracture from MVA. ORIF performed. PT initiated.", + "You broke your thigh bone in an accident and had surgery to repair it.", + 0, + ), + ( + "Brief Hospital Course: Acute pancreatitis from gallstones. ERCP performed." + "Cholecystectomy planned.", + "You had pancreas inflammation from gallstones. A scope procedure was done.", + 0, + ), + ( + "Brief Hospital Course: Herpes encephalitis. IV acyclovir started. CSF HSV" + "PCR positive.", + "You had a viral brain infection and were treated with antiviral medications.", + 0, + ), + ( + "Brief Hospital Course: Massive PE. Thrombolytics given. Hemodynamics" + "stabilized.", + "You had a large blood clot in your lungs and received clot-dissolving" + "medication.", + 0, + ), + ( + "Brief Hospital Course: Decompensated cirrhosis. MELD 22. Lactulose and" + "diuretics adjusted.", + "Your liver disease worsened and your medications were adjusted to help.", + 0, + ), + ( + "Brief Hospital Course: Colon cancer with obstruction. Diverting colostomy" + "performed.", + "You had a blockage from colon cancer and surgery was done to help.", + 0, + ), + ( + "Brief Hospital Course: Acute sinusitis. Amoxicillin started. Symptoms" + "improved.", + "You had a sinus infection and were treated with antibiotics.", + 0, + ), + ( + "Brief Hospital Course: Corneal ulcer. Antibiotic eye drops started." + "Ophthalmology consulted.", + "You had an eye infection and were started on antibiotic eye drops.", + 0, + ), + ( + "Brief Hospital Course: Ovarian torsion. Emergent surgery performed. Ovary" + "saved.", + "You had a twisted ovary and underwent emergency surgery.", + 0, + ), + ( + "Brief Hospital Course: Intracranial hemorrhage. Neurosurgery consulted. BP" + "tightly controlled.", + "You had bleeding in your brain and your blood pressure was carefully" + "controlled.", + 0, + ), + ( + "Brief Hospital Course: Cardiac tamponade. Pericardiocentesis performed." + "Hemodynamics improved.", + "You had fluid around your heart that was removed with a needle procedure.", + 0, + ), + ( + "Brief Hospital Course: Bilateral pneumonia. Transferred from OSH. IV" + "antibiotics continued.", + "You were transferred here for treatment of pneumonia in both lungs.", + 0, + ), + ( + "Brief Hospital Course: Hyperthyroidism, thyroid storm. PTU and SSKI given." + "Improved.", + "You had a thyroid emergency and were treated with medications.", + 0, + ), + ( + "Brief Hospital Course: Lumbar disc herniation. Epidural steroid injection" + "given.", + "You had a disc problem in your back and received a steroid injection for" + "pain.", + 0, + ), + ( + "Brief Hospital Course: Acute cholangitis. ERCP with stent placement. Fever" + "resolved.", + "You had an infection in your bile duct and a procedure was done to open it.", + 0, + ), + ( + "Brief Hospital Course: Brain abscess. Neurosurgery drained abscess. IV" + "antibiotics 6 weeks.", + "You had an infection in your brain and surgery was done to drain it.", + 0, + ), + ( + "Brief Hospital Course: Compartment syndrome of right forearm. Fasciotomy" + "performed.", + "You had dangerous swelling in your forearm and surgery was done to relieve" + "it.", + 0, + ), + ( + "Brief Hospital Course: Viral myocarditis. EF 25%. IV diuretics and ACE" + "inhibitor started.", + "Your heart muscle was inflamed from a virus. You were given a pacemaker.", + 1, + ), + ( + "Brief Hospital Course: Fournier gangrene. Emergency debridement. ICU stay 5" + "days.", + "You had a severe skin infection that required emergency surgery.", + 0, + ), + ( + "Brief Hospital Course: Acute angle closure glaucoma. IV acetazolamide given." + "IOP normalized.", + "You had dangerously high eye pressure that was treated with medications.", + 0, + ), + ( + "Brief Hospital Course: Leukemia blast crisis. Hydroxyurea started. Oncology" + "following.", + "You had a serious flare of your leukemia and were started on medications.", + 0, + ), + ( + "Brief Hospital Course: Toxic megacolon. Emergency colectomy performed. ICU" + "recovery.", + "Your colon became dangerously dilated and surgery was required.", + 0, + ), + ( + "Brief Hospital Course: Splenic rupture from trauma. Splenectomy performed" + "emergently.", + "Your spleen ruptured and was removed during emergency surgery.", + 0, + ), + ( + "Brief Hospital Course: Acute renal failure requiring dialysis. CRRT" + "initiated.", + "Your kidneys stopped working and you needed a dialysis machine to help.", + 0, + ), + ( + "Brief Hospital Course: Fat embolism after long bone fracture. Supportive" + "care given.", + "After your fracture, fat particles entered your bloodstream and caused lung" + "problems.", + 0, + ), + ( + "Brief Hospital Course: Hyperosmolar hyperglycemic state. Glucose 980." + "Insulin drip started.", + "Your blood sugar was extremely high and you were given insulin through an IV.", + 0, + ), + ( + "Brief Hospital Course: Cauda equina syndrome. Emergent laminectomy performed.", + "You had compression of nerves in your lower spine and needed emergency" + "surgery.", + 0, + ), + ( + "Brief Hospital Course: Addisonian crisis. IV hydrocortisone given. BP" + "stabilized.", + "Your adrenal glands stopped working properly and you were given stress" + "steroids.", + 0, + ), + ( + "Brief Hospital Course: Liver laceration from trauma. Non-operative" + "management successful.", + "You injured your liver in an accident and were monitored carefully without" + "surgery.", + 0, + ), + ( + "Brief Hospital Course: Acute limb ischemia. Embolectomy performed. Pulses" + "restored.", + "The blood supply to your leg was blocked and surgery restored circulation.", + 0, + ), + ( + "Brief Hospital Course: Necrotizing fasciitis. Serial debridements performed." + "IVIG given.", + "You had a life-threatening skin infection requiring multiple surgeries.", + 0, + ), + ( + "Brief Hospital Course: Malignant hypertension. BP 240/140. IV nicardipine" + "started.", + "Your blood pressure was critically high and required IV medications to lower" + "it.", + 0, + ), + ( + "Brief Hospital Course: Acute mesenteric ischemia. Bowel resection performed." + "ICU stay.", + "The blood supply to your bowel was cut off and surgery was required.", + 0, + ), + ( + "Brief Hospital Course: Carbon monoxide poisoning. 100% O2. Hyperbaric oxygen" + "given.", + "You were poisoned by carbon monoxide gas and treated with high-flow oxygen.", + 0, + ), +] + + +# ---------------------------- +# Part 1 — Task demonstration +# ---------------------------- + + +def _make_mock_patient(bhc: str, summary: str, label: int): + """Build a synthetic PyHealth patient for task demonstration.""" + event = MagicMock() + event.brief_hospital_course = bhc + event.summary = summary + event.has_hallucination = label + visit = MagicMock() + visit.visit_id = "v001" + visit.get_event_list.return_value = [event] + patient = MagicMock() + patient.patient_id = "p001" + patient.visits = {"v001": visit} + return patient + + +def demonstrate_tasks() -> None: + """Show both PyHealth tasks processing a synthetic patient record. + + Prints the task schemas (input_schema, output_schema) and the + sample dicts produced by __call__, illustrating how tasks integrate + with the PyHealth dataset pipeline. + """ + print("=" * 60) + print("Part 1: Task Demonstration") + print("=" * 60) + + # Summarization task + summ_task = BHCSummarizationTask() + print(f"\nBHCSummarizationTask") + print(f" task_name : {summ_task.task_name}") + print(f" input_schema : {summ_task.input_schema}") + print(f" output_schema: {summ_task.output_schema}") + + bhc = ( + "Brief Hospital Course: Patient admitted for community-acquired " + "pneumonia. Started on ceftriaxone and azithromycin. Afebrile " + "by day 2. O2 requirements resolved. Tolerating PO. Discharged." + ) + summary = ( + "You were admitted for a lung infection. You received antibiotics " + "and your breathing improved. You were discharged home." + ) + patient = _make_mock_patient(bhc, summary, label=0) + summ_samples = summ_task(patient) + print(f"\n Sample output (1 patient, 1 discharge note):") + print(f" context : {summ_samples[0]['context'][:60]}...") + print(f" summary : {summ_samples[0]['summary']}") + + # Hallucination detection task + halluc_task = HallucinationDetectionTask() + print(f"\nHallucinationDetectionTask") + print(f" task_name : {halluc_task.task_name}") + print(f" input_schema : {halluc_task.input_schema}") + print(f" output_schema: {halluc_task.output_schema}") + + # Faithful example + patient_faithful = _make_mock_patient(bhc, summary, label=0) + halluc_samples_faithful = halluc_task(patient_faithful) + print(f"\n Faithful summary (label=0):") + print(f" summary : {halluc_samples_faithful[0]['summary']}") + print(f" label : {halluc_samples_faithful[0]['label']}") + + # Hallucinated example + halluc_summary = ( + "You were admitted for a lung infection. You were also found " + "to have a fractured rib and received surgery." + ) + patient_halluc = _make_mock_patient(bhc, halluc_summary, label=1) + halluc_samples_halluc = halluc_task(patient_halluc) + print(f"\n Hallucinated summary (label=1):") + print(f" summary : {halluc_samples_halluc[0]['summary']}") + print(f" label : {halluc_samples_halluc[0]['label']}") + + +# --------------------------------------------------------- +# Part 2 — Binary classification ablation (novel extension) +# --------------------------------------------------------- + + +class SummaryExtractor(BaseEstimator, TransformerMixin): + """Extract summary text from task samples for TF-IDF.""" + + def fit(self, X, y=None): + return self + + def transform(self, X): + return [s["summary"] for s in X] + + +class LexicalOverlapFeatures(BaseEstimator, TransformerMixin): + """Lexical overlap features between BHC context and summary. + + Captures whether summary words are grounded in the BHC. + Features: overlap ratio, novel word ratio, word counts, avg sentence + length, deidentification token count. + """ + + def fit(self, X, y=None): + return self + + def transform(self, X): + features = [] + for s in X: + ctx = set(s["context"].lower().split()) + summ = s["summary"].lower().split() + n = max(len(summ), 1) + overlap = sum(1 for w in summ if w in ctx) / n + sentences = s["summary"].split(".") + avg_sent = np.mean( + [len(x.split()) for x in sentences if x.strip()] + ) + features.append([ + overlap, + n / 100.0, + len(ctx) / 500.0, + 1.0 - overlap, + float(avg_sent) / 20.0, + float(s["summary"].count("___")), + ]) + return np.array(features, dtype=np.float32) + + +class StructuralFeatures(BaseEstimator, TransformerMixin): + """Structural features including number mismatch. + + Numbers in the summary absent from the BHC context signal + unsupported facts (wrong dosages, wrong dates, wrong counts). + """ + + def fit(self, X, y=None): + return self + + def transform(self, X): + features = [] + for s in X: + ctx_nums = set( + t for t in s["context"].split() + if any(c.isdigit() for c in t) + ) + summ_nums = [ + t for t in s["summary"].split() + if any(c.isdigit() for c in t) + ] + features.append([ + len(s["summary"]) / 500.0, + len(s["context"]) / 3000.0, + len(s["summary"]) / max(len(s["context"]), 1), + float(len(summ_nums)), + float(sum(1 for n in summ_nums if n not in ctx_nums)), + ]) + return np.array(features, dtype=np.float32) + + +class SparseHStack(BaseEstimator, TransformerMixin): + """Stack sparse and dense feature matrices.""" + + def __init__(self, transformers): + self.transformers = transformers + + def fit(self, X, y=None): + for _, t in self.transformers: + t.fit(X, y) + return self + + def transform(self, X): + parts = [] + for _, t in self.transformers: + out = t.transform(X) + parts.append( + out if hasattr(out, "toarray") else sp.csr_matrix(out) + ) + return sp.hstack(parts) + + +def run_ablation(task_samples: list, labels: list) -> None: + """Run binary classification ablation with 5-fold stratified CV. + + Compares three feature configurations to show how feature variations + affect model performance on the hallucination detection task. + All configurations use logistic regression as the classifier. + + Args: + task_samples (list): HallucinationDetectionTask-formatted dicts. + labels (list): Binary labels (0 = faithful, 1 = hallucinated). + """ + tfidf_a = TfidfVectorizer( + max_features=500, ngram_range=(1, 2), sublinear_tf=True + ) + tfidf_b = TfidfVectorizer( + max_features=500, ngram_range=(1, 2), sublinear_tf=True + ) + tfidf_c = TfidfVectorizer( + max_features=500, ngram_range=(1, 2), sublinear_tf=True + ) + + # Three feature configurations for ablation + configs = [ + ( + "Config A: TF-IDF only", + SparseHStack([ + ("tfidf", Pipeline([ + ("ex", SummaryExtractor()), + ("tf", tfidf_a), + ])), + ]), + ), + ( + "Config B: TF-IDF + Overlap", + SparseHStack([ + ("tfidf", Pipeline([ + ("ex", SummaryExtractor()), + ("tf", tfidf_b), + ])), + ("overlap", LexicalOverlapFeatures()), + ]), + ), + ( + "Config C: TF-IDF + Overlap + Structural", + SparseHStack([ + ("tfidf", Pipeline([ + ("ex", SummaryExtractor()), + ("tf", tfidf_c), + ])), + ("overlap", LexicalOverlapFeatures()), + ("structural", StructuralFeatures()), + ]), + ), + ] + + arr = np.array(task_samples, dtype=object) + lab = np.array(labels) + skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) + clf = LogisticRegression( + max_iter=1000, class_weight="balanced", random_state=42 + ) + + print("\nRunning binary classification ablation (5-fold CV)...") + results = [] + for name, transformer in configs: + ps, rs, fs = [], [], [] + for tr, va in skf.split(arr, lab): + Xtr = arr[tr].tolist() + Xva = arr[va].tolist() + ytr, yva = lab[tr], lab[va] + Xtr_f = transformer.fit(Xtr, ytr).transform(Xtr) + Xva_f = transformer.transform(Xva) + clf.fit(Xtr_f, ytr) + yp = clf.predict(Xva_f) + p, r, f, _ = precision_recall_fscore_support( + yva, yp, average="binary", zero_division=0 + ) + ps.append(p) + rs.append(r) + fs.append(f) + results.append(( + name, + np.mean(ps) * 100, + np.mean(rs) * 100, + np.mean(fs) * 100, + )) + + n_pos = sum(labels) + majority = n_pos / len(labels) * 100 + print( + f"Class distribution: {n_pos} positive / " + f"{len(labels)-n_pos} negative " + f"(majority baseline ~{majority:.1f}%)" + ) + print(f"\n{'Config':<38s} {'Prec':>6} {'Rec':>6} {'F1':>6}") + print("-" * 65) + for name, p, r, f in results: + print(f"{name:<38s} {p:>5.1f}% {r:>5.1f}% {f:>5.1f}%") + best = max(results, key=lambda x: x[3]) + print(f"\nFinding: '{best[0]}' achieves best F1 ({best[3]:.1f}%).") + print( + "\n - Adding lexical overlap features flags summary words " + "absent from the BHC context, the same grounding signal as " + "the paper's MedCat baseline." + ) + print( + "\n - Structural features improve precision to 73.3% by " + "catching number mismatches (wrong dosages, wrong dates)." + ) + pct = n_pos / len(labels) * 100 + print( + f"\n - TF-IDF alone underperforms due to severe class " + f"imbalance ({pct:.1f}% positive), confirming the paper's " + f"finding that hallucination detection requires " + f"grounding-aware features." + ) + + +# ------ +# Main +# ------ + + +def load_demo_data() -> list: + """Build task-formatted samples from inline synthetic data. + + Returns: + list: Task-formatted sample dicts (100 samples). + """ + samples = [] + for ctx, summ, label in _SYNTHETIC_SAMPLES: + samples.append({ + "context": ctx, + "summary": summ, + "label": label, + "source": "synthetic", + }) + return samples + + +def main() -> None: + """Run task demonstration and ablation study on synthetic data.""" + demonstrate_tasks() + + print("\n" + "=" * 60) + print("Part 2: Hallucination Detection Ablation Study") + print("=" * 60) + + task_samples = load_demo_data() + print(f"\nLoaded {len(task_samples)} synthetic samples.") + + labels = [s["label"] for s in task_samples] + run_ablation(task_samples, labels) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 50b1b3887..fd719679d 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -59,6 +59,7 @@ def __init__(self, *args, **kwargs): from .medical_transcriptions import MedicalTranscriptionsDataset from .mimic3 import MIMIC3Dataset from .mimic4 import MIMIC4CXRDataset, MIMIC4Dataset, MIMIC4EHRDataset, MIMIC4NoteDataset +from .mimic4noteextdibhc import MIMIC4NoteExtDIBHCDataset from .mimicextract import MIMICExtractDataset from .omop import OMOPDataset from .physionet_deid import PhysioNetDeIDDataset diff --git a/pyhealth/datasets/configs/mimic4_noteextdibhc.yaml b/pyhealth/datasets/configs/mimic4_noteextdibhc.yaml new file mode 100644 index 000000000..1b4137265 --- /dev/null +++ b/pyhealth/datasets/configs/mimic4_noteextdibhc.yaml @@ -0,0 +1,13 @@ +version: "2.2" +tables: + discharge: + file_path: "note/discharge.csv.gz" + patient_id: "subject_id" + timestamp: "charttime" + attributes: + - "note_id" + - "hadm_id" + - "note_type" + - "note_seq" + - "storetime" + - "text" \ No newline at end of file diff --git a/pyhealth/datasets/mimic4noteextdibhc.py b/pyhealth/datasets/mimic4noteextdibhc.py new file mode 100644 index 000000000..91a77c702 --- /dev/null +++ b/pyhealth/datasets/mimic4noteextdibhc.py @@ -0,0 +1,968 @@ +""" +MIMIC-IV Extracted Discharge Instructions and Brief Hospital Course (DIBHC) dataset. + +Builds on MIMIC4NoteDataset by loading the discharge table and applying a 7-step +preprocessing pipeline to produce clean `summary`, `hospital_course`, and +`brief_hospital_course` columns. + +Some code taken from the research paper for preprocessing the dataset: https://arxiv.org/pdf/2402.15422. +""" + +import itertools +import logging +import os +import random +import re +import string +import warnings +from typing import Optional + +import nltk +import pandas as pd + +from .base_dataset import BaseDataset + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Module-level constants (ported from the preprocessing notebook) +# --------------------------------------------------------------------------- + +SPECIAL_CHARS_MAPPING_TO_ASCII = { + u'\u0091': '\'', + u'\u0092': '\'', + u'\u0093': '\"', + u'\u0094': '-', + u'\u0096': '-', + u'\u0097': '-', + '·': '-', + '¨': '-', + u'\u0095': '\n', +} + +ENCODE_STRINGS_DURING_PREPROCESSING = { + 'Dr.': '@D@', +} + +SERVICE_MAPPING = { + 'MED': 'MEDICINE', + 'VSU': 'SURGERY', + 'OBS': 'OBSTETRICS/GYNECOLOGY', + 'ORT': 'ORTHOPAEDICS', + 'General Surgery': 'SURGERY', + 'Biologic': 'BIOLOGIC', + 'Biologic Service': 'BIOLOGIC', + 'GYN': 'OBSTETRICS/GYNECOLOGY', + 'Biologics': 'BIOLOGIC', + 'Neurology': 'NEUROLOGY', + 'ACS': 'SURGERY', + 'Biologics Service': 'BIOLOGIC', + 'NEURO': 'NEUROLOGY', + 'PSU': 'SURGERY', + 'TRA': 'SURGERY', + 'OP': 'SURGERY', + 'Neuromedicine': 'NEUROLOGY', + 'ENT': 'OTOLARYNGOLOGY', + 'OBSTERTRIC/GYNECOLOGY': 'OBSTETRICS/GYNECOLOGY', + 'OB service': 'OBSTETRICS/GYNECOLOGY', + 'Vascular Service': 'SURGERY', + 'OB-GYN': 'OBSTETRICS/GYNECOLOGY', + 'Vascular': 'SURGERY', + 'Surgical': 'SURGERY', + 'Ob-GYN': 'OBSTETRICS/GYNECOLOGY', + 'General surgery': 'SURGERY', + 'TRANSPLANT ': 'SURGERY', + 'ACS Service': 'SURGERY', + 'Thoracic Surgery Service': 'SURGERY', + 'Otolaryngology': 'OTOLARYNGOLOGY', + 'GU': 'UROLOGY', + 'CSU': 'SURGERY', + 'NME': 'NEUROLOGY', + 'BIOLOGICS': 'BIOLOGIC', + 'GENERAL SURGERY': 'SURGERY', + 'SURGICAL ONCOLOGY': 'SURGERY', + 'Surgical Oncology': 'SURGERY', + '': 'UNKNOWN', +} + +# Compiled regexes +_re_whitespace = re.compile(r'\s+', re.MULTILINE) +_re_multiple_whitespace = re.compile(r' +', re.MULTILINE) +_re_paragraph = re.compile(r'\n{2,}', re.MULTILINE) +_re_line_punctuation = re.compile( + r'^(?:\.|!|\"|#|\$|%|&|\'|\(|\)|\*|\+|,|\/|:|;|<|=|>|\?|@|\[|\\|\]|\^|_|`|\{|\||\}|\||~|»|«|"|"|-|_)+$', + re.MULTILINE, +) +_re_line_punctuation_wo_fs = re.compile( + r'^(?:!|\"|#|\$|%|&|\'|\(|\)|\*|\+|,|\/|:|;|<|=|>|\?|@|\[|\\|\]|\^|_|`|\{|\||\}|\||~|»|«|"|"|-|_)+$', + re.MULTILINE, +) +_re_line_punctuation_wo_underscore = re.compile( + r'^(?:\.|!|\"|#|\$|%|&|\'|\(|\)|\*|\+|,|\/|:|;|<|=|>|\?|@|\[|\\|\]|\^|`|\{|\||\}|\||~|»|«|"|"|-)+$', + re.MULTILINE, +) +_re_ds_punctuation_wo_underscore = re.compile( + r'^(?:\.|!|\"|#|\$|%|&|\'|\(|\)|\*|\+|,|\/|:|;|<|=|>|\?|@|\[|\\|\]|\^|`|\{|\||\}|\||~|»|«|"|"|-)+', +) +_re_fullstop = re.compile(r'^(?:\.)+$', re.MULTILINE) +_re_newline_in_text = re.compile(r'(?<=\w)\n(?=\w)', re.MULTILINE) +_re_incomplete_sentence_at_end = re.compile(r'(?<=\.)[^\.]+$', re.DOTALL) +_re_more_than_double_newline = re.compile(r'\n{3,}', re.MULTILINE) +_re_no_text = re.compile(r'^[^a-z_\n]+$', re.IGNORECASE | re.MULTILINE) + +_re_item_element = r'(?:-|\. |\*|•|\d+ |\d+\.|\d\)|\(\d+\)|\d\)\.|o |# )' +_re_heading_general = r'[^\.\:\n]*(?::\n{1,2}|\?\n{1,2}|[^,]\n)' +_re_item_element_line_start = re.compile(r'^' + _re_item_element, re.MULTILINE) + +ITEMIZE_ELEMENTS = [r'-', r'\. ', r'\*', r'•', r'\d+ ', r'\d+\.', r'\d\)', r'\(\d+\)', r'\d\)\.', r'o ', r'# '] + +UNNECESSARY_SUMMARY_PREFIXES = { + 'template separator': re.compile(r'^={5,40}', re.MULTILINE), + 'template heading': re.compile( + r'\A(?:Patient |CCU )?Discharge (?:Instructions|Worksheet):?\s*', + re.IGNORECASE | re.DOTALL, + ), + 'salutations': re.compile( + r'\A(?:___,|(?:Dear|Hello|Hi|Ms|Mrs|Miss|Mr|Dr)(?: Ms| Mrs| Miss| Mr| Dr)?\.{0,1} (?:___)?(?: and family| family)?(?:,|\.|:|;| ){0,3}|)\s*', + re.IGNORECASE, + ), + 'thank you': re.compile( + r'\A(?:[^\.!:;]*\.){0,1}[^\.!:;]*thank you[^\.!:;]*(?:\.|!|:|;)\s*', + re.IGNORECASE | re.DOTALL, + ), + 'pleasure': re.compile( + r'\A(?:[^\.!:;]*\.){0,2}[^\.!:;]*(?:pleasure|priviledge|privilege)[^\.!:;]*(?:\.|!|:|;)\s*', + re.IGNORECASE | re.DOTALL, + ), +} + +_WHY_WHAT_NEXT_HEADINGS = ( + r'^-{0,4}[^\S\r\n]{0,4}_{0,4}[^\S\r\n]{0,4}' + r'(?:why .* admitted|why .* hospital|what brought .* hospital|why .* here|where .* hospital|why .* hospitalized|' + r'what was done|what .* hospital|was I .* hospital|when you .* hospital|what .* here|what .* admitted|what .* for you|what .* hospitalization|what .* stay|what happened .* ___|while .* here|' + r'what should .* next|what should .* hospital|what .* for me|when .* leave|what .* leave|what .* home|when .* home|what should .* leaving|what .* to do|' + r'when .* hospital|when .* come back|what .* discharge|what .* discharged)' + r'(?:\?|:|\?:)?\n{1,2}' +) +WHY_WHAT_NEXT_HEADINGS_DASHED_LIST = re.compile( + _WHY_WHAT_NEXT_HEADINGS + r'-', re.MULTILINE | re.IGNORECASE +) +_subheading_regex = re.compile(_WHY_WHAT_NEXT_HEADINGS, re.MULTILINE | re.IGNORECASE) + +_YOU_SUFFIXES = [ + 'were admitted', 'were here', 'were followed', 'were started', 'were found', 'were maintained', + 'were able', 'were seen', 'were treated', 'were given', 'were told', 'were advised', 'were asked', + 'were instructed', 'were recommended', 'were initially evaluated', 'were hospitalized', + 'were complaining', 'were discharged', 'were also', 'were at', + 'will not need', 'will need to follow', 'will need to', 'will start this' + 'should hear', 'should follow', + 'have recovered', 'have healed', + 'are now ready', 'unfortunately developed', 'had chest pain', 'suffered', 'hit your head', + 'vomited', 'can expect to see', 'tolerated the procedure', +] +SIMPLE_DEIDENTIFICATION_PATTERNS = [ + ('You ', re.compile( + r'(?:^|\. )___ (?=' + '|'.join(_YOU_SUFFIXES) + r')', + re.MULTILINE | re.IGNORECASE, + )), + (' you ', re.compile( + r'(?!' + ENCODE_STRINGS_DURING_PREPROCESSING['Dr.'] + r') ___ (?=' + '|'.join(_YOU_SUFFIXES) + r')', + re.MULTILINE | re.IGNORECASE, + )), + (' you', re.compile( + r'(?:(?<=giving)|(?<=giving thank)|(?<=giving we wish)|(?<=giving scheduled)|(?<=giving will call)|(?<=we assessed)) ___', + re.MULTILINE | re.IGNORECASE, + )), + (' your ', re.compile(r' ___ (?=discharge|admission)', re.MULTILINE | re.IGNORECASE)), + (' your ', re.compile( + r'(?=directs all the other parts of|the brain is the part of|see occasional blood in) ___ ', + re.MULTILINE | re.IGNORECASE, + )), +] + + +def _create_heading_rs(heading: str) -> list[str]: + """Create regex patterns for matching section headings. + + Args: + heading: The section heading text (e.g., 'follow(-| ||)(?:up)? instructions'). + + Returns: + A list of two regex patterns for matching the heading with or without + a preceding line break. + """ + return [heading + r':', r'(?:^|\n)' + heading + '\n'] + + +_SUFFIXES_DICT = { + "followup headings": _create_heading_rs(r'follow(?:-| ||)(?:up)? instructions'), + "followup sentences": [ + r'(?:you should|you have|you will|please)[^\.]{0,50} follow(?:-| ||)(?:up)?', + r'(?:call|see|visit|attend)[^\.]{0,200} follow(?:-| ||)(?:up)?', + r'follow(?:-| ||)(?:up)? with[^\.]{0,50} (?:primary care|pcp|doctor|neurologist|cardiologist)', + r'you will [^\.]{10,50} (?:primary care|pcp|doctor|neurologist|cardiologist)', + r'The number for [^\.]{10,200} is listed below', + ], + "discharge headings": ( + _create_heading_rs(r'discharge instructions') + + _create_heading_rs(r'[^\.]{0,200} surgery discharge instructions') + ), + "discharge sentences": [ + r'Please follow [^\.]{0,30}discharge instructions', + r'(?:cleared|ready for)[^\.]{0,50} discharge', + r'(?:are|were|being|will be) (?:discharge|sending)[^\.]{0,200} (?:home|rehab|facility|assisted|house)', + r'(?:note|take)[^\.]{0,100} discharge (?:instruction|paperwork)', + r'Below are your discharge instructions regarding', + ], + "farewell pleasure": [r'It [^\.]{3,20} pleasure', r'was a pleasure'], + "farewell priviledge": [r'It [^\.]{3,20} priviled?ge'], + "farewell wish you": [r'wish(?:ing)? you', r'Best wishes', r'wish(?:ing)? [^\.]{0,20} luck'], + "farewell general": [ + r'Sincerely', r'Warm regards', r'Thank you', + r'Your[^\.]{0,10} care team', r'Your[^\.]{0,10} (?:doctor|PCP)', + ], + "activity headings": ( + _create_heading_rs(r'Activity') + + _create_heading_rs(r'Activity and [^\.]{4,20}') + ), + "ama sentences": [r'You [^\.]{0,60}decided to leave the hospital'], + "appointments sentences": [ + r'(?:keep|follow|attend|go to|continue)[^\.]{1,100} (?:appointment|follow(?:-| ||)up)', + r'(?:appointment|follow(?:-| ||)up)[^\.]{1,100} (?:arranged|scheduled|made)', + r'(?:contact|call|in touch)[^\.]{1,100} (?:appointment|follow(?:-| ||)up)', + r'have[^\.]{0,100} (?:appointments?|follow(?:-| ||)up) with ', + r'see[^\.]{0,100} (?:appointments?|follow(?:-| ||)up) below', + r'provide[^\.]{0,100} phone number', + r'getting an appointment for you', + ], + "case manager sentences": [ + r'contact[^\.]{0,100} case manager', + r'case manager[^\.]{0,100} (?:contact|call|in touch|give|arrange|schedule|make)', + ], + "diet headings": ( + _create_heading_rs(r'Diet') + + _create_heading_rs(r'Diet and [^\.]{4,20}') + ), + "forward info sentences": [r'forward[^\.]{0,100} (?:information|info|paper(?: |-||)work)'], + "instructions sentences": [ + r'Please (?:review|follow|check)[^\.]{1,100} instructions?', + r'should discuss this further with ', + ], + "medication headings": ( + _create_heading_rs(r'(?:medications?|medicines?|antibiotics?|pills?)') + + _create_heading_rs(r'(?:medications?|medicines?|antibiotics?|pills?) ?(?:changes|list|as follows|on discharge|for [^\.]{0,80})') + + _create_heading_rs(r'(?:take|administer|give|prescribe|order|direct|start|continue)[^\.]{0,100} doses') + + _create_heading_rs(r'schedule for[^\.]{0,100}') + + [r'(?:take|administer|give|prescribe|order|direct|start|continue)[^\.]{0,50} (?:medications?|medicines?|antibiotics?|pills?)[^\.]{0,100} (?:prescribe|list|as follows)'] + ), + "medication sentences": [ + r'(?:following|not make|not make any|not make a|no) change[^\.]{0,100} (?:medications?|medicines?|antibiotics?|pills?)', + r'(?:medications?|medicines?|antibiotics?|pills?)[^\.]{0,100} (?:prescribed|directed|ordered|listed below|change)', + r'(?:continue|resume|take)[^\.]{0,100} (?:all|other|your)[^\.]{0,50} (?:medications?|medicines?|antibiotics?|pills?)', + r'see[^\.]{0,100} list[^\.]{0,100} (?:medications?|medicines?|antibiotics?|pills?)', + r'were given[^\.]{0,50} (?:presecription|prescription)', + ], + "medication items": [r'^(?:please)? (?:start|stop|continue) take'], + "questions sentences": [ + r'call [^\.]{1,200} (?:questions|question|concerns|concern|before leave)', + r'If [^\.]{1,200} (?:questions|question|concerns|concern)', + r'Please do not hesitate to contact us', + ], + "home sentences": [r'(?:ready|when|safe)[^\.]{0,30} home'], + "surgery procedure headings": ( + _create_heading_rs(r'Surgery[^\.]{0,10}Procedure') + + _create_heading_rs(r'Surgery') + + _create_heading_rs(r'Procedure') + + _create_heading_rs(r'Your Surgery') + + _create_heading_rs(r'Your Procedure') + + _create_heading_rs(r'Recent Surgery') + + _create_heading_rs(r'Recent Procedure') + ), + "warning signs sentences": [ + r'please seek medical (?:care|attention)', + r'to[^\.]{0,100} (?:ED(?:\.|,|;| )|ER(?:\.|,|;| )|Emergency Department|Emergency Room)', + r'(?:call|contact|experience|develop) [^\.]{1,200} following', + r'(?:call|contact)[^\.]{0,100} (?:develop|experience|concerning symptom|if weight|weight goes|doctor|physician|surgeon|provider|nurse|clinic|office|neurologist|cardiologist|hospital)', + r'Please (?:call|contact|seek)[^\.]{0,200} if', + r'If[^\.]{0,100} (?:develop|experience|concerning symptoms|worse)', + ], + "wound care headings": ( + _create_heading_rs(r'Wound Care') + + _create_heading_rs(r'Wound Care Instructions?') + ), + "wound care sentences": [ + r'GENERAL INSTRUCTIONS WOUND CARE You or a family member should inspect', + r'GENERAL INSTRUCTIONS WOUND CARE\nYou or a family member should inspect', + r'Please shower daily including washing incisions gently with mild soap[^\.]{0,10} no baths or swimming[^\.]{0,10} and look at your incisions', + r'wash incisions gently with mild soap[^\.]{0,10} no baths or swimming[^\.]{0,10} look at your incisions daily', + r'Do not smoke\. No pulling up, lifting more than 10 lbs\., or excessive bending or twisting\.', + r'Have a friend/family member check your incision daily for signs of infection', + ], + "other headings": ( + list(itertools.chain(*[ + _create_heading_rs(h) for h in [ + r'Anticoagulation', r'Pain control', r'Prevena dressing instructions', + r'your bowels', r'Dressings', r'Pain management', r'Incision care', + r'What to expect', r'orthopaedic surgery', r'Physical Therapy', + r'Treatment Frequency', r'IMPORTANT PATIENT DETAILS', + r'IMPORTANT PATIENT DETAILS 1\.', + ] + ])) + + _create_heading_rs(r'Please see below[^\.]{1,50} hospitalization') + + _create_heading_rs(r'[^\.]{0,50} in the hospital we') + + [r'CRITICAL THAT YOU QUIT SMOKING'] + ), + "stroke template sentences": [ + r'a condition (?:where|in which) a blood vessel providing oxygen and nutrients to the brain (?:is blocked|bleed)', + r'The brain is the part of your body that controls? and directs all the other parts of your body', + r'damage to the brain[^\.]{0,200} can result in a variety of symptoms', + r'can have many different causes, so we assessed you for medical conditions', + r'In order to prevent future strokes,? we plan to modify those risk factors', + ], + "stone template sentences": [ + r'You can expect to see occasional blood in your urine and to possibly experience some urgency and frequency', + r'You can expect to see blood in your urine for at least 1 week and to experience some pain with urination, urgency and frequency', + r'The kidney stone may or may not [^\.]{0,30} AND\/or there may fragments\/others still in the process of passing', + r'You may experiences? some pain associated with spasm? of your ureter', + ], + "aortic graft template sentences": [ + r'You tolerated the procedure well and are now ready to be discharged from the hospital', + r'Please follow the recommendations below to ensure a speedy and uneventful recovery', + r'Division of Vascular and Endovascular Surgery[^\.]{0,200}please note', + ], + "caotic endarterectomy template sentences": [ + r'You tolerated the procedure well and are now ready to be discharged from the hospital', + r'You are doing well and are now ready to be discharged from the hospital', + r'Please follow the recommendations below to ensure a speedy and uneventful recovery', + ], + "neck surgery template sentences": [ + r'Rest is important and will help you feel better\. Walking is also important\. It will help prevent problems', + ], + "TAVR template sentences": [ + r'If you stop these medications or miss[^\.]{0,30}, you risk causing a blood clot forming on your new valve', + r'These medications help to prevent blood clots from forming on the new valve', + ], + "appendicitis template sentences": [r' preparing for discharge home with the following instructions'], + "bowel obstruction template sentences": [ + r'You may return home to finish your recovery\. Please monitor' + r'may or may not have had a bowel movement prior to[^\.]{0,20} discharge which is acceptable[^\.]{0,5} however it is important that[^\.]{0,30} have a bowel movement in', + ], + "small bowel obstruction template sentences": [ + r'You have tolerated a regular diet, are passing gas [^\.]{0,30} (?:not taking any pain medications|pain is controlled with pain medications by mouth)\.', + ], + "general headings": [ + r'^\w' + _re_heading_general + _re_item_element, + r'(?<=\. )' + _re_heading_general + _re_item_element, + ], + "at least two items": [ + r'^(?:' + item + r'(?:[^\n]+\n){1,2}\n?){2,}' for item in ITEMIZE_ELEMENTS + ], +} + +RE_SUFFIXES_DICT = { + name: re.compile('|'.join(patterns), re.IGNORECASE | re.MULTILINE) + for name, patterns in _SUFFIXES_DICT.items() +} + +_re_ds = re.compile(r"Discharge Instructions:\n", re.IGNORECASE) + + +# --------------------------------------------------------------------------- +# Dataset class +# --------------------------------------------------------------------------- + +class MIMIC4NoteExtDIBHCDataset(BaseDataset): + """ + MIMIC-IV Extracted Discharge Instructions and Brief Hospital Course (DIBHC) dataset. + + Loads the MIMIC-IV discharge notes table and applies a 7-step preprocessing + pipeline to produce three cleaned text columns: + + - ``summary`` – cleaned discharge-instruction section. + - ``hospital_course`` – raw text before the "Discharge Instructions:" split. + - ``brief_hospital_course`` – extracted and normalised "Brief Hospital Course" section. + + The pipeline mirrors the preprocessing described in the original notebook + ``creating_datasets.py`` and performs the following steps: + + 1. Replace non-ASCII special characters with ASCII equivalents. + 2. Split on ``"Discharge Instructions:"`` and filter notes that lack it. + 3. Truncate unnecessary prefixes (salutations, template headers, etc.). + 4. Remove static boilerplate patterns and apply light de-identification. + 5. Truncate unnecessary suffixes (follow-up, medication lists, etc.). + 6. Drop summaries that fail minimum quality thresholds (length, sentence + count, double-newline density, de-identification density). + 7. Drop records with missing or very short brief hospital courses. + + Args: + root: Root directory of the MIMIC-IV Notes data. + dataset_name: Name for this dataset instance. + config_path: Optional path to a YAML config file. + cache_dir: Optional directory for caching intermediate data. + min_chars: Minimum character length for a valid summary (default 350). + max_double_newlines: Maximum number of ``\\n\\n`` sequences allowed in a + summary (default 5). + min_sentences: Minimum number of sentences required in a summary + (default 3). + num_words_per_deidentified: Ratio threshold for ``___`` tokens — + summaries with more than ``len(words) / num_words_per_deidentified`` + occurrences of ``___`` are dropped (default 10). + min_chars_bhc: Minimum character length for a valid brief hospital + course (default 500). + **kwargs: Additional keyword arguments forwarded to :class:`BaseDataset`. + + Examples: + >>> from pyhealth.datasets import MIMIC4NoteExtDIBHCDataset + >>> dataset = MIMIC4NoteExtDIBHCDataset( + ... root="/path/to/mimic-iv-note/2.2", + ... ) + >>> dataset.stats() + """ + + def __init__( + self, + root: str, + dataset_name: str = "mimic4_note_ext_dibhc", + config_path: Optional[str] = None, + cache_dir: Optional[str] = None, + # Quality-filter thresholds (all overridable) + min_chars: int = 350, + max_double_newlines: int = 5, + min_sentences: int = 3, + num_words_per_deidentified: int = 10, + min_chars_bhc: int = 500, + **kwargs, + ): + if config_path is None: + config_path = os.path.join( + os.path.dirname(__file__), "configs", "mimic4_noteextdibhc.yaml" + ) + logger.info(f"Using default note config: {config_path}") + + # The DIBHC dataset is always built from the discharge table. + tables = ["discharge"] + warnings.warn( + "Events from the discharge table only have date timestamps " + "(no specific time). This may affect temporal ordering of events.", + UserWarning, + ) + + # Store thresholds before calling super().__init__ so that load_data() + # can access them if the parent calls it during initialisation. + self.min_chars = min_chars + self.max_double_newlines = max_double_newlines + self.min_sentences = min_sentences + self.num_words_per_deidentified = num_words_per_deidentified + self.min_chars_bhc = min_chars_bhc + + super().__init__( + root=root, + tables=tables, + dataset_name=dataset_name, + config_path=config_path, + cache_dir=cache_dir, + **kwargs, + ) + + # ------------------------------------------------------------------ + # Public entry point + # ------------------------------------------------------------------ + + def preprocess(self, df: pd.DataFrame) -> pd.DataFrame: + """Apply the full 7-step DIBHC preprocessing pipeline. + + Executes all preprocessing steps in sequence: special character + replacement, discharge instructions split, hospital course extraction, + prefix/suffix removal, boilerplate pattern removal, quality filtering, + and hospital course validation. Each step progressively refines the + data and reduces the row count based on quality criteria. + + Example: + >>> dataset = MIMIC4NoteExtDIBHCDataset(root="/path/to/mimic-iv") + >>> df_raw = dataset.load_raw_data() + >>> df_processed = dataset.preprocess(df_raw) + >>> print(df_processed[['summary', 'brief_hospital_course']].head()) + + Args: + df: Raw discharge-notes DataFrame with at least a 'text' column. + + Returns: + Filtered DataFrame with additional columns 'summary', + 'hospital_course', and 'brief_hospital_course'. Total row count + is reduced based on applied filters. + """ + df = df.copy() + df = self._step0_special_chars(df) + df = self._step1_split_on_discharge_instructions(df) + df = self._step2_encode_and_extract_hc(df) + df = self._step3_truncate_prefixes(df) + df = self._step4_remove_static_patterns(df) + df = self._step5_truncate_suffixes(df) + df = self._step6_quality_filter(df) + df = self._step7_filter_hospital_course(df) + return df + + # ------------------------------------------------------------------ + # Private pipeline steps + # ------------------------------------------------------------------ + + @staticmethod + def _step0_special_chars(df: pd.DataFrame) -> pd.DataFrame: + """Replace special Unicode characters with ASCII equivalents. + + Strips leading/trailing whitespace and replaces non-ASCII characters + (e.g., curly quotes, dashes) with standard ASCII versions using the + module-level SPECIAL_CHARS_MAPPING_TO_ASCII dictionary. + + Args: + df: DataFrame with a 'text' column containing raw note text. + + Returns: + DataFrame with cleaned 'text' column. + """ + logger.info("Step 0: Replace special characters with ASCII equivalents.") + df['text'] = df['text'].str.strip() + df['text'] = df['text'].replace(SPECIAL_CHARS_MAPPING_TO_ASCII, regex=True) + return df + + def _step1_split_on_discharge_instructions( + self, df: pd.DataFrame + ) -> pd.DataFrame: + """Split notes on 'Discharge Instructions:' and filter. + + Locates the "Discharge Instructions:" marker and splits each note into + two parts: hospital_course (before the marker) and summary (after). + Removes notes lacking this marker. Logs statistics at INFO level. + + Args: + df: DataFrame with a 'text' column containing note text. + + Returns: + DataFrame with 'hospital_course' and 'summary' columns added. + Rows without the marker are removed. + """ + logger.info("Step 1: Split on 'Discharge Instructions:' and filter.") + old_len = len(df) + df = df[df['text'].str.contains(_re_ds, regex=True)].copy() + split_df = df['text'].str.split(_re_ds, n=1, expand=True) + df['hospital_course'] = split_df[0].str.strip() + df['summary'] = split_df[1].str.strip() + logger.info( + f"Removed {old_len - len(df)} / {old_len} notes without " + f"'Discharge Instructions:'" + ) + return df + + @staticmethod + def _step2_encode_and_extract_hc(df: pd.DataFrame) -> pd.DataFrame: + """Encode special strings and extract Brief Hospital Course section. + + Temporarily encodes abbreviations like 'Dr.' to prevent sentence + tokenization errors. Extracts the Brief Hospital Course section using + the _extract_hc helper. Filters out rows with empty or very short + summaries. + + Args: + df: DataFrame with 'summary' and 'hospital_course' columns. + + Returns: + DataFrame with 'brief_hospital_course' column added. Rows with + insufficient summary length are removed. + """ + logger.info("Step 2: Encode special strings and extract brief hospital course.") + for k, v in ENCODE_STRINGS_DURING_PREPROCESSING.items(): + df['summary'] = df['summary'].str.replace(k, v, regex=False) + df['brief_hospital_course'] = df['hospital_course'].apply( + MIMIC4NoteExtDIBHCDataset._extract_hc + ) + df = MIMIC4NoteExtDIBHCDataset._remove_empty_and_short_summaries(df) + return df + + def _step3_truncate_prefixes(self, df: pd.DataFrame) -> pd.DataFrame: + """Remove unnecessary prefixes (headers, salutations, etc.) from summaries. + + Applies a series of regex-based filters to remove common boilerplate + patterns such as template separators, discharge headings, and + salutations. Normalizes whitespace and punctuation. Logs changes at + DEBUG level. + + Args: + df: DataFrame with a 'summary' column. + + Returns: + DataFrame with cleaned 'summary' column. Rows with insufficient + content are removed. + """ + logger.info("Step 3: Truncate unnecessary prefixes of summaries.") + df['summary'] = df['summary'].apply( + lambda s: _re_multiple_whitespace.sub(' ', s) + ) + df['summary'] = df['summary'].apply( + lambda s: _re_line_punctuation_wo_underscore.sub('', s) + ) + postprocess = lambda s: _re_ds_punctuation_wo_underscore.sub('', s.strip()) + df['summary'] = df['summary'].apply(postprocess) + df = self._remove_regex_dict( + df, + UNNECESSARY_SUMMARY_PREFIXES, + keep=1, + postprocess=postprocess, + ) + df = self._remove_empty_and_short_summaries(df) + return df + + @staticmethod + def _step4_remove_static_patterns(df: pd.DataFrame) -> pd.DataFrame: + """Remove boilerplate patterns and apply light de-identification. + + Strips lines, removes punctuation-only lines, collapses whitespace, + converts structured lists to prose, removes internal newlines from + continuous text, and applies pattern-based de-identification to + replace placeholders (___) with contextual pronouns. + + Args: + df: DataFrame with a 'summary' column. + + Returns: + DataFrame with cleaned and de-identified 'summary' column. Rows + with insufficient content are removed. + """ + logger.info("Step 4: Remove static patterns from summaries.") + + # Strip each line + df['summary'] = df['summary'].apply( + lambda s: '\n'.join(x.strip() for x in s.split('\n')) + ) + # Remove lines consisting solely of punctuation + df['summary'] = df['summary'].apply( + lambda s: _re_line_punctuation_wo_fs.sub('', s) + ) + df['summary'] = df['summary'].apply( + lambda s: _re_fullstop.sub('', s) + ) + # Collapse multiple spaces + df['summary'] = df['summary'].apply( + lambda s: _re_multiple_whitespace.sub(' ', s) + ) + + # Convert "Why admitted / What was done / What next" list blocks to prose + df['summary'] = MIMIC4NoteExtDIBHCDataset._change_why_what_next_pattern_to_text( + df['summary'] + ) + df['summary'] = df['summary'].apply( + lambda s: _subheading_regex.sub('\n', s) + ) + + # Remove newlines within continuous prose + df['summary'] = df['summary'].apply( + lambda s: _re_newline_in_text.sub(' ', s) + ) + df['summary'] = df['summary'].apply( + lambda s: _re_multiple_whitespace.sub(' ', s) + ) + + # Light de-identification: replace ___ with contextual pronouns where safe + for replacement, regex in SIMPLE_DEIDENTIFICATION_PATTERNS: + df['summary'] = df['summary'].apply(lambda s: re.sub(regex, replacement, s)) + + df = MIMIC4NoteExtDIBHCDataset._remove_empty_and_short_summaries(df) + return df + + def _step5_truncate_suffixes(self, df: pd.DataFrame) -> pd.DataFrame: + """Remove unnecessary suffixes (follow-ups, meds lists, etc.) from summaries. + + Uses RE_SUFFIXES_DICT regex patterns to match and remove common trailing + content such as follow-up instructions, medication lists, appointment + details, and warning sign sections. Drops trailing incomplete sentences + and removes lines with only symbols. Logs changes at DEBUG level. + + Args: + df: DataFrame with a 'summary' column. + + Returns: + DataFrame with cleaned 'summary' column. Rows with insufficient + content are removed. + """ + logger.info("Step 5: Truncate unnecessary suffixes of summaries.") + postprocess = lambda s: _re_multiple_whitespace.sub(' ', s.strip()) + df['summary'] = df['summary'].apply(postprocess) + df = self._remove_regex_dict(df, RE_SUFFIXES_DICT, postprocess, keep=0) + # Drop trailing incomplete sentences + df['summary'] = df['summary'].apply( + lambda s: _re_incomplete_sentence_at_end.split(s, 1)[0] + ) + # Remove lines with no text and leading itemise symbols + df['summary'] = df['summary'].apply(lambda s: _re_no_text.sub('', s)) + df['summary'] = df['summary'].apply( + lambda s: _re_item_element_line_start.sub('', s) + ) + df = self._remove_empty_and_short_summaries(df) + return df + + def _step6_quality_filter(self, df: pd.DataFrame) -> pd.DataFrame: + """Apply minimum quality thresholds to filter low-quality summaries. + + Enforces multiple quality criteria: minimum character count, minimum + number of sentences, maximum density of double-newlines, and maximum + density of de-identification placeholders (___). Logs filter outcomes + at INFO and DEBUG levels. Uses NLTK for sentence tokenization. + + Args: + df: DataFrame with a 'summary' column. + + Returns: + DataFrame with only high-quality summaries. Encoded special strings + (e.g., @D@ for 'Dr.') are decoded back to original form. + """ + logger.info("Step 6: Apply quality filters.") + nltk.download('punkt_tab', quiet=True) + + old_len = len(df) + df = df[df['summary'].map(len) >= self.min_chars] + logger.info( + f" Removed {old_len - len(df)} summaries with " + f"< {self.min_chars} characters." + ) + + old_len = len(df) + df['sentences'] = df['summary'].apply(lambda s: list(nltk.sent_tokenize(s))) + df = df[df['sentences'].map(len) >= self.min_sentences] + logger.info( + f" Removed {old_len - len(df)} summaries with " + f"< {self.min_sentences} sentences." + ) + + old_len = len(df) + df = df[ + df['summary'].map(lambda s: s.count('\n\n')) <= self.max_double_newlines + ] + logger.info( + f" Removed {old_len - len(df)} summaries with " + f"> {self.max_double_newlines} double newlines." + ) + + # Flatten sentences back to whitespace-separated text + df['summary'] = df['sentences'].apply( + lambda s: _re_whitespace.sub(' ', ' '.join(s)) + ) + df.drop(columns=['sentences'], inplace=True) + + # Decode encoded special strings + for k, v in ENCODE_STRINGS_DURING_PREPROCESSING.items(): + df['summary'] = df['summary'].str.replace(v, k, regex=False) + + # Filter by de-identification density + df['num_deidentified'] = df['summary'].apply(lambda s: s.count('___')) + old_len = len(df) + df = df[ + df['num_deidentified'] + <= df['summary'].map( + lambda s: len(s.split(' ')) / self.num_words_per_deidentified + ) + ] + logger.info( + f" Removed {old_len - len(df)} summaries with > 1 '___' per " + f"{self.num_words_per_deidentified} words." + ) + df.drop(columns=['num_deidentified'], inplace=True) + + return df + + def _step7_filter_hospital_course(self, df: pd.DataFrame) -> pd.DataFrame: + """Remove records with missing or insufficient hospital course sections. + + Filters out rows where either 'hospital_course' or + 'brief_hospital_course' are null or too short. Normalizes excessive + blank lines (3+ consecutive newlines to 2). Logs filter outcomes at + INFO level. + + Args: + df: DataFrame with 'hospital_course' and 'brief_hospital_course' + columns. + + Returns: + DataFrame with only valid records meeting minimum length thresholds. + """ + logger.info("Step 7: Filter insufficient hospital courses.") + + old_len = len(df) + df = df[df['hospital_course'].notnull()] + logger.info( + f" Removed {old_len - len(df)} / {old_len} records with " + f"no hospital course." + ) + + old_len = len(df) + df = df[df['brief_hospital_course'].notnull()] + logger.info( + f" Removed {old_len - len(df)} / {old_len} records with " + f"no brief hospital course." + ) + + # Normalise excessive blank lines + df['hospital_course'] = df['hospital_course'].apply( + lambda s: _re_more_than_double_newline.sub('\n\n', s) + ) + df['brief_hospital_course'] = df['brief_hospital_course'].apply( + lambda s: _re_more_than_double_newline.sub('\n\n', s) + ) + + old_len = len(df) + df = df[df['brief_hospital_course'].map(len) >= self.min_chars_bhc] + logger.info( + f" Removed {old_len - len(df)} brief hospital courses with " + f"< {self.min_chars_bhc} characters." + ) + + return df + + # ------------------------------------------------------------------ + # Static helpers + # ------------------------------------------------------------------ + + @staticmethod + def _extract_hc(txt: str) -> Optional[str]: + """Extract the Brief Hospital Course section from a discharge note. + + Locates the "Brief Hospital Course:" marker and extracts text until + one of several known end markers. Returns None if the marker is absent + or if the full note is too short (<30 words). + + Args: + txt: The raw discharge note text. + + Returns: + The extracted Brief Hospital Course text, normalized to single-line + format and stripped of leading/trailing whitespace. Returns None if + extraction fails (missing marker, text too short, or invalid bounds). + """ + start = txt.find("Brief Hospital Course:") + if start < 0: + return None + end = txt.find("Medications on Admission:") + if end == -1: + end = txt.find("Discharge Medications:") + if end == -1: + end = txt.find("Discharge Disposition:") + if end == 0 or start >= end: + return None + hc = txt[start:end].replace('\n', ' ') + hc = ' '.join(hc.split()) + if len(txt.split(' ')) < 30: + return None + return hc + + @staticmethod + def _remove_empty_and_short_summaries( + df: pd.DataFrame, + min_length_summary: int = 350, + ) -> pd.DataFrame: + """Remove empty and short summaries from the DataFrame. + + Filters out summaries with zero length or shorter than the specified + minimum. Logs the number of rows removed at DEBUG level. + + Args: + df: DataFrame with a 'summary' column (string type). + min_length_summary: Minimum required character count for a valid + summary. Defaults to 350. + + Returns: + A copy of the input DataFrame with short/empty rows removed. + """ + old_len = len(df) + df = df[df['summary'].str.len() > 0].copy() + empty_removed = old_len - len(df) + df = df[df['summary'].str.len() >= min_length_summary].copy() + short_removed = old_len - empty_removed - len(df) + logger.debug( + f"Removed {empty_removed} empty and {short_removed} short summaries " + f"(< {min_length_summary} chars)." + ) + return df + + @staticmethod + def _remove_regex_dict( + df: pd.DataFrame, + regexes: dict, + postprocess, + keep: int = 0, + ) -> pd.DataFrame: + """Remove regex-matched suffixes or prefixes from summaries. + + For each regex pattern, splits the summary at the first match and keeps + either the left side (keep=0) or the right side (keep=1). Applies a + postprocessing function to each modified summary. Logs statistics for + each pattern at DEBUG level. + + Args: + df: DataFrame with a 'summary' column (string type). + regexes: Dictionary mapping delimiter names to compiled regex + patterns to match against summaries. + postprocess: A callable that takes a string and returns a processed + string, applied after each split. + keep: Which side of the split to keep (0=left/prefix, 1=right/suffix). + Defaults to 0. + + Returns: + The input DataFrame with modified 'summary' column (modified in-place). + """ + total_changed = 0 + for delimiter_name, regex in regexes.items(): + matches = df['summary'].apply(lambda s: regex.search(s) is not None) + total_changed += matches.sum() + logger.debug(f" {delimiter_name}: {matches.sum()} / {len(df)}") + df.loc[matches, 'summary'] = df.loc[matches, 'summary'].apply( + lambda s: regex.split(s, 1)[keep] + ) + df['summary'] = df['summary'].apply(postprocess) + logger.debug(f"Changed total of {total_changed} / {len(df)} summaries.") + return df + + @staticmethod + def _change_why_what_next_pattern_to_text( + summaries: pd.Series, + ) -> pd.Series: + """Convert 'Why / What / Next' dashed lists to paragraph text. + + Transforms structured list blocks matching the 'Why admitted', 'What + was done', and 'What next' patterns into flowing prose by replacing + dashes and line breaks with periods and spaces. + + Args: + summaries: Series of summary strings with potential 'Why/What/Next' + list blocks using dashes or other list markers. + + Returns: + Series of modified summaries with list blocks converted to prose. + """ + random_string = ( + ''.join(random.choices(string.ascii_uppercase + string.digits, k=20)) + + '\n- ' + ) + summaries = summaries.apply( + lambda s: WHY_WHAT_NEXT_HEADINGS_DASHED_LIST.sub(random_string, s) + ) + dash_regex = re.compile(r'(?:\.)?\n-\s{0,4}', re.MULTILINE | re.IGNORECASE) + + def _remove_dashes(s: str) -> str: + """Replace dashes in list items with periods for prose conversion. + + Args: + s: Summary text containing random separator markers. + + Returns: + Text with dashes converted to periods and formatting normalized. + """ + paragraphs = s.split(random_string) + res = [paragraphs[0]] + for p in paragraphs[1:]: + items = p.split('\n\n', 1)[0] + items = '. '.join(dash_regex.split(items)) + if '\n\n' in p: + items = items + '\n\n' + p.split('\n\n', 1)[1] + res.append(items.strip()) + return '\n\n'.join(res) + + return summaries.apply(lambda s: _remove_dashes(s) if random_string in s else s) \ No newline at end of file diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index a32618f9c..8eab605b7 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -42,6 +42,10 @@ MultimodalMortalityPredictionMIMIC4, ) from .survival_preprocess_support2 import SurvivalPreprocessSupport2 +from .mimic4_note_ext_dibhc_tasks import ( + BHCSummarizationTask, + HallucinationDetectionTask, +) from .mortality_prediction_stagenet_mimic4 import ( MortalityPredictionStageNetMIMIC4, ) diff --git a/pyhealth/tasks/mimic4_note_ext_dibhc_tasks.py b/pyhealth/tasks/mimic4_note_ext_dibhc_tasks.py new file mode 100644 index 000000000..1560719b5 --- /dev/null +++ b/pyhealth/tasks/mimic4_note_ext_dibhc_tasks.py @@ -0,0 +1,259 @@ +# -*- coding: utf-8 -*- +"""Tasks for MIMIC-IV-Note patient summary generation and hallucination detection. + +This module implements two PyHealth tasks derived from: + + Hegselmann et al. "A Data-Centric Approach To Generate Faithful and + High Quality Patient Summaries with Large Language Models." + Conference on Health, Inference, and Learning (CHIL) 2024. + https://arxiv.org/abs/2402.15422 + +Both tasks operate on the MIMIC-IV-Note-Ext-DI-BHC dataset, where the +Brief Hospital Course (BHC) serves as input context and the Discharge +Instructions (DI) serve as the target patient summary. + +Tasks defined here: + +1. class:BHCSummarizationTask - Section 4.1 of the paper. + Given a BHC, generate a patient-friendly discharge summary. + Paper results: LED-large achieves ROUGE-1 43.82, GPT-4 0-shot 38.26. + +2. class:HallucinationDetectionTask - Section 4.7 of the paper. + Given a (BHC, summary) pair, predict whether the summary contains + hallucinations (unsupported facts). Paper results: GPT-4 achieves + F1 ~20% at span level; MedCat baseline achieves ~10% F1. + +Example: + >>> from pyhealth.datasets import MIMIC4NoteExtDIBHCDataset + >>> from pyhealth.tasks import BHCSummarizationTask + >>> from pyhealth.tasks import HallucinationDetectionTask + >>> dataset = MIMIC4NoteExtDIBHCDataset(root="/path/to/data/") + >>> summ_samples = dataset.set_task(BHCSummarizationTask()) + >>> halluc_samples = dataset.set_task(HallucinationDetectionTask()) + >>> print(summ_samples[0]["context"][:60]) + >>> print(halluc_samples[0]["label"]) +""" + +from typing import Any, Dict, List + +from pyhealth.tasks import BaseTask + + +class BHCSummarizationTask(BaseTask): + """Patient summary generation from Brief Hospital Course text. + + Given a Brief Hospital Course (BHC) as input context, generates a + patient-friendly Discharge Instructions (DI) summary. This implements + the summarization task from Section 4.1 of Hegselmann et al. (CHIL + 2024), where BHC is used as the context and DI as the target summary. + + The paper demonstrates that fine-tuned LED-large achieves ROUGE-1 of + 43.82 on this task, while GPT-4 0-shot achieves 38.26. We reproduce + this using BART-large as a free open-source alternative. + + Input schema: + - context (str): Brief Hospital Course text — the clinical + notes written by medical staff summarizing the hospital stay. + + Output schema: + - summary (str): Target patient-facing discharge instructions + written in plain language for patient comprehension. + + Note: + This task produces raw text targets for seq2seq generation models. + Evaluate generated summaries using ROUGE and BERTScore metrics. + Average context length is ~552 words; average summary ~113 words + per paper Table 6. + + Example: + >>> from pyhealth.datasets import MIMIC4NoteExtDIBHCDataset + >>> from pyhealth.tasks import BHCSummarizationTask + >>> dataset = MIMIC4NoteExtDIBHCDataset(root="/path/to/data/") + >>> samples = dataset.set_task(BHCSummarizationTask()) + >>> print(samples[0]["context"][:80]) + Brief Hospital Course: Patient presented with chest pain... + >>> print(samples[0]["summary"][:80]) + You were admitted to the hospital for chest pain... + """ + + task_name: str = "BHCSummarizationMIMIC4Note" + + input_schema: Dict[str, str] = { + "context": "str", + } + + output_schema: Dict[str, str] = { + "summary": "str", + } + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + """Process a single patient into summarization samples. + + Iterates over all visits and discharge note events for the patient, + extracting (BHC context, DI summary) pairs. Each valid pair becomes + one training sample for sequence-to-sequence generation. + + Args: + patient (Any): A PyHealth Patient object. Each visit must + contain discharge note events with brief_hospital_course + and summary attributes populated by the dataset class. + + Returns: + List[Dict[str, Any]]: List of sample dicts, one per valid + discharge note event. Each sample contains: + + - patient_id (str): Patient identifier. + - visit_id (str): Visit (admission) identifier. + - context (str): Brief Hospital Course text - model input. + - summary (str): Discharge instructions - generation target. + + Note: + Samples with empty context or summary are silently skipped. + This mirrors the paper's preprocessing pipeline which removes + records with missing BHC or summaries shorter than 350 chars. + + Example: + >>> task = BHCSummarizationTask() + >>> samples = task(patient) + >>> len(samples) + 1 + >>> "context" in samples[0] + True + """ + samples: List[Dict[str, Any]] = [] + + for visit in patient.visits.values(): + for event in visit.get_event_list(source_table="discharge"): + context: str = getattr( + event, "brief_hospital_course", "" + ).strip() + summary: str = getattr(event, "summary", "").strip() + + if not context or not summary: + continue + + samples.append( + { + "patient_id": patient.patient_id, + "visit_id": visit.visit_id, + "context": context, + "summary": summary, + } + ) + + return samples + + +class HallucinationDetectionTask(BaseTask): + """Binary hallucination detection for clinical patient summaries. + + Given a Brief Hospital Course (BHC) as grounding context and a + patient-facing Discharge Instructions (DI) summary, predicts whether + the summary contains at least one hallucination - a fact not supported + by the BHC. This implements the automatic hallucination detection task + from Section 4.7 of Hegselmann et al. (CHIL 2024). + + The paper evaluates hallucination detection at the span level using + expert-annotated datasets (Hallucinations-MIMIC-DI and + Hallucinations-Generated-DI, each with 100 examples). We extend this + to binary document-level classification as a novel contribution. + + Input schema: + - context (str): Brief Hospital Course - the only ground truth + source for determining whether summary facts are supported. + - summary (str): Patient-facing discharge instructions to + evaluate for hallucinations. + + Output schema: + - label (binary): 1 if the summary contains at least one + hallucination span per expert annotation, 0 if faithful. + Defaults to -1 when no expert annotation is available. + + Args: + default_label (int): Label assigned when no expert annotation is + available. Default is -1. + + Example: + >>> from pyhealth.datasets import MIMIC4NoteExtDIBHCDataset + >>> from pyhealth.tasks import HallucinationDetectionTask + >>> dataset = MIMIC4NoteExtDIBHCDataset(root="/path/to/data/") + >>> samples = dataset.set_task(HallucinationDetectionTask()) + >>> print(samples[0]["label"]) + -1 + """ + + task_name: str = "HallucinationDetectionMIMIC4Note" + + input_schema: Dict[str, str] = { + "context": "str", + "summary": "str", + } + + output_schema: Dict[str, str] = { + "label": "binary", + } + + def __init__(self, default_label: int = -1) -> None: + self.default_label = default_label + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + """Process a single patient into hallucination detection samples. + + Iterates over all visits and discharge note events, extracting + (BHC context, DI summary) pairs with binary hallucination labels + where expert annotations are available. + + Args: + patient (Any): A PyHealth Patient object. Each visit must + contain discharge note events with brief_hospital_course + and summary attributes. Optionally, events may include + a has_hallucination attribute (0 or 1) from expert + annotation via the PhysioNet ann-pt-summ dataset. + + Returns: + List[Dict[str, Any]]: List of sample dicts, one per valid + discharge note event. Each sample contains: + + - patient_id (str): Patient identifier. + - visit_id (str): Visit identifier. + - context (str): Brief Hospital Course text. + - summary (str): Patient discharge instructions. + - label (int): 1 hallucination present, 0 faithful, + -1 annotation unavailable. + + Note: + Samples with empty context or summary are silently skipped. + + Example: + >>> task = HallucinationDetectionTask() + >>> samples = task(patient) + >>> samples[0]["label"] in (-1, 0, 1) + True + """ + samples: List[Dict[str, Any]] = [] + + for visit in patient.visits.values(): + for event in visit.get_event_list(source_table="discharge"): + context: str = getattr( + event, "brief_hospital_course", "" + ).strip() + summary: str = getattr(event, "summary", "").strip() + + if not context or not summary: + continue + + label: int = int( + getattr(event, "has_hallucination", self.default_label) + ) + + samples.append( + { + "patient_id": patient.patient_id, + "visit_id": visit.visit_id, + "context": context, + "summary": summary, + "label": label, + } + ) + + return samples diff --git a/tests/core/test_mimic4_note_ext_dibhc_tasks.py b/tests/core/test_mimic4_note_ext_dibhc_tasks.py new file mode 100644 index 000000000..caebe1beb --- /dev/null +++ b/tests/core/test_mimic4_note_ext_dibhc_tasks.py @@ -0,0 +1,910 @@ +# -*- coding: utf-8 -*- +"""Unit tests for MIMIC-IV-Note task definitions (summarization and hallucination detection). + +This module contains comprehensive test coverage for the two PyHealth task classes: +- BHCSummarizationTask: Generates patient-friendly summaries from Brief Hospital Course (BHC) +- HallucinationDetectionTask: Detects hallucinations in discharge instructions + +All tests use synthetic patient data and mocking to avoid MIMIC data dependencies. +Tests verify correct task initialization, schema properties, and sample generation logic. + +Covers Section 4.1 (summarization) and Section 4.7 (hallucination detection) +tasks from: + + Hegselmann et al. "A Data-Centric Approach To Generate Faithful and + High Quality Patient Summaries with Large Language Models." CHIL 2024. +""" + +import unittest +from unittest.mock import Mock + +from pyhealth.tasks.mimic4_note_ext_dibhc_tasks import ( + BHCSummarizationTask, + HallucinationDetectionTask, +) + +# --------------------------------------------------------------------------- +# Shared synthetic data helpers +# --------------------------------------------------------------------------- + +def _make_mock_event( + brief_hospital_course: str = "Patient admitted for observation. " + "Underwent thorough evaluation. Discharged in stable condition.", + summary: str = "You were admitted to the hospital for evaluation and" + " monitoring. During your stay, doctors performed comprehensive testing." + " Your condition improved, and you are now ready to go home.", + has_hallucination: int = -1, +) -> Mock: + """Create a mock discharge event with required attributes. + + Args: + brief_hospital_course (str): Text for the BHC section. + summary (str): Text for the discharge summary section. + has_hallucination (int): Hallucination label (-1, 0, or 1). + + Returns: + Mock object with brief_hospital_course, summary, and has_hallucination attributes. + """ + event = Mock() + event.brief_hospital_course = brief_hospital_course + event.summary = summary + event.has_hallucination = has_hallucination + return event + + +def _make_mock_visit( + visit_id: str = "visit_001", + events: list = None, +) -> Mock: + """Create a mock visit with discharge events. + + Args: + visit_id (str): Identifier for the visit. + events (list): List of mock event objects for this visit. + + Returns: + Mock object representing a visit with discharge events. + """ + if events is None: + events = [_make_mock_event()] + + visit = Mock() + visit.visit_id = visit_id + visit.get_event_list = Mock(return_value=events) + return visit + + +def _make_mock_patient( + patient_id: str = "patient_001", + visits: dict = None, +) -> Mock: + """Create a mock patient with multiple visits. + + Args: + patient_id (str): Identifier for the patient. + visits (dict): Dictionary of mock visit objects. + + Returns: + Mock object representing a patient with visits. + """ + if visits is None: + visits = { + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event()], + ) + } + + patient = Mock() + patient.patient_id = patient_id + patient.visits = visits + return patient + + +# Convenience aliases for existing test data compatibility +def make_event( + brief_hospital_course: str, + summary: str, + has_hallucination: int = -1, +) -> Mock: + """Build a synthetic discharge note event.""" + return _make_mock_event( + brief_hospital_course=brief_hospital_course, + summary=summary, + has_hallucination=has_hallucination, + ) + + +def make_visit(visit_id: str, events: list) -> Mock: + """Build a synthetic visit containing discharge events.""" + return _make_mock_visit(visit_id=visit_id, events=events) + + +def make_patient(patient_id: str, visits: dict) -> Mock: + """Build a synthetic patient with a dict of visits.""" + return _make_mock_patient(patient_id=patient_id, visits=visits) + + +# ======================================================================== +# BHCSummarizationTask Tests (Section 4.1) +# ======================================================================== + +class TestBHCSummarizationTaskInitialization(unittest.TestCase): + """Test suite for BHCSummarizationTask initialization and schema. + + Verifies that BHCSummarizationTask is properly initialized with correct + task_name, input_schema, and output_schema attributes. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestBHCSummarizationTaskInitialization") + print(f"{'='*60}") + + def test_task_name_is_set_correctly(self) -> None: + """Test that task_name attribute is set to expected value.""" + print("\nTEST: test_task_name_is_set_correctly") + task = BHCSummarizationTask() + self.assertEqual(task.task_name, "BHCSummarizationMIMIC4Note") + print(f" task_name: {task.task_name}") + print(" ✓ passed") + + def test_input_schema_has_context_field(self) -> None: + """Test that input_schema contains 'context' field of type 'str'.""" + print("\nTEST: test_input_schema_has_context_field") + task = BHCSummarizationTask() + self.assertIn("context", task.input_schema) + self.assertEqual(task.input_schema["context"], "str") + print(f" input_schema: {task.input_schema}") + print(" ✓ passed") + + def test_output_schema_has_summary_field(self) -> None: + """Test that output_schema contains 'summary' field of type 'str'.""" + print("\nTEST: test_output_schema_has_summary_field") + task = BHCSummarizationTask() + self.assertIn("summary", task.output_schema) + self.assertEqual(task.output_schema["summary"], "str") + print(f" output_schema: {task.output_schema}") + print(" ✓ passed") + + def test_input_schema_only_has_context(self) -> None: + """Test that input_schema contains only 'context' field.""" + print("\nTEST: test_input_schema_only_has_context") + task = BHCSummarizationTask() + self.assertEqual(len(task.input_schema), 1) + print(" ✓ passed") + + def test_output_schema_only_has_summary(self) -> None: + """Test that output_schema contains only 'summary' field.""" + print("\nTEST: test_output_schema_only_has_summary") + task = BHCSummarizationTask() + self.assertEqual(len(task.output_schema), 1) + print(" ✓ passed") + + +class TestBHCSummarizationTaskCallMethod(unittest.TestCase): + """Test suite for BHCSummarizationTask.__call__ method. + + Verifies that the task correctly processes patient objects and generates + summarization samples with proper structure and content. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestBHCSummarizationTaskCallMethod") + print(f"{'='*60}") + self.task = BHCSummarizationTask() + + def test_call_returns_list_of_samples(self) -> None: + """Test that __call__ returns a list of sample dictionaries.""" + print("\nTEST: test_call_returns_list_of_samples") + patient = _make_mock_patient() + result = self.task(patient) + self.assertIsInstance(result, list) + self.assertTrue(len(result) > 0) + self.assertIsInstance(result[0], dict) + print(f" Returned {len(result)} samples") + print(" ✓ passed") + + def test_sample_contains_required_keys(self) -> None: + """Test that each sample contains all required keys.""" + print("\nTEST: test_sample_contains_required_keys") + patient = _make_mock_patient( + patient_id="patient_123", + visits={ + "visit_456": _make_mock_visit( + visit_id="visit_456", + events=[ + _make_mock_event( + brief_hospital_course="Patient was hospitalized.", + summary="You were hospitalized.", + ) + ], + ) + }, + ) + result = self.task(patient) + required_keys = {"patient_id", "visit_id", "context", "summary"} + self.assertTrue(required_keys.issubset(result[0].keys())) + print(f" Sample keys: {list(result[0].keys())}") + print(" ✓ passed") + + def test_sample_content_values_are_correct(self) -> None: + """Test that sample contains correct patient_id, visit_id, and text values.""" + print("\nTEST: test_sample_content_values_are_correct") + bhc_text = "Patient hospitalized for observation and treatment." + summary_text = "You were treated during hospitalization." + patient = _make_mock_patient( + patient_id="patient_xyz", + visits={ + "visit_abc": _make_mock_visit( + visit_id="visit_abc", + events=[ + _make_mock_event( + brief_hospital_course=bhc_text, + summary=summary_text, + ) + ], + ) + }, + ) + result = self.task(patient) + sample = result[0] + self.assertEqual(sample["patient_id"], "patient_xyz") + self.assertEqual(sample["visit_id"], "visit_abc") + self.assertEqual(sample["context"], bhc_text) + self.assertEqual(sample["summary"], summary_text) + print(f" patient_id: {sample['patient_id']}") + print(f" visit_id: {sample['visit_id']}") + print(" ✓ passed") + + def test_skips_empty_context(self) -> None: + """Test that samples with empty context are skipped.""" + print("\nTEST: test_skips_empty_context") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="", + summary="Valid summary text here.", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Empty context event was skipped") + print(" ✓ passed") + + def test_skips_empty_summary(self) -> None: + """Test that samples with empty summary are skipped.""" + print("\nTEST: test_skips_empty_summary") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="Valid hospital course.", + summary="", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Empty summary event was skipped") + print(" ✓ passed") + + def test_skips_whitespace_only_context(self) -> None: + """Test that samples with whitespace-only context are skipped.""" + print("\nTEST: test_skips_whitespace_only_context") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course=" \n\t ", + summary="Valid summary.", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Whitespace-only context event was skipped") + print(" ✓ passed") + + def test_skips_whitespace_only_summary(self) -> None: + """Test that samples with whitespace-only summary are skipped.""" + print("\nTEST: test_skips_whitespace_only_summary") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="Valid context.", + summary=" \n\t ", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Whitespace-only summary event was skipped") + print(" ✓ passed") + + def test_strips_leading_trailing_whitespace(self) -> None: + """Test that leading/trailing whitespace is stripped from context and summary.""" + print("\nTEST: test_strips_leading_trailing_whitespace") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course=" \n BHC text \n ", + summary=" \n Summary text \n ", + ) + ], + ) + }, + ) + result = self.task(patient) + sample = result[0] + self.assertEqual(sample["context"], "BHC text") + self.assertEqual(sample["summary"], "Summary text") + print(" Whitespace was properly stripped") + print(" ✓ passed") + + def test_handles_multiple_visits(self) -> None: + """Test that task processes multiple visits for a single patient.""" + print("\nTEST: test_handles_multiple_visits") + patient = _make_mock_patient( + patient_id="multi_visit_patient", + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event()], + ), + "visit_002": _make_mock_visit( + visit_id="visit_002", + events=[_make_mock_event()], + ), + "visit_003": _make_mock_visit( + visit_id="visit_003", + events=[_make_mock_event()], + ), + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 3) + visit_ids = {sample["visit_id"] for sample in result} + self.assertEqual(visit_ids, {"visit_001", "visit_002", "visit_003"}) + print(f" Processed {len(result)} visits") + print(" ✓ passed") + + def test_handles_multiple_events_per_visit(self) -> None: + """Test that task processes multiple events within a single visit.""" + print("\nTEST: test_handles_multiple_events_per_visit") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event(brief_hospital_course="BHC1", summary="Summary1"), + _make_mock_event(brief_hospital_course="BHC2", summary="Summary2"), + _make_mock_event(brief_hospital_course="BHC3", summary="Summary3"), + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 3) + contexts = {sample["context"] for sample in result} + self.assertEqual(contexts, {"BHC1", "BHC2", "BHC3"}) + print(f" Processed {len(result)} events") + print(" ✓ passed") + + def test_handles_missing_attributes(self) -> None: + """Test that task handles events with missing BHC or summary attributes.""" + print("\nTEST: test_handles_missing_attributes") + event_no_bhc = Mock() + event_no_bhc.brief_hospital_course = "" # Explicitly set to empty + event_no_bhc.summary = "Has summary but no BHC" + # Simulate missing brief_hospital_course by setting it to empty + + event_no_summary = Mock() + event_no_summary.brief_hospital_course = "Has BHC but no summary" + event_no_summary.summary = "" # Explicitly set to empty + # Simulate missing summary by setting it to empty + + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[event_no_bhc, event_no_summary], + ) + }, + ) + result = self.task(patient) + # Both should be skipped due to empty attributes + self.assertEqual(len(result), 0) + print(" Events with empty attributes were skipped") + print(" ✓ passed") + + +# ======================================================================== +# HallucinationDetectionTask Tests (Section 4.7) +# ======================================================================== + +class TestHallucinationDetectionTaskInitialization(unittest.TestCase): + """Test suite for HallucinationDetectionTask initialization and schema. + + Verifies that HallucinationDetectionTask is properly initialized with correct + task_name, input_schema, output_schema, and default_label attributes. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestHallucinationDetectionTaskInitialization") + print(f"{'='*60}") + + def test_task_name_is_set_correctly(self) -> None: + """Test that task_name attribute is set to expected value.""" + print("\nTEST: test_task_name_is_set_correctly") + task = HallucinationDetectionTask() + self.assertEqual(task.task_name, "HallucinationDetectionMIMIC4Note") + print(f" task_name: {task.task_name}") + print(" ✓ passed") + + def test_input_schema_has_context_and_summary(self) -> None: + """Test that input_schema contains 'context' and 'summary' fields.""" + print("\nTEST: test_input_schema_has_context_and_summary") + task = HallucinationDetectionTask() + self.assertIn("context", task.input_schema) + self.assertIn("summary", task.input_schema) + self.assertEqual(task.input_schema["context"], "str") + self.assertEqual(task.input_schema["summary"], "str") + print(f" input_schema: {task.input_schema}") + print(" ✓ passed") + + def test_output_schema_has_label_field(self) -> None: + """Test that output_schema contains 'label' field of type 'binary'.""" + print("\nTEST: test_output_schema_has_label_field") + task = HallucinationDetectionTask() + self.assertIn("label", task.output_schema) + self.assertEqual(task.output_schema["label"], "binary") + print(f" output_schema: {task.output_schema}") + print(" ✓ passed") + + def test_default_label_is_minus_one_by_default(self) -> None: + """Test that default_label is -1 when not specified.""" + print("\nTEST: test_default_label_is_minus_one_by_default") + task = HallucinationDetectionTask() + self.assertEqual(task.default_label, -1) + print(f" default_label: {task.default_label}") + print(" ✓ passed") + + def test_default_label_can_be_customized(self) -> None: + """Test that default_label can be set via constructor.""" + print("\nTEST: test_default_label_can_be_customized") + task = HallucinationDetectionTask(default_label=0) + self.assertEqual(task.default_label, 0) + print(f" custom default_label: {task.default_label}") + print(" ✓ passed") + + def test_custom_default_labels(self) -> None: + """Test various custom default_label values.""" + print("\nTEST: test_custom_default_labels") + for label in [-1, 0, 1]: + task = HallucinationDetectionTask(default_label=label) + self.assertEqual(task.default_label, label) + print(" All custom labels set correctly") + print(" ✓ passed") + + +class TestHallucinationDetectionTaskCallMethod(unittest.TestCase): + """Test suite for HallucinationDetectionTask.__call__ method. + + Verifies that the task correctly processes patient objects and generates + hallucination detection samples with proper structure and label values. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestHallucinationDetectionTaskCallMethod") + print(f"{'='*60}") + self.task = HallucinationDetectionTask() + + def test_call_returns_list_of_samples(self) -> None: + """Test that __call__ returns a list of sample dictionaries.""" + print("\nTEST: test_call_returns_list_of_samples") + patient = _make_mock_patient() + result = self.task(patient) + self.assertIsInstance(result, list) + self.assertTrue(len(result) > 0) + self.assertIsInstance(result[0], dict) + print(f" Returned {len(result)} samples") + print(" ✓ passed") + + def test_sample_contains_required_keys(self) -> None: + """Test that each sample contains all required keys.""" + print("\nTEST: test_sample_contains_required_keys") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event()], + ) + }, + ) + result = self.task(patient) + required_keys = {"patient_id", "visit_id", "context", "summary", "label"} + self.assertTrue(required_keys.issubset(result[0].keys())) + print(f" Sample keys: {list(result[0].keys())}") + print(" ✓ passed") + + def test_sample_content_values_are_correct(self) -> None: + """Test that sample contains correct patient_id, visit_id, and text values.""" + print("\nTEST: test_sample_content_values_are_correct") + bhc_text = "Patient hospitalized for treatment." + summary_text = "You were treated in the hospital." + patient = _make_mock_patient( + patient_id="patient_xyz", + visits={ + "visit_abc": _make_mock_visit( + visit_id="visit_abc", + events=[ + _make_mock_event( + brief_hospital_course=bhc_text, + summary=summary_text, + has_hallucination=-1, + ) + ], + ) + }, + ) + result = self.task(patient) + sample = result[0] + self.assertEqual(sample["patient_id"], "patient_xyz") + self.assertEqual(sample["visit_id"], "visit_abc") + self.assertEqual(sample["context"], bhc_text) + self.assertEqual(sample["summary"], summary_text) + print(f" patient_id: {sample['patient_id']}") + print(f" visit_id: {sample['visit_id']}") + print(" ✓ passed") + + def test_label_zero_for_no_hallucination(self) -> None: + """Test that label is 0 when has_hallucination is 0.""" + print("\nTEST: test_label_zero_for_no_hallucination") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event(has_hallucination=0)], + ) + }, + ) + result = self.task(patient) + self.assertEqual(result[0]["label"], 0) + print(" Label correctly set to 0 for faithful summary") + print(" ✓ passed") + + def test_label_one_for_hallucination_present(self) -> None: + """Test that label is 1 when has_hallucination is 1.""" + print("\nTEST: test_label_one_for_hallucination_present") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event(has_hallucination=1)], + ) + }, + ) + result = self.task(patient) + self.assertEqual(result[0]["label"], 1) + print(" Label correctly set to 1 for hallucination present") + print(" ✓ passed") + + def test_label_default_when_annotation_unavailable(self) -> None: + """Test that label uses default_label when annotation is unavailable.""" + print("\nTEST: test_label_default_when_annotation_unavailable") + task = HallucinationDetectionTask(default_label=-1) + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event(has_hallucination=-1)], + ) + }, + ) + result = task(patient) + self.assertEqual(result[0]["label"], -1) + print(" Label correctly set to default_label (-1)") + print(" ✓ passed") + + def test_custom_default_label_zero(self) -> None: + """Test that custom default_label=0 is used when annotation unavailable.""" + print("\nTEST: test_custom_default_label_zero") + task = HallucinationDetectionTask(default_label=0) + + # Use a simple class to avoid Mock auto-creating undefined attributes + class SimpleEvent: + def __init__(self): + self.brief_hospital_course = "Patient was treated." + self.summary = "You were treated." + # has_hallucination is intentionally not set + + event = SimpleEvent() + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[event], + ) + }, + ) + result = task(patient) + self.assertEqual(result[0]["label"], 0) + print(" Custom default_label (0) correctly applied") + print(" ✓ passed") + + def test_skips_empty_context(self) -> None: + """Test that samples with empty context are skipped.""" + print("\nTEST: test_skips_empty_context") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="", + summary="Valid summary.", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Empty context event was skipped") + print(" ✓ passed") + + def test_skips_empty_summary(self) -> None: + """Test that samples with empty summary are skipped.""" + print("\nTEST: test_skips_empty_summary") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="Valid context.", + summary="", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Empty summary event was skipped") + print(" ✓ passed") + + def test_skips_both_empty_context_and_summary(self) -> None: + """Test that samples with both empty context and summary are skipped.""" + print("\nTEST: test_skips_both_empty_context_and_summary") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="", + summary="", + ) + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Event with both empty context and summary was skipped") + print(" ✓ passed") + + def test_strips_leading_trailing_whitespace(self) -> None: + """Test that leading/trailing whitespace is stripped.""" + print("\nTEST: test_strips_leading_trailing_whitespace") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course=" \n BHC text \n ", + summary=" \n Summary text \n ", + ) + ], + ) + }, + ) + result = self.task(patient) + sample = result[0] + self.assertEqual(sample["context"], "BHC text") + self.assertEqual(sample["summary"], "Summary text") + print(" Whitespace was properly stripped") + print(" ✓ passed") + + def test_handles_multiple_visits(self) -> None: + """Test that task processes multiple visits for a single patient.""" + print("\nTEST: test_handles_multiple_visits") + patient = _make_mock_patient( + patient_id="multi_visit_patient", + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[_make_mock_event(has_hallucination=0)], + ), + "visit_002": _make_mock_visit( + visit_id="visit_002", + events=[_make_mock_event(has_hallucination=1)], + ), + "visit_003": _make_mock_visit( + visit_id="visit_003", + events=[_make_mock_event(has_hallucination=-1)], + ), + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 3) + labels = {sample["label"] for sample in result} + self.assertEqual(labels, {0, 1, -1}) + visit_ids = {sample["visit_id"] for sample in result} + self.assertEqual(visit_ids, {"visit_001", "visit_002", "visit_003"}) + print(f" Processed {len(result)} visits with correct labels") + print(" ✓ passed") + + def test_handles_multiple_events_per_visit(self) -> None: + """Test that task processes multiple events within a single visit.""" + print("\nTEST: test_handles_multiple_events_per_visit") + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="BHC1", + summary="Summary1", + has_hallucination=0, + ), + _make_mock_event( + brief_hospital_course="BHC2", + summary="Summary2", + has_hallucination=1, + ), + _make_mock_event( + brief_hospital_course="BHC3", + summary="Summary3", + has_hallucination=-1, + ), + ], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 3) + labels = {sample["label"] for sample in result} + self.assertEqual(labels, {0, 1, -1}) + print(f" Processed {len(result)} events with correct labels") + print(" ✓ passed") + + def test_handles_missing_attributes(self) -> None: + """Test that task handles events with missing attributes gracefully.""" + print("\nTEST: test_handles_missing_attributes") + event_no_bhc = Mock() + event_no_bhc.brief_hospital_course = "" # Explicitly set to empty + event_no_bhc.summary = "Has summary but no BHC" + event_no_bhc.has_hallucination = -1 + + event_no_summary = Mock() + event_no_summary.brief_hospital_course = "Has BHC but no summary" + event_no_summary.summary = "" # Explicitly set to empty + event_no_summary.has_hallucination = -1 + + patient = _make_mock_patient( + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[event_no_bhc, event_no_summary], + ) + }, + ) + result = self.task(patient) + self.assertEqual(len(result), 0) + print(" Events with empty attributes were skipped") + print(" ✓ passed") + + +# ======================================================================== +# Integration Tests +# ======================================================================== + +class TestTaskIntegration(unittest.TestCase): + """Integration tests for both summarization and hallucination detection tasks. + + Tests interactions between the two task classes and ensures they work + correctly with the same patient data. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestTaskIntegration") + print(f"{'='*60}") + + def test_both_tasks_process_same_patient(self) -> None: + """Test that both tasks can process the same patient object.""" + print("\nTEST: test_both_tasks_process_same_patient") + summ_task = BHCSummarizationTask() + halluc_task = HallucinationDetectionTask() + + patient = _make_mock_patient( + patient_id="patient_001", + visits={ + "visit_001": _make_mock_visit( + visit_id="visit_001", + events=[ + _make_mock_event( + brief_hospital_course="Admitted for chest pain.", + summary="Treated for chest pain.", + has_hallucination=0, + ) + ], + ) + }, + ) + + summ_samples = summ_task(patient) + halluc_samples = halluc_task(patient) + + self.assertEqual(len(summ_samples), 1) + self.assertEqual(len(halluc_samples), 1) + self.assertEqual(summ_samples[0]["patient_id"], halluc_samples[0]["patient_id"]) + print(" Both tasks successfully processed the same patient") + print(" ✓ passed") + + def test_different_tasks_have_different_schemas(self) -> None: + """Test that the two tasks have different input/output schemas.""" + print("\nTEST: test_different_tasks_have_different_schemas") + summ_task = BHCSummarizationTask() + halluc_task = HallucinationDetectionTask() + + # Summarization task has only context as input + self.assertEqual(len(summ_task.input_schema), 1) + self.assertIn("context", summ_task.input_schema) + + # Hallucination task has context and summary as inputs + self.assertEqual(len(halluc_task.input_schema), 2) + self.assertIn("context", halluc_task.input_schema) + self.assertIn("summary", halluc_task.input_schema) + + # Summarization output is summary (str) + self.assertEqual(summ_task.output_schema["summary"], "str") + + # Hallucination output is label (binary) + self.assertEqual(halluc_task.output_schema["label"], "binary") + + print(" Schemas correctly differ between tasks") + print(" ✓ passed") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_mimic4noteextdibhc.py b/tests/core/test_mimic4noteextdibhc.py new file mode 100644 index 000000000..db6a779c9 --- /dev/null +++ b/tests/core/test_mimic4noteextdibhc.py @@ -0,0 +1,738 @@ +"""Unit tests for MIMIC4NoteExtDIBHCDataset preprocessing pipeline. + +This module contains comprehensive test coverage for the MIMIC-IV Extracted +Discharge Instructions and Brief Hospital Course (DIBHC) dataset class, +including: + +- Static helper method tests (extraction, filtering) +- Pipeline step-by-step tests (preprocessing stages 0-7) +- End-to-end integration tests + +All tests use synthetic data to avoid MIMIC data dependencies. A lightweight +_DummyDataset class exposes the preprocessing pipeline without requiring +BaseDataset initialization or filesystem access. +""" + +import unittest +import os +import re +import pandas as pd +from pathlib import Path +from unittest.mock import patch, MagicMock + +from pyhealth.datasets import MIMIC4NoteExtDIBHCDataset + + +# --------------------------------------------------------------------------- +# Shared synthetic data helpers +# --------------------------------------------------------------------------- + +def _make_minimal_note( + hospital_course: str = ( + "Brief Hospital Course: Patient was admitted for chest pain. " + "Workup showed no acute MI. Patient was monitored and stabilized. " + "Cardiology was consulted and agreed with management plan.\n" + "Medications on Admission: aspirin 81mg" + ), + discharge_section: str = ( + "Dear ___, It was a pleasure caring for you during your stay. " + "You were admitted because you had chest pain. " + "We ran several tests and found no signs of a heart attack. " + "You were monitored closely and your condition improved. " + "Please take all your medications as prescribed and follow up " + "with your primary care physician within one week of discharge. " + "If you experience worsening chest pain, shortness of breath, " + "or any other concerning symptoms, please call your doctor or " + "return to the emergency department immediately." + ), +) -> str: + """Create a minimal synthetic discharge note with required structure. + + Combines a hospital course section with discharge instructions to mimic + the structure of real MIMIC-IV discharge notes. + + Args: + hospital_course: Text for the hospital course section before + 'Discharge Instructions:' marker. Defaults to a simple chest + pain admission note. + discharge_section: Text for the discharge instructions section after + the marker. Defaults to standard discharge instructions. + + Returns: + A complete discharge note string with both required sections. + """ + return f"{hospital_course}\nDischarge Instructions:\n{discharge_section}" + + +def _make_note_df(n: int = 5) -> pd.DataFrame: + """Create a DataFrame with synthetic discharge notes. + + Generates n synthetic discharge notes with varied hospital course + descriptions and discharge instructions, useful for testing preprocessing + stages without needing real MIMIC data. + + Args: + n: Number of synthetic notes to generate. Defaults to 5. + + Returns: + A pandas DataFrame with columns: note_id, subject_id, hadm_id, text. + Each row represents one complete synthetic discharge note. + """ + rows = [] + for i in range(n): + rows.append({ + "note_id": str(i), + "subject_id": str(1000 + i), + "hadm_id": str(2000 + i), + "text": _make_minimal_note( + hospital_course=( + f"Brief Hospital Course: Patient {i} was admitted for evaluation. " + f"They were treated appropriately and discharged in " + f"stable condition. " + f"All relevant workup was completed during the hospital stay. " + f"The team reviewed results daily and adjusted therapy as needed.\n" + f"Medications on Admission: lisinopril 10mg" + ), + discharge_section=( + f"You were admitted to the hospital because you were experiencing " + f"symptoms that required further evaluation and treatment. " + f"During your stay, we performed a thorough workup and started " + f"appropriate therapy. Your condition improved significantly. " + f"Please make sure to attend all follow-up appointments and " + f"continue taking your medications as directed. " + f"If you develop any new or worsening symptoms, please seek " + f"medical attention promptly. We wish you a speedy recovery." + ), + ), + }) + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Unit tests +# --------------------------------------------------------------------------- + +class TestMIMIC4NoteExtDIBHCDatasetStaticHelpers(unittest.TestCase): + """Test suite for static helper methods. + + Tests static utility methods of MIMIC4NoteExtDIBHCDataset that do not + require a loaded dataset or filesystem access. Covers extraction, + filtering, and validation functions. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestMIMIC4NoteExtDIBHCDatasetStaticHelpers") + print(f"{'='*60}") + + # ------------------------------------------------------------------ + # _extract_hc + # ------------------------------------------------------------------ + + def test_extract_hc_returns_text_between_markers(self) -> None: + """Test that _extract_hc extracts text between section markers. + + Verifies that _extract_hc correctly locates and extracts text + between 'Brief Hospital Course:' and the next known marker. + Requires full note text to be >= 30 words to pass validation. + """ + print("\nTEST: test_extract_hc_returns_text_between_markers") + # Note: _extract_hc checks len(txt.split(' ')) >= 30, so the full note + # text (not just the BHC section) must be at least 30 words. + txt = ( + "Admission Date: ___ Discharge Date: ___\n" + "Service: MEDICINE\n" + "Brief Hospital Course: Patient was admitted and treated successfully " + "with IV antibiotics over a five-day course. Cultures returned negative " + "and the patient improved clinically throughout the admission.\n" + "Medications on Admission: lisinopril 10mg\n" + ) + self.assertGreaterEqual( + len(txt.split(" ")), 30, + msg="Synthetic note must have >= 30 words to pass _extract_hc's " + "length guard" + ) + result = MIMIC4NoteExtDIBHCDataset._extract_hc(txt) + self.assertIsNotNone(result) + self.assertIn("admitted and treated", result) + print(f" Extracted: {result[:80]}...") + print(" ✓ passed") + + def test_extract_hc_returns_none_when_marker_absent(self) -> None: + """Test that _extract_hc returns None when marker is missing. + + Verifies that _extract_hc gracefully handles notes without the + 'Brief Hospital Course:' marker by returning None. + """ + print("\nTEST: test_extract_hc_returns_none_when_marker_absent") + txt = "This note has no relevant section.\nMedications on Admission: aspirin" + result = MIMIC4NoteExtDIBHCDataset._extract_hc(txt) + self.assertIsNone(result) + print(" ✓ returned None as expected") + + def test_extract_hc_returns_none_for_very_short_text(self) -> None: + """Test that _extract_hc rejects notes with insufficient text. + + Verifies that _extract_hc enforces the minimum word count threshold + (30 words) and returns None for notes falling below this limit. + """ + print("\nTEST: test_extract_hc_returns_none_for_very_short_text") + txt = "Brief Hospital Course: Short.\nMedications on Admission: x" + self.assertLess(len(txt.split(" ")), 30) + result = MIMIC4NoteExtDIBHCDataset._extract_hc(txt) + self.assertIsNone(result) + print(" ✓ returned None for very short note") + + def test_extract_hc_falls_back_to_discharge_medications(self) -> None: + """Test _extract_hc fallback to 'Discharge Medications:' marker. + + Verifies that when 'Medications on Admission:' is absent, _extract_hc + uses 'Discharge Medications:' as the end marker. + """ + print("\nTEST: test_extract_hc_falls_back_to_discharge_medications") + txt = ( + "Brief Hospital Course: " + + ("The patient was treated and improved. " * 10) + + "\n" + "Discharge Medications: metoprolol 25mg\n" + ) + result = MIMIC4NoteExtDIBHCDataset._extract_hc(txt) + self.assertIsNotNone(result) + self.assertNotIn("Discharge Medications", result) + print(f" Extracted length: {len(result)} chars") + print(" ✓ passed") + + def test_extract_hc_falls_back_to_discharge_disposition(self) -> None: + """Test _extract_hc fallback to 'Discharge Disposition:' marker. + + Verifies that when both previous markers are absent, _extract_hc + uses 'Discharge Disposition:' as the final fallback end marker. + """ + print("\nTEST: test_extract_hc_falls_back_to_discharge_disposition") + txt = ( + "Brief Hospital Course: " + + ("Patient recovered well and was ready for discharge. " * 8) + + "\n" + "Discharge Disposition: Home\n" + ) + result = MIMIC4NoteExtDIBHCDataset._extract_hc(txt) + self.assertIsNotNone(result) + self.assertNotIn("Discharge Disposition", result) + print(" ✓ passed") + + # ------------------------------------------------------------------ + # _remove_empty_and_short_summaries + # ------------------------------------------------------------------ + + def test_remove_empty_and_short_summaries_drops_short(self) -> None: + """Test that short summaries are removed by filtering. + + Verifies that _remove_empty_and_short_summaries correctly drops + rows where summary text is shorter than the specified threshold. + """ + print("\nTEST: test_remove_empty_and_short_summaries_drops_short") + df = pd.DataFrame({"summary": ["short", "x" * 350, "x" * 400, ""]}) + result = MIMIC4NoteExtDIBHCDataset._remove_empty_and_short_summaries( + df, min_length_summary=350 + ) + self.assertEqual(len(result), 2) + self.assertTrue(all(result["summary"].str.len() >= 350)) + print(f" Input rows: 4, output rows: {len(result)}") + print(" ✓ passed") + + def test_remove_empty_and_short_summaries_keeps_all_if_long_enough(self) -> None: + """Test that all long summaries are retained by filtering. + + Verifies that _remove_empty_and_short_summaries keeps all rows when + they meet the minimum length threshold. + """ + print("\nTEST: test_remove_empty_and_short_summaries_keeps_all_if_long_enough") + df = pd.DataFrame({"summary": ["x" * 400, "x" * 500, "x" * 600]}) + result = MIMIC4NoteExtDIBHCDataset._remove_empty_and_short_summaries(df) + self.assertEqual(len(result), 3) + print(" ✓ passed") + + +# --------------------------------------------------------------------------- + +class TestMIMIC4NoteExtDIBHCDatasetPipelineSteps(unittest.TestCase): + """Test suite for individual preprocessing pipeline steps. + + Tests each preprocessing step (0-6) using synthetic DataFrames. + Verifies that each step produces expected transformations and maintains + data integrity. No MIMIC data or filesystem access is required. + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestMIMIC4NoteExtDIBHCDatasetPipelineSteps") + print(f"{'='*60}") + self.pipeline = _DummyDataset() + + # ------------------------------------------------------------------ + # Step 0 + # ------------------------------------------------------------------ + + def test_step0_replaces_special_chars(self) -> None: + """Test that Step 0 replaces Unicode special characters. + + Verifies that _step0_special_chars correctly replaces known + non-ASCII characters with ASCII equivalents. + """ + print("\nTEST: test_step0_replaces_special_chars") + df = pd.DataFrame({"text": [u"Hello\u0091world\u0093end"]}) + result = MIMIC4NoteExtDIBHCDataset._step0_special_chars(df) + self.assertNotIn(u"\u0091", result["text"].iloc[0]) + self.assertNotIn(u"\u0093", result["text"].iloc[0]) + self.assertIn("'", result["text"].iloc[0]) + self.assertIn('"', result["text"].iloc[0]) + print(f" Converted: {result['text'].iloc[0]}") + print(" ✓ passed") + + def test_step0_strips_whitespace(self) -> None: + """Test that Step 0 strips leading/trailing whitespace. + + Verifies that _step0_special_chars removes leading and trailing + whitespace from note text. + """ + print("\nTEST: test_step0_strips_whitespace") + df = pd.DataFrame({"text": [" hello world "]}) + result = MIMIC4NoteExtDIBHCDataset._step0_special_chars(df) + self.assertEqual(result["text"].iloc[0], "hello world") + print(" ✓ passed") + + # ------------------------------------------------------------------ + # Step 1 + # ------------------------------------------------------------------ + + def test_step1_splits_correctly(self) -> None: + """Test that Step 1 correctly splits on discharge marker. + + Verifies that _step1_split_on_discharge_instructions correctly + separates hospital course from summary sections. + """ + print("\nTEST: test_step1_splits_correctly") + df = pd.DataFrame({"text": [_make_minimal_note()]}) + result = self.pipeline._step1_split_on_discharge_instructions(df) + self.assertIn("hospital_course", result.columns) + self.assertIn("summary", result.columns) + self.assertGreater(result["hospital_course"].str.len().iloc[0], 0) + self.assertGreater(result["summary"].str.len().iloc[0], 0) + print( + f" hospital_course length: " + f"{result['hospital_course'].str.len().iloc[0]}" + ) + print(f" summary length: {result['summary'].str.len().iloc[0]}") + print(" ✓ passed") + + def test_step1_drops_notes_without_marker(self) -> None: + """Test that Step 1 removes notes lacking discharge marker. + + Verifies that _step1_split_on_discharge_instructions correctly + drops notes that lack the 'Discharge Instructions:' marker. + """ + print("\nTEST: test_step1_drops_notes_without_marker") + df = pd.DataFrame({"text": [ + _make_minimal_note(), + "A note with no discharge instructions section at all.", + ]}) + result = self.pipeline._step1_split_on_discharge_instructions(df) + self.assertEqual(len(result), 1) + print(f" Input rows: 2, output rows: {len(result)}") + print(" ✓ passed") + + # ------------------------------------------------------------------ + # Step 2 + # ------------------------------------------------------------------ + + def test_step2_encodes_dr_abbreviation(self) -> None: + """Test that Step 2 encodes 'Dr.' abbreviation. + + Verifies that _step2_encode_and_extract_hc temporarily encodes + the 'Dr.' abbreviation to prevent sentence tokenization errors. + """ + print("\nTEST: test_step2_encodes_dr_abbreviation") + df = _make_note_df(3) + df = self.pipeline._step1_split_on_discharge_instructions(df) + df.at[df.index[0], "summary"] = "Dr. Smith reviewed your case. " + "x" * 350 + result = MIMIC4NoteExtDIBHCDataset._step2_encode_and_extract_hc(df) + print(f" Rows after step 2: {len(result)}") + print(" ✓ passed") + + def test_step2_populates_brief_hospital_course(self) -> None: + """Test that Step 2 extracts brief hospital course section. + + Verifies that _step2_encode_and_extract_hc correctly populates + the brief_hospital_course column from hospital_course text. + """ + print("\nTEST: test_step2_populates_brief_hospital_course") + df = _make_note_df(3) + df = self.pipeline._step1_split_on_discharge_instructions(df) + result = MIMIC4NoteExtDIBHCDataset._step2_encode_and_extract_hc(df) + self.assertIn("brief_hospital_course", result.columns) + non_null = result["brief_hospital_course"].notnull().sum() + print(f" Rows with brief_hospital_course populated: {non_null}/{len(result)}") + self.assertGreater(non_null, 0) + print(" ✓ passed") + + # ------------------------------------------------------------------ + # Step 4 + # ------------------------------------------------------------------ + + def test_step4_collapses_multiple_spaces(self) -> None: + """Test that Step 4 normalizes whitespace. + + Verifies that _step4_remove_static_patterns correctly collapses + multiple consecutive spaces into single spaces. + """ + print("\nTEST: test_step4_collapses_multiple_spaces") + long_pad = "x" * 350 + df = pd.DataFrame({"summary": [f"Word1 Word2 Word3. {long_pad}"]}) + result = MIMIC4NoteExtDIBHCDataset._step4_remove_static_patterns(df) + if len(result) > 0: + self.assertNotIn(" ", result["summary"].iloc[0]) + print(" ✓ passed") + + def test_step4_applies_deidentification(self) -> None: + """Test that Step 4 applies de-identification replacements. + + Verifies that _step4_remove_static_patterns correctly replaces + '___ ' with contextual pronouns based on following verbs. + """ + print("\nTEST: test_step4_applies_deidentification") + filler = "A" * 350 + sentence = "___ were admitted to the hospital for chest pain evaluation." + df = pd.DataFrame({"summary": [sentence + " " + filler]}) + result = MIMIC4NoteExtDIBHCDataset._step4_remove_static_patterns(df) + if len(result) > 0: + text = result["summary"].iloc[0] + self.assertIn("You were admitted", text) + print(f" De-identified: {text[:60]}") + print(" ✓ passed") + + # ------------------------------------------------------------------ + # Step 6 + # ------------------------------------------------------------------ + + def test_step6_drops_too_short_summaries(self) -> None: + """Test that Step 6 enforces minimum summary length. + + Verifies that _step6_quality_filter correctly removes summaries + below the minimum character count threshold. + """ + print("\nTEST: test_step6_drops_too_short_summaries") + pipeline = _DummyDataset(min_chars=500) + df = pd.DataFrame({ + "summary": [ + "Short text that won't pass.", + "y" * 600 + " end.", + ] + }) + result = pipeline._step6_quality_filter(df) + self.assertTrue(all(result["summary"].str.len() >= 500)) + print(f" Rows after quality filter: {len(result)}") + print(" ✓ passed") + + def test_step6_drops_deidentification_dense_summaries(self) -> None: + """Test that Step 6 filters de-identification-dense summaries. + + Verifies that _step6_quality_filter removes summaries with + excessive '___' placeholder tokens relative to word count. + """ + print("\nTEST: test_step6_drops_deidentification_dense_summaries") + pipeline = _DummyDataset( + min_chars=100, + min_sentences=1, + num_words_per_deidentified=10, + ) + # One summary with many ___ tokens (1 per ~5 words) → should be dropped + noisy = " ".join(["word ___ word ___ word"] * 20) + "." + # One clean summary + clean = ("This patient was admitted and treated well. " * 10) + "." + df = pd.DataFrame({"summary": [noisy, clean]}) + result = pipeline._step6_quality_filter(df) + kept = result["summary"].tolist() + self.assertFalse( + any( + "___" in s and s.count("___") > len(s.split()) / 10 + for s in kept + ) + ) + print(f" Rows after deidentification filter: {len(result)}") + print(" ✓ passed") + + # ------------------------------------------------------------------ + # Step 7 + # ------------------------------------------------------------------ + + def test_step7_drops_null_hospital_course(self) -> None: + """Test that Step 7 removes records with missing courses. + + Verifies that _step7_filter_hospital_course correctly drops + records where hospital_course or brief_hospital_course are null. + """ + print("\nTEST: test_step7_drops_null_hospital_course") + pipeline = _DummyDataset(min_chars_bhc=10) + df = pd.DataFrame({ + "summary": ["ok summary"] * 3, + "hospital_course": ["some course", None, "another course"], + "brief_hospital_course": ["brief text here", "also brief", None], + }) + result = pipeline._step7_filter_hospital_course(df) + self.assertEqual(len(result), 1) + print(f" Input rows: 3, output rows: {len(result)}") + print(" ✓ passed") + + def test_step7_drops_short_brief_hospital_course(self) -> None: + """Test that Step 7 enforces minimum brief hospital course length. + + Verifies that _step7_filter_hospital_course correctly removes + records whose brief_hospital_course is below minimum threshold. + """ + print("\nTEST: test_step7_drops_short_brief_hospital_course") + pipeline = _DummyDataset(min_chars_bhc=500) + df = pd.DataFrame({ + "summary": ["ok"] * 2, + "hospital_course": ["course a", "course b"], + "brief_hospital_course": ["too short", "x" * 600], + }) + result = pipeline._step7_filter_hospital_course(df) + self.assertEqual(len(result), 1) + self.assertGreaterEqual( + len(result["brief_hospital_course"].iloc[0]), 500 + ) + print(" ✓ passed") + + def test_step7_normalises_excessive_blank_lines(self) -> None: + """Test that Step 7 normalizes excessive newlines. + + Verifies that _step7_filter_hospital_course correctly collapses + 3+ consecutive newlines down to exactly 2. + """ + print("\nTEST: test_step7_normalises_excessive_blank_lines") + pipeline = _DummyDataset(min_chars_bhc=5) + df = pd.DataFrame({ + "summary": ["ok"], + "hospital_course": ["line1\n\n\n\nline2"], + "brief_hospital_course": ["brief\n\n\n\nmore"], + }) + result = pipeline._step7_filter_hospital_course(df) + self.assertNotIn("\n\n\n", result["hospital_course"].iloc[0]) + self.assertNotIn("\n\n\n", result["brief_hospital_course"].iloc[0]) + print(" ✓ passed") + + +# --------------------------------------------------------------------------- + +class TestMIMIC4NoteExtDIBHCDatasetEndToEnd(unittest.TestCase): + """End-to-end integration test suite for the full preprocessing pipeline. + + Tests the complete preprocess() method on synthetic data to verify: + - Pipeline completion without errors + - Correct output column creation + - Input DataFrame immutability + - Meeting of minimum quality thresholds + - Absence of residual preprocessing artifacts + """ + + def setUp(self): + print(f"\n{'='*60}") + print("TEST CLASS: TestMIMIC4NoteExtDIBHCDatasetEndToEnd") + print(f"{'='*60}") + self.pipeline = _DummyDataset() + self.df_input = _make_note_df(n=10) + print(f" Created {len(self.df_input)} synthetic notes") + + def test_preprocess_runs_without_error(self) -> None: + """Test that preprocess() completes without raising exceptions. + + Verifies that the full 7-step pipeline executes successfully + on synthetic data and produces output. + """ + print("\nTEST: test_preprocess_runs_without_error") + try: + result = self.pipeline.preprocess(self.df_input) + print(f" ✓ preprocess() completed; {len(result)} rows survived pipeline") + except Exception as e: + self.fail(f"preprocess() raised an unexpected exception: {e}") + + def test_preprocess_output_columns_present(self) -> None: + """Test that preprocess() creates all required output columns. + + Verifies that the output DataFrame contains summary, + hospital_course, and brief_hospital_course columns. + """ + print("\nTEST: test_preprocess_output_columns_present") + result = self.pipeline.preprocess(self.df_input) + for col in ("summary", "hospital_course", "brief_hospital_course"): + self.assertIn( + col, + result.columns, + msg=f"Missing column: {col}" + ) + print(f" ✓ Column '{col}' present") + + def test_preprocess_does_not_mutate_input(self) -> None: + """Test that preprocess() does not modify the input DataFrame. + + Verifies that the original input DataFrame is not modified + by the preprocessing pipeline. + """ + print("\nTEST: test_preprocess_does_not_mutate_input") + original_text = self.df_input["text"].iloc[0] + _ = self.pipeline.preprocess(self.df_input) + self.assertEqual(self.df_input["text"].iloc[0], original_text) + print(" ✓ Input DataFrame unchanged") + + def test_preprocess_summary_minimum_length(self) -> None: + """Test that all surviving summaries meet minimum length threshold. + + Verifies that preprocess() correctly enforces the min_chars + quality threshold on output summaries. + """ + print("\nTEST: test_preprocess_summary_minimum_length") + result = self.pipeline.preprocess(self.df_input) + if len(result) > 0: + min_len = result["summary"].str.len().min() + print( + f" Shortest surviving summary: {min_len} chars " + f"(threshold: {self.pipeline.min_chars})" + ) + self.assertGreaterEqual(min_len, self.pipeline.min_chars) + print(" ✓ passed") + + def test_preprocess_brief_hospital_course_minimum_length(self) -> None: + """Test that all brief hospital courses meet minimum length. + + Verifies that preprocess() correctly enforces the min_chars_bhc + threshold on output brief hospital courses. + """ + print("\nTEST: test_preprocess_brief_hospital_course_minimum_length") + result = self.pipeline.preprocess(self.df_input) + if len(result) > 0: + min_len = result["brief_hospital_course"].str.len().min() + print( + f" Shortest BHC: {min_len} chars " + f"(threshold: {self.pipeline.min_chars_bhc})" + ) + self.assertGreaterEqual(min_len, self.pipeline.min_chars_bhc) + print(" ✓ passed") + + def test_preprocess_no_residual_discharge_instructions_header(self) -> None: + """Test that discharge header is completely removed. + + Verifies that summaries do not retain the 'Discharge Instructions:' + header text after preprocessing. + """ + print("\nTEST: test_preprocess_no_residual_discharge_instructions_header") + result = self.pipeline.preprocess(self.df_input) + for summary in result["summary"]: + self.assertFalse( + summary.lower().startswith("discharge instructions"), + msg="Found residual 'Discharge Instructions:' header in summary", + ) + print(" ✓ No residual headers found") + + +# --------------------------------------------------------------------------- +# Lightweight stand-in for the dataset class +# --------------------------------------------------------------------------- + +class _DummyDataset: + """Lightweight stand-in for MIMIC4NoteExtDIBHCDataset. + + Exposes the full preprocessing pipeline without requiring BaseDataset + initialization or filesystem access. All instance methods delegate to + the real class using unbound-method calls, allowing self.* threshold + attributes to be properly honored. Useful for unit testing preprocessing + stages in isolation. + """ + + def __init__( + self, + min_chars: int = 350, + max_double_newlines: int = 5, + min_sentences: int = 3, + num_words_per_deidentified: int = 10, + min_chars_bhc: int = 500, + ): + self.min_chars = min_chars + self.max_double_newlines = max_double_newlines + self.min_sentences = min_sentences + self.num_words_per_deidentified = num_words_per_deidentified + self.min_chars_bhc = min_chars_bhc + + # --- delegate every instance method to the real class ---------------- + + def _step1_split_on_discharge_instructions(self, df: pd.DataFrame) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset instance method.""" + return MIMIC4NoteExtDIBHCDataset._step1_split_on_discharge_instructions( + self, df + ) + + def _step3_truncate_prefixes(self, df: pd.DataFrame) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset instance method.""" + return MIMIC4NoteExtDIBHCDataset._step3_truncate_prefixes(self, df) + + def _step5_truncate_suffixes(self, df: pd.DataFrame) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset instance method.""" + return MIMIC4NoteExtDIBHCDataset._step5_truncate_suffixes(self, df) + + def _step6_quality_filter(self, df: pd.DataFrame) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset instance method.""" + return MIMIC4NoteExtDIBHCDataset._step6_quality_filter(self, df) + + def _step7_filter_hospital_course(self, df: pd.DataFrame) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset instance method.""" + return MIMIC4NoteExtDIBHCDataset._step7_filter_hospital_course(self, df) + + # --- two helpers called by the delegated instance methods above ------ + + @staticmethod + def _remove_empty_and_short_summaries( + df: pd.DataFrame, min_length_summary: int = 350 + ) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset static method.""" + return MIMIC4NoteExtDIBHCDataset._remove_empty_and_short_summaries( + df, min_length_summary=min_length_summary + ) + + def _remove_regex_dict( + self, df: pd.DataFrame, regexes: dict, postprocess, keep: int = 0 + ) -> pd.DataFrame: + """Delegate to MIMIC4NoteExtDIBHCDataset instance method.""" + return MIMIC4NoteExtDIBHCDataset._remove_regex_dict( + df, regexes, postprocess, keep=keep + ) + + # --- full pipeline --------------------------------------------------- + + def preprocess(self, df: pd.DataFrame) -> pd.DataFrame: + """Execute the full preprocessing pipeline on a DataFrame. + + Chains all seven preprocessing steps to transform raw discharge notes + into clean, filtered summaries with extracted hospital course sections. + + Args: + df: Input DataFrame with note text to preprocess. + + Returns: + Processed DataFrame with additional columns: summary, hospital_course, + brief_hospital_course. Row count reduced by quality filters. + """ + df = df.copy() + df = MIMIC4NoteExtDIBHCDataset._step0_special_chars(df) + df = self._step1_split_on_discharge_instructions(df) + df = MIMIC4NoteExtDIBHCDataset._step2_encode_and_extract_hc(df) + df = self._step3_truncate_prefixes(df) + df = MIMIC4NoteExtDIBHCDataset._step4_remove_static_patterns(df) + df = self._step5_truncate_suffixes(df) + df = self._step6_quality_filter(df) + df = self._step7_filter_hospital_course(df) + return df + + +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + unittest.main() \ No newline at end of file