feat(tasks): generalized note section extraction by note type

examples/multimodal_mimic4_task_tutorial.ipynb

@@ -74,10 +74,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "SEED = 42\n",
-    "random.seed(SEED)\n",
-    "np.random.seed(SEED)\n",
-    "torch.manual_seed(SEED)"
+    "SEED = 1"
    ]
   },
   {
@@ -108,7 +105,7 @@
     "        print(f\"  admittime:     {adm.timestamp}\")\n",
     "        print(f\"  dischtime:     {adm.dischtime}\")\n",
     "\n",
-    "def print_patient_note_info(sample, note_type=\"discharge\", char_limit=80):\n",
+    "def print_patient_note_info(sample, note_type=\"discharge\", char_limit=500):\n",
     "    patient = dataset.get_patient(sample['patient_id'])\n",
     "    notes = patient.get_events(\n",
     "        event_type=note_type,\n",
@@ -144,7 +141,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "PYHEALTH_REPO_ROOT = '/Users/wpang/Desktop/PyHealth'\n",
+    "PYHEALTH_REPO_ROOT = '/Users/williampang/Desktop/PyHealth'\n",
     "\n",
     "EHR_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n",
     "NOTE_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n",
@@ -204,6 +201,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "np.random.seed(SEED)\n",
     "random_patient_record = np.random.randint(0, len(samples)-1)\n",
     "sample = samples[random_patient_record]\n",
     "sample_patient_id = sample['patient_id']\n",
@@ -226,7 +224,7 @@
    "id": "fd7d6767-ae51-4278-bb55-d12bdf3a4ff2",
    "metadata": {},
    "source": [
-    "### 3.2 Random Patient Window"
+    "### 3.2 Patient with Filtered Time Window"
    ]
   },
   {
@@ -265,12 +263,47 @@
     "print_patient_note_info(sample, note_type=\"radiology\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "e1dd556a",
+   "metadata": {},
+   "source": [
+    "### 4. Inspect Task Output (before processor)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9113eab7-20f5-455b-8e99-570de8622ba6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "task = ClinicalNotesICDLabsMIMIC4()\n",
+    "samples = dataset.set_task(task, num_workers=1)\n",
+    "sample = samples[random_patient_record]\n",
+    "sample_patient_id = sample['patient_id']\n",
+    "\n",
+    "patient = dataset.get_patient(sample_patient_id)\n",
+    "sample_task_output = task(patient)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23399090-175d-4492-8657-4a2c51293992",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sample_task_output"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "a447c4e2-cfae-4e15-a2ac-0fac5754c1f2",
    "metadata": {},
    "source": [
-    "## 4. Cleanup"
+    "## 5. Cleanup"
    ]
   },
   {
@@ -294,9 +327,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "pyhealth",
    "language": "python",
-   "name": "python3"
+   "name": "pyhealth"
   },
   "language_info": {
    "codemirror_mode": {
@@ -307,7 +340,8 @@
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3"
+   "pygments_lexer": "ipython3",
+   "version": "3.13.3"
   }
  },
  "nbformat": 4,

pyhealth/tasks/multimodal_mimic4.py

@@ -69,99 +69,61 @@ class BaseMultimodalMIMIC4Task(BaseTask):
         item for itemids in VITAL_CATEGORIES.values() for item in itemids
     ]
 
+    RADIOLOGY_CLINICAL_HEADERS: ClassVar[List[str]] = [
+        "indication",
+        "impression",
+        # "findings",
+        # "clinical history",
+        # "history",
+        # "comparison",
+        # "technique",
+        # "conclusion",
+        # "summary"
+    ]
+
+    DISCHARGE_CLINICAL_HEADERS: ClassVar[List[str]] = [
+        "chief complaint",
+        # "history of present illness",
+        # "hpi",
+        # "past medical history",
+        # "past medical and surgical history",
+        # "past medical/surgical history",
+        # "past surgical history",
+        # "medications on admission",
+        # "admission medications",
+        # "home medications",
+        # "social history",
+        # "family history",
+        # "allergies",
+        # "review of systems",
+    ]
+
     def __init__(
         self,
         window_hours: Optional[float] = None,
     ):
         self.window_hours = window_hours
 
