From 73ebc050635b093658471852ed6bb8965e4cafd7 Mon Sep 17 00:00:00 2001 From: John Wu Date: Sun, 14 Sep 2025 16:47:29 -0500 Subject: [PATCH 1/2] added task caching with parquet and pickle formats --- pyhealth/datasets/base_dataset.py | 115 ++++++++++--- tests/core/test_caching.py | 271 ++++++++++++++++++++++++++++++ 2 files changed, 359 insertions(+), 27 deletions(-) create mode 100644 tests/core/test_caching.py diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 2ac1e1d63..221fc336d 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -1,5 +1,6 @@ import logging import os +import pickle from abc import ABC from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path @@ -64,13 +65,14 @@ def scan_csv_gz_or_csv_tsv(path: str) -> pl.LazyFrame: Returns: pl.LazyFrame: The LazyFrame for the CSV.gz, CSV, TSV.gz, or TSV file. """ + def scan_file(file_path: str) -> pl.LazyFrame: - separator = '\t' if '.tsv' in file_path else ',' + separator = "\t" if ".tsv" in file_path else "," return pl.scan_csv(file_path, separator=separator, infer_schema=False) - + if path_exists(path): return scan_file(path) - + # Try the alternative extension if path.endswith(".csv.gz"): alt_path = path[:-3] # Remove .gz -> try .csv @@ -82,13 +84,14 @@ def scan_file(file_path: str) -> pl.LazyFrame: alt_path = f"{path}.gz" # Add .gz -> try .tsv.gz else: raise FileNotFoundError(f"Path does not have expected extension: {path}") - + if path_exists(alt_path): logger.info(f"Original path does not exist. Using alternative: {alt_path}") return scan_file(alt_path) - + raise FileNotFoundError(f"Neither path exists: {path} or {alt_path}") + class BaseDataset(ABC): """Abstract base class for all PyHealth datasets. @@ -352,7 +355,11 @@ def default_task(self) -> Optional[BaseTask]: return None def set_task( - self, task: Optional[BaseTask] = None, num_workers: int = 1 + self, + task: Optional[BaseTask] = None, + num_workers: int = 1, + cache_dir: Optional[str] = None, + cache_format: str = "parquet", ) -> SampleDataset: """Processes the base dataset to generate the task-specific sample dataset. @@ -361,6 +368,10 @@ def set_task( num_workers (int): Number of workers for multi-threading. Default is 1. This is because the task function is usually CPU-bound. And using multi-threading may not speed up the task function. + cache_dir (Optional[str]): Directory to cache processed samples. + Default is None (no caching). + cache_format (str): Format for caching ('parquet' or 'pickle'). + Default is 'parquet'. Returns: SampleDataset: The generated sample dataset. @@ -378,27 +389,77 @@ def set_task( filtered_global_event_df = task.pre_filter(self.collected_global_event_df) - logger.info(f"Generating samples with {num_workers} worker(s)...") - - samples = [] - - if num_workers == 1: - # single-threading (by default) - for patient in tqdm( - self.iter_patients(filtered_global_event_df), - total=filtered_global_event_df["patient_id"].n_unique(), - desc=f"Generating samples for {task.task_name} with 1 worker", - smoothing=0, - ): - samples.extend(task(patient)) - else: - # multi-threading (not recommended) - logger.info(f"Generating samples for {task.task_name} with {num_workers} workers") - patients = list(self.iter_patients(filtered_global_event_df)) - with ThreadPoolExecutor(max_workers=num_workers) as executor: - futures = [executor.submit(task, patient) for patient in patients] - for future in tqdm(as_completed(futures), total=len(futures), desc=f"Collecting samples for {task.task_name} from {num_workers} workers"): - samples.extend(future.result()) + # Check for cached data if cache_dir is provided + samples = None + if cache_dir is not None: + cache_path = Path(cache_dir) / f"{task.task_name}.{cache_format}" + if cache_path.exists(): + logger.info(f"Loading cached samples from {cache_path}") + try: + if cache_format == "parquet": + # Load samples from parquet file + cached_df = pl.read_parquet(cache_path) + samples = cached_df.to_dicts() + elif cache_format == "pickle": + # Load samples from pickle file + with open(cache_path, "rb") as f: + samples = pickle.load(f) + else: + raise ValueError(f"Unsupported cache format: {cache_format}") + logger.info(f"Loaded {len(samples)} cached samples") + except Exception as e: + logger.warning(f"Failed to load cached data: {e}. Regenerating...") + samples = None + + # Generate samples if not loaded from cache + if samples is None: + logger.info(f"Generating samples with {num_workers} worker(s)...") + + samples = [] + + if num_workers == 1: + # single-threading (by default) + for patient in tqdm( + self.iter_patients(filtered_global_event_df), + total=filtered_global_event_df["patient_id"].n_unique(), + desc=f"Generating samples for {task.task_name} with 1 worker", + smoothing=0, + ): + samples.extend(task(patient)) + else: + # multi-threading (not recommended) + logger.info( + f"Generating samples for {task.task_name} with {num_workers} workers" + ) + patients = list(self.iter_patients(filtered_global_event_df)) + with ThreadPoolExecutor(max_workers=num_workers) as executor: + futures = [executor.submit(task, patient) for patient in patients] + for future in tqdm( + as_completed(futures), + total=len(futures), + desc=f"Collecting samples for {task.task_name} from {num_workers} workers", + ): + samples.extend(future.result()) + + # Cache the samples if cache_dir is provided + if cache_dir is not None: + cache_path = Path(cache_dir) / f"{task.task_name}.{cache_format}" + cache_path.parent.mkdir(parents=True, exist_ok=True) + logger.info(f"Caching samples to {cache_path}") + try: + if cache_format == "parquet": + # Save samples as parquet file + samples_df = pl.DataFrame(samples) + samples_df.write_parquet(cache_path) + elif cache_format == "pickle": + # Save samples as pickle file + with open(cache_path, "wb") as f: + pickle.dump(samples, f) + else: + raise ValueError(f"Unsupported cache format: {cache_format}") + logger.info(f"Successfully cached {len(samples)} samples") + except Exception as e: + logger.warning(f"Failed to cache samples: {e}") sample_dataset = SampleDataset( samples, diff --git a/tests/core/test_caching.py b/tests/core/test_caching.py new file mode 100644 index 000000000..d6fc039fd --- /dev/null +++ b/tests/core/test_caching.py @@ -0,0 +1,271 @@ +import unittest +import tempfile +import os +from pathlib import Path +from unittest.mock import patch +import polars as pl + +from tests.base import BaseTestCase +from pyhealth.datasets.base_dataset import BaseDataset +from pyhealth.tasks.base_task import BaseTask +from pyhealth.datasets.sample_dataset import SampleDataset +from pyhealth.data import Patient + + +class MockTask(BaseTask): + """Mock task for testing purposes.""" + + def __init__(self, task_name="test_task"): + self.task_name = task_name + self.input_schema = {"test_attribute": "raw"} + self.output_schema = {"test_label": "binary"} + + def __call__(self, patient): + """Return mock samples based on patient data.""" + # Extract patient's test data from the patient's data source + patient_data = patient.data_source + + samples = [] + for row in patient_data.iter_rows(named=True): + sample = { + "test_attribute": row["test/test_attribute"], + "test_label": row["test/test_label"], + "patient_id": row["patient_id"], + } + samples.append(sample) + + return samples + + +class MockDataset(BaseDataset): + """Mock dataset for testing purposes.""" + + def __init__(self): + # Initialize without calling parent __init__ to avoid file dependencies + self.dataset_name = "TestDataset" + self.dev = False + + # Create realistic test data with patient_id, test_attribute, and test_label + self._collected_global_event_df = pl.DataFrame( + { + "patient_id": ["1", "2", "1", "2"], + "event_type": ["test", "test", "test", "test"], + "timestamp": [None, None, None, None], + "test/test_attribute": [ + "pat_1_attr_1", + "pat_2_attr_1", + "pat_1_attr_2", + "pat_2_attr_2", + ], + "test/test_label": [0, 1, 1, 0], + } + ) + self._unique_patient_ids = ["1", "2"] + + @property + def collected_global_event_df(self): + return self._collected_global_event_df + + @property + def unique_patient_ids(self): + return self._unique_patient_ids + + def iter_patients(self, df=None): + """Mock patient iterator that returns real Patient objects.""" + if df is None: + df = self.collected_global_event_df + + grouped = df.group_by("patient_id") + for patient_id, patient_df in grouped: + patient_id = patient_id[0] + yield Patient(patient_id=patient_id, data_source=patient_df) + + +class TestCachingFunctionality(BaseTestCase): + """Test cases for caching functionality in BaseDataset.set_task().""" + + def setUp(self): + """Set up test fixtures.""" + self.dataset = MockDataset() + self.task = MockTask() + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up test fixtures.""" + # Clean up temporary directory + import shutil + + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_set_task_signature(self): + """Test that set_task has the correct method signature.""" + import inspect + + sig = inspect.signature(BaseDataset.set_task) + params = list(sig.parameters.keys()) + + expected_params = ["self", "task", "num_workers", "cache_dir", "cache_format"] + self.assertEqual(params, expected_params) + + # Check default values + self.assertEqual(sig.parameters["task"].default, None) + self.assertEqual(sig.parameters["num_workers"].default, 1) + self.assertEqual(sig.parameters["cache_dir"].default, None) + self.assertEqual(sig.parameters["cache_format"].default, "parquet") + + def test_set_task_no_caching(self): + """Test set_task without caching (cache_dir=None).""" + sample_dataset = self.dataset.set_task(self.task) + + self.assertIsInstance(sample_dataset, SampleDataset) + self.assertEqual(len(sample_dataset), 4) # Two patients, two samples each + self.assertEqual(sample_dataset.dataset_name, "TestDataset") + + # Check that samples have the correct structure + sample = sample_dataset[0] + self.assertIn("test_attribute", sample) + self.assertIn("test_label", sample) + self.assertIn("patient_id", sample) + + def test_full_parquet_caching_cycle(self): + """Test complete save and load cycle with parquet caching.""" + cache_path = Path(self.temp_dir) / f"{self.task.task_name}.parquet" + + # Step 1: First call - should generate samples and save to cache + self.assertFalse(cache_path.exists(), "Cache file should not exist initially") + + sample_dataset_1 = self.dataset.set_task( + self.task, cache_dir=self.temp_dir, cache_format="parquet" + ) + + # Verify cache file was created + self.assertTrue( + cache_path.exists(), "Cache file should be created after first call" + ) + + # Verify the sample dataset is correct + self.assertIsInstance(sample_dataset_1, SampleDataset) + self.assertEqual( + len(sample_dataset_1), 4 + ) # Should have 4 samples from our mock data + + # Step 2: Second call - should load from cache (not regenerate) + sample_dataset_2 = self.dataset.set_task( + self.task, cache_dir=self.temp_dir, cache_format="parquet" + ) + + # Verify the loaded dataset matches the original + self.assertIsInstance(sample_dataset_2, SampleDataset) + self.assertEqual(len(sample_dataset_2), 4) + + # Step 3: Verify the actual cached data is correct + # Load the parquet file directly to check its contents + cached_df = pl.read_parquet(cache_path) + cached_samples = cached_df.to_dicts() + + self.assertEqual(len(cached_samples), 4) + + # Verify sample content matches expected structure + for sample in cached_samples: + self.assertIn("test_attribute", sample) + self.assertIn("test_label", sample) + self.assertIn("patient_id", sample) + self.assertIn(sample["patient_id"], ["1", "2"]) + self.assertIn(sample["test_label"], [0, 1]) + + def test_full_pickle_caching_cycle(self): + """Test complete save and load cycle with pickle caching.""" + cache_path = Path(self.temp_dir) / f"{self.task.task_name}.pickle" + + # Step 1: First call - should generate samples and save to cache + self.assertFalse(cache_path.exists(), "Cache file should not exist initially") + + sample_dataset_1 = self.dataset.set_task( + self.task, cache_dir=self.temp_dir, cache_format="pickle" + ) + + # Verify cache file was created + self.assertTrue( + cache_path.exists(), "Cache file should be created after first call" + ) + + # Verify the sample dataset is correct + self.assertIsInstance(sample_dataset_1, SampleDataset) + self.assertEqual( + len(sample_dataset_1), 4 + ) # Should have 4 samples from our mock data + + # Step 2: Second call - should load from cache (not regenerate) + sample_dataset_2 = self.dataset.set_task( + self.task, cache_dir=self.temp_dir, cache_format="pickle" + ) + + # Verify the loaded dataset matches the original + self.assertIsInstance(sample_dataset_2, SampleDataset) + self.assertEqual(len(sample_dataset_2), 4) + + # Step 3: Verify the actual cached data is correct + # Load the pickle file directly to check its contents + import pickle + + with open(cache_path, "rb") as f: + cached_samples = pickle.load(f) + + self.assertEqual(len(cached_samples), 4) + + # Verify sample content matches expected structure + for sample in cached_samples: + self.assertIn("test_attribute", sample) + self.assertIn("test_label", sample) + self.assertIn("patient_id", sample) + self.assertIn(sample["patient_id"], ["1", "2"]) + self.assertIn(sample["test_label"], [0, 1]) + + def test_set_task_invalid_cache_format(self): + """Test set_task with invalid cache format.""" + # This should not raise an error during set_task call, + # but should log a warning when trying to save + sample_dataset = self.dataset.set_task( + self.task, cache_dir=self.temp_dir, cache_format="invalid_format" + ) + + self.assertIsInstance(sample_dataset, SampleDataset) + self.assertEqual(len(sample_dataset), 4) # Generated samples + + @patch("polars.read_parquet") + def test_set_task_cache_load_failure_fallback(self, mock_read_parquet): + """Test fallback to generation when cache loading fails.""" + # Make read_parquet raise an exception + mock_read_parquet.side_effect = Exception("Failed to read cache") + + # Create a dummy cache file + cache_path = Path(self.temp_dir) / f"{self.task.task_name}.parquet" + cache_path.touch() + + sample_dataset = self.dataset.set_task( + self.task, cache_dir=self.temp_dir, cache_format="parquet" + ) + + # Should still work by falling back to generation + self.assertIsInstance(sample_dataset, SampleDataset) + self.assertEqual(len(sample_dataset), 4) # Generated samples + + def test_cache_directory_creation(self): + """Test that cache directory is created if it doesn't exist.""" + nested_cache_dir = os.path.join(self.temp_dir, "nested", "cache", "dir") + + # Ensure the nested directory doesn't exist + self.assertFalse(os.path.exists(nested_cache_dir)) + + with patch("polars.DataFrame.write_parquet"): + sample_dataset = self.dataset.set_task( + self.task, cache_dir=nested_cache_dir, cache_format="parquet" + ) + + # Directory should be created + self.assertTrue(os.path.exists(nested_cache_dir)) + self.assertIsInstance(sample_dataset, SampleDataset) + + +if __name__ == "__main__": + unittest.main() From 2f4e11b411abcb27dfa69c2b03fe0e0cefb93db2 Mon Sep 17 00:00:00 2001 From: John Wu Date: Sun, 14 Sep 2025 19:42:29 -0500 Subject: [PATCH 2/2] added notebook with more realistic testing --- .../mimic3_mortality_prediction_cached.ipynb | 291 ++++++++++++++++++ 1 file changed, 291 insertions(+) create mode 100644 examples/mimic3_mortality_prediction_cached.ipynb diff --git a/examples/mimic3_mortality_prediction_cached.ipynb b/examples/mimic3_mortality_prediction_cached.ipynb new file mode 100644 index 000000000..8f908bf56 --- /dev/null +++ b/examples/mimic3_mortality_prediction_cached.ipynb @@ -0,0 +1,291 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "cc001dde", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Adding PyHealth to sys.path: /home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth\n" + ] + } + ], + "source": [ + "import os\n", + "import sys\n", + "\n", + "pyhealth_path = os.path.dirname(os.getcwd())\n", + "if pyhealth_path not in sys.path:\n", + " print(f\"Adding PyHealth to sys.path: {pyhealth_path}\")\n", + " sys.path.insert(0, pyhealth_path)" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "9887d8db", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth'" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pyhealth_path" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "3e05a68b", + "metadata": {}, + "outputs": [], + "source": [ + "from pyhealth.tasks import BaseTask\n", + "from typing import Any, Dict, List, Optional\n", + "from datetime import datetime\n", + "\n", + "class MortalityPredictionMIMIC3Heterogeneous(BaseTask):\n", + " \"\"\"Task for predicting mortality using MIMIC-III dataset with text data.\n", + "\n", + " This task aims to predict whether the patient will decease in the next hospital\n", + " visit based on clinical information from the current visit.\n", + " \"\"\"\n", + "\n", + " task_name: str = \"MortalityPredictionMIMIC3\"\n", + " input_schema: Dict[str, str] = {\n", + " \"conditions\": \"sequence\",\n", + " \"procedures\": \"sequence\",\n", + " \"drugs\": \"sequence\",\n", + " }\n", + " output_schema: Dict[str, str] = {\"mortality\": \"binary\"}\n", + "\n", + " def __call__(self, patient: Any) -> List[Dict[str, Any]]:\n", + " \"\"\"Processes a single patient for the mortality prediction task.\"\"\"\n", + " samples = []\n", + "\n", + " # We will drop the last visit\n", + " visits = patient.get_events(event_type=\"admissions\")\n", + "\n", + " if len(visits) <= 1:\n", + " return []\n", + "\n", + " for i in range(len(visits) - 1):\n", + " visit = visits[i]\n", + " next_visit = visits[i + 1]\n", + "\n", + " # Check discharge status for mortality label - more robust handling\n", + " if next_visit.hospital_expire_flag not in [0, 1, \"0\", \"1\"]:\n", + " mortality_label = 0\n", + " else:\n", + " mortality_label = int(next_visit.hospital_expire_flag)\n", + "\n", + " # Convert string timestamps to datetime objects\n", + " try:\n", + " # Check the type and convert if necessary\n", + " if isinstance(visit.dischtime, str):\n", + " discharge_time = datetime.strptime(\n", + " visit.dischtime, \"%Y-%m-%d %H:%M:%S\"\n", + " )\n", + " else:\n", + " discharge_time = visit.dischtime\n", + " except (ValueError, AttributeError):\n", + " # If conversion fails, skip this visit\n", + " print(\"Error parsing discharge time:\", visit.dischtime)\n", + " continue\n", + "\n", + " # Get clinical codes\n", + " diagnoses = patient.get_events(\n", + " event_type=\"diagnoses_icd\",\n", + " start=visit.timestamp,\n", + " end=discharge_time, # Now using a datetime object\n", + " )\n", + " procedures = patient.get_events(\n", + " event_type=\"procedures_icd\",\n", + " start=visit.timestamp,\n", + " end=discharge_time, # Now using a datetime object\n", + " )\n", + " prescriptions = patient.get_events(\n", + " event_type=\"prescriptions\",\n", + " start=visit.timestamp,\n", + " end=discharge_time, # Now using a datetime object\n", + " )\n", + "\n", + " conditions = [event.icd9_code for event in diagnoses]\n", + " procedures_list = [event.icd9_code for event in procedures]\n", + " drugs = [event.drug for event in prescriptions]\n", + "\n", + " # Exclude visits without condition, procedure, or drug code\n", + " samples.append(\n", + " {\n", + " \"hadm_id\": visit.hadm_id,\n", + " \"patient_id\": patient.patient_id,\n", + " \"conditions\": conditions,\n", + " \"procedures\": procedures_list,\n", + " \"drugs\": drugs,\n", + " \"mortality\": mortality_label,\n", + " }\n", + " )\n", + "\n", + " return samples\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "da8360f8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "No config path provided, using default config\n", + "Initializing mimic3 dataset from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III (dev mode: False)\n", + "Scanning table: patients from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/johnwu3/projects/PyHealth_Branch_Testing/PyHealth/pyhealth/datasets/mimic3.py:50: UserWarning: Events from prescriptions table only have date timestamp (no specific time). This may affect temporal ordering of events.\n", + " warnings.warn(\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Some column names were converted to lowercase\n", + "Scanning table: admissions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n", + "Some column names were converted to lowercase\n", + "Scanning table: icustays from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv\n", + "Some column names were converted to lowercase\n", + "Scanning table: diagnoses_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv\n", + "Some column names were converted to lowercase\n", + "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n", + "Scanning table: procedures_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv\n", + "Some column names were converted to lowercase\n", + "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n", + "Scanning table: prescriptions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv\n", + "Some column names were converted to lowercase\n", + "Scanning table: noteevents from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv.gz\n", + "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/NOTEEVENTS.csv\n", + "Some column names were converted to lowercase\n", + "Preprocessing table: noteevents with preprocess_noteevents\n" + ] + } + ], + "source": [ + "from pyhealth.datasets import MIMIC3Dataset\n", + "dataset = MIMIC3Dataset(\n", + " root=\"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\",\n", + " tables=[\"diagnoses_icd\", \"procedures_icd\", \"prescriptions\", \"noteevents\"]\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "dd5c282f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Setting task MortalityPredictionMIMIC3 for mimic3 base dataset...\n", + "Loading cached samples from cache/MortalityPredictionMIMIC3.parquet\n", + "Loaded 2776 cached samples\n", + "Label mortality vocab: {0: 0, 1: 1}\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Processing samples: 100%|██████████| 2776/2776 [00:00<00:00, 63155.03it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generated 2776 samples for task MortalityPredictionMIMIC3\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "from pyhealth.tasks.mortality_prediction import MortalityPredictionMIMIC3\n", + "from pyhealth.datasets import split_by_patient, get_dataloader\n", + "mimic3_mortality_prediction = MortalityPredictionMIMIC3Heterogeneous()\n", + "samples = dataset.set_task(mimic3_mortality_prediction, num_workers=1, cache_dir=\"cache/\") # use default task" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "06b5a9a3", + "metadata": {}, + "outputs": [], + "source": [ + "from pyhealth.datasets import split_by_sample\n", + "\n", + "\n", + "train_dataset, val_dataset, test_dataset = split_by_sample(\n", + " dataset=samples,\n", + " ratios=[0.7, 0.1, 0.2]\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "medical_coding_demo", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.16" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}