data/load_data.py

from datasets import load_dataset

def load_discharge_notes():
    notes = load_dataset(
        "ntphuc149/MIMIC-III-Clinical-Database",
        "NOTEEVENTS",
        split="train"
    )

    # Filter only discharge summaries
    notes = notes.filter(
        lambda x: x["CATEGORY"] == "Discharge summary"
    )

    return notes