-    _ADMISSION_SECTION_TARGETS: ClassVar[frozenset] = frozenset({
-        "chief complaint",
-        "history of present illness",
-        "hpi",
-        "past medical history",
-        "past medical and surgical history",
-        "past medical/surgical history",
-        "past surgical history",
-        "medications on admission",
-        "admission medications",
-        "home medications",
-    })
-
-    # Matches a line that is purely a section header: words + colon + nothing else.
-    # E.g. "Chief Complaint:", "Past Medical History:", "HPI:".
-    _SECTION_HEADER_RE: ClassVar[re.Pattern] = re.compile(
-        r"^\s*([A-Za-z][A-Za-z\s,/\-\.]{1,60}?)\s*:\s*$"
-    )
-
     @staticmethod
     def _clean_text(text: Optional[str]) -> Optional[str]:
         """Return text if non-empty, otherwise None."""
         return text if text else None
 
     @staticmethod
-    def _extract_admission_sections(text: str) -> str:
-        """Extract admission-context sections from a MIMIC-IV discharge note.
-
-        Parses Chief Complaint, HPI, Past Medical History, and Medications on
-        Admission sections — information clinically available at admission time.
-        Falls back to the first 1 024 characters if no target sections are found.
-        """
-        sections: Dict[str, str] = {}
-        current_key: Optional[str] = None
-        current_lines: List[str] = []
-
-        for line in text.split("\n"):
-            m = BaseMultimodalMIMIC4Task._SECTION_HEADER_RE.match(line)
-            if m:
-                if current_key is not None:
-                    sections[current_key] = "\n".join(current_lines).strip()
-                current_key = m.group(1).strip().lower()
-                current_lines = []
-            elif current_key is not None:
-                current_lines.append(line)
-
-        if current_key is not None:
-            sections[current_key] = "\n".join(current_lines).strip()
-
-        targets = BaseMultimodalMIMIC4Task._ADMISSION_SECTION_TARGETS
-        extracted = [v for k, v in sections.items() if k in targets and v]
-        return " [SEP] ".join(extracted) if extracted else text[:1024]
-
-    def _collect_admission_note_sections(
-        self,
-        patient: Any,
-        hadm_id: Any,
-        admission_time: datetime,
-    ) -> Tuple[List[str], List[float]]:
-        """Collect admission-context text from the discharge note.
-
-        No time filter is applied because the target sections (CC, HPI, PMH,
-        Medications on Admission) describe the patient's state *at admission*
-        regardless of when the note was finalised. Returned timestamps are 0.0
-        so downstream models treat the text as admission-time context.
-        """
-        notes = patient.get_events(
-            event_type="discharge",
-            filters=[("hadm_id", "==", hadm_id)],
-        )
-
-        texts: List[str] = []
-        for note in notes:
-            try:
-                raw = note.text
-                if not raw:
-                    continue
-                extracted = self._extract_admission_sections(raw)
-                if extracted:
-                    texts.append(extracted)
-            except AttributeError:
-                pass
-
-        if not texts:
-            return [self.MISSING_TEXT_TOKEN], [self.MISSING_FLOAT_TOKEN]
-
-        return texts, [0.0] * len(texts)
+    def _parse_note_sections(text: str, note_type: str) -> Dict[str, str]:
+        """Split a note into {lowercased_header: content_text} pairs."""
+        ext_text = text + '\n\n'
+        if note_type == "radiology":
+            section_re = re.compile(r'([a-zA-Z ]+):[ \t\n]+(.+?)\n{2,}', re.DOTALL)
+        elif note_type == "discharge":
+            section_re = re.compile(r'([a-zA-Z ]+):\n+(.+?)\n{2,}', re.DOTALL)
+        else:
+            raise ValueError(f"Note Type '{note_type}' not supported.")
+        return {
+            m.group(1).strip().lower(): m.group(2).strip()
+            for m in section_re.finditer(ext_text)
+            if m.end() - m.start() > 0
+        }
 
     @staticmethod
     def _parse_datetime(value: Any) -> Optional[datetime]:
@@ -424,6 +386,8 @@ def _collect_notes(
         admission_time: datetime,
         start_time: Optional[datetime] = None,
         end_time: Optional[datetime] = None,
+        section_headers: Optional[List[str]] = None,
+        fallback_to_full_note: bool = False,
     ) -> Tuple[List[str], List[float]]:
         """Collect notes of a given type for one admission.
 
@@ -432,8 +396,13 @@ def _collect_notes(
             note_event_type: Event type string (e.g. "discharge", "radiology").
             hadm_id: Admission ID to filter by.
             admission_time: Admission start time; used to compute time offsets.
-            start_time: Optional start of the time window
-            end_time: Optional end of the time window
+            start_time: Optional start of the time window.
+            end_time: Optional end of the time window.
+            section_headers: When provided, extract only these named sections
+                from each note (lowercased match against parsed headers).
+            fallback_to_full_note: When True (default), falls back to the full
+                note text if no matching sections are found. When False, notes
+                with no matching sections are dropped entirely.
 
         Returns:
             Tuple of (texts, hours_from_admission). Falls back to
@@ -453,6 +422,14 @@ def _collect_notes(
             try:
                 note_text = self._clean_text(note.text)
                 if note_text:
+                    if section_headers is not None:
+                        parsed = self._parse_note_sections(note_text, note_type=note_event_type)
+                        extracted = [f"{k}: {v}" for k, v in parsed.items() if k in section_headers and v]
+                        if extracted:
+                            note_text = " [SEP] ".join(extracted)
+                        elif not fallback_to_full_note:
+                            continue
+
                     time_from_admission = self._to_hours(
                         (note.timestamp - admission_time).total_seconds()
                     )
@@ -463,12 +440,6 @@ def _collect_notes(
             ):  # note object is missing .text or .timestamp attribute (e.g. malformed note)
                 pass
 
-        if (
-            not notes or not texts
-        ):  # If we get an empty list or all notes were malformed
-            return [self.MISSING_TEXT_TOKEN], [
-                self.MISSING_FLOAT_TOKEN
-            ]  # Token representing missing text/time
         return texts, note_times
 
 
@@ -549,6 +520,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
                 admission_time,
                 start_time=effective_start,
                 end_time=effective_end,
+                section_headers=self.DISCHARGE_CLINICAL_HEADERS,
             )
             all_discharge_texts.extend(discharge_texts)
             all_discharge_times_from_admission.extend(discharge_times)
@@ -560,10 +532,18 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
                 admission_time,
                 start_time=effective_start,
                 end_time=effective_end,
+                section_headers=self.RADIOLOGY_CLINICAL_HEADERS,
             )
             all_radiology_texts.extend(radiology_texts)
             all_radiology_times_from_admission.extend(radiology_times)
 
+        if not all_discharge_texts:
+            all_discharge_texts = [self.MISSING_TEXT_TOKEN]
+            all_discharge_times_from_admission = [self.MISSING_FLOAT_TOKEN]
+        if not all_radiology_texts:
+            all_radiology_texts = [self.MISSING_TEXT_TOKEN]
+            all_radiology_times_from_admission = [self.MISSING_FLOAT_TOKEN]
+
         discharge_note_times_from_admission = (
             all_discharge_texts,
             all_discharge_times_from_admission,
@@ -697,6 +677,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
                 admission_time,
                 start_time=effective_start,
                 end_time=effective_end,
+                section_headers=self.DISCHARGE_CLINICAL_HEADERS,
             )
             all_discharge_texts.extend(discharge_texts)
             all_discharge_times_from_admission.extend(discharge_times)
@@ -708,6 +689,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
                 admission_time,
                 start_time=effective_start,
                 end_time=effective_end,
+                section_headers=self.RADIOLOGY_CLINICAL_HEADERS,
             )
             all_radiology_texts.extend(radiology_texts)
             all_radiology_times_from_admission.extend(radiology_times)
@@ -746,6 +728,13 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
             all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
             all_lab_times.append(self.MISSING_FLOAT_TOKEN)
 
+        if not all_discharge_texts:
+            all_discharge_texts = [self.MISSING_TEXT_TOKEN]
+            all_discharge_times_from_admission = [self.MISSING_FLOAT_TOKEN]
+        if not all_radiology_texts:
+            all_radiology_texts = [self.MISSING_TEXT_TOKEN]
+            all_radiology_times_from_admission = [self.MISSING_FLOAT_TOKEN]
+
         discharge_note_times_from_admission = (
             all_discharge_texts,
             all_discharge_times_from_admission,
@@ -992,9 +981,12 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
             if admission_dischtime < admission_time:
                 admission_dischtime = admission_time
 
-            # Admission-context sections — no time window applied to notes
-            note_texts, note_times = self._collect_admission_note_sections(
-                patient, admission.hadm_id, admission_time
+            note_texts, note_times = self._collect_notes(
+                patient,
+                "discharge",
+                admission.hadm_id,
+                admission_time,
+                section_headers=self.DISCHARGE_CLINICAL_HEADERS,
             )
             all_note_texts.extend(note_texts)
             all_note_times.extend(note_times)
@@ -1056,6 +1048,10 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]:
             all_vital_masks.append([False] * len(self.VITAL_CATEGORY_NAMES))
             all_vital_times.append(self.MISSING_FLOAT_TOKEN)
 
+        if not all_note_texts:
+            all_note_texts = [self.MISSING_TEXT_TOKEN]
+            all_note_times = [self.MISSING_FLOAT_TOKEN]
+
         record: Dict[str, Any] = {
             "patient_id": patient.patient_id,
             "admission_note_times": (all_note_texts, all_note_times),