move collate_fn_dict to datasets

README.rst

@@ -133,7 +133,7 @@ An ML Pipeline Example
  from pyhealth.tasks import drug_recommendation_mimic3_fn
  from pyhealth.datasets.splitter import split_by_patient
  from torch.utils.data import DataLoader
- from pyhealth.utils import collate_fn_dict
+ from pyhealth.datasets.utils import collate_fn_dict
 
  mimic3dataset.set_task(task_fn=drug_recommendation_mimic3_fn) # use default drugrec task
  train_ds, val_ds, test_ds = split_by_patient(mimic3dataset, [0.8, 0.1, 0.1])

docs/index.rst

@@ -109,7 +109,7 @@ An ML Pipeline Example
  from pyhealth.tasks import drug_recommendation_mimic3_fn
  from pyhealth.datasets.splitter import split_by_patient
  from torch.utils.data import DataLoader
- from pyhealth.utils import collate_fn_dict
+ from pyhealth.datasets.utils import collate_fn_dict
 
  mimic3dataset.set_task(task_fn=drug_recommendation_mimic3_fn) # use default drugrec task
  train_ds, val_ds, test_ds = split_by_patient(mimic3dataset, [0.8, 0.1, 0.1])

leaderboard/leaderboard_gen.py

@@ -5,7 +5,7 @@
 from pyhealth.models import *
 from pyhealth.datasets.splitter import split_by_patient
 from pyhealth.tasks import *
-from pyhealth.utils import collate_fn_dict
+from pyhealth.datasets.utils import collate_fn_dict
 from pyhealth.trainer import Trainer
 from pyhealth.evaluator import evaluate
 from pyhealth.metrics import *

pyhealth/datasets/__init__.py

@@ -4,3 +4,4 @@
 from .omop import OMOPDataset
 from .eicu import eICUDataset
 from .splitter import split_by_patient, split_by_visit
+from .utils import collate_fn_dict, get_dataloader

pyhealth/datasets/utils.py

@@ -1,8 +1,10 @@
+import hashlib
 import os
 from datetime import datetime
 from typing import Optional
-import hashlib
+
 from dateutil.parser import parse as dateutil_parse
+from torch.utils.data import DataLoader
 
 from pyhealth import BASE_CACHE_PATH
 from pyhealth.utils import create_directory
@@ -28,3 +30,14 @@ def strptime(s: str) -> Optional[datetime]:
  if s != s:
  return None
  return dateutil_parse(s)
+
+
+def collate_fn_dict(batch):
+ return {key: [d[key] for d in batch] for key in batch[0]}
+
+
+def get_dataloader(dataset, batch_size, shuffle=False):
+ dataloader = DataLoader(
+ dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn_dict
+ )
+ return dataloader

pyhealth/models/test/test_gamenet.py

@@ -3,7 +3,7 @@
 from pyhealth.datasets import MIMIC3Dataset
 from pyhealth.models import GAMENet
 from pyhealth.tasks import drug_recommendation_mimic3_fn
-from pyhealth.utils import collate_fn_dict
+from pyhealth.datasets.utils import collate_fn_dict
 
 dataset = MIMIC3Dataset(
  root="/srv/local/data/physionet.org/files/mimiciii/1.4",

pyhealth/models/test/test_general_model.py

@@ -1,7 +1,7 @@
 from torch.utils.data import DataLoader
 
 from pyhealth.datasets import MIMIC3Dataset
-from pyhealth.utils import collate_fn_dict
+from pyhealth.datasets.utils import collate_fn_dict
 
 # from pyhealth.models import CNN as Model
 # from pyhealth.models import RNN as Model

pyhealth/models/test/test_micron.py

@@ -3,7 +3,7 @@
 from pyhealth.datasets import MIMIC3Dataset
 from pyhealth.models import MICRON
 from pyhealth.tasks import drug_recommendation_mimic3_fn
-from pyhealth.utils import collate_fn_dict
+from pyhealth.datasets.utils import collate_fn_dict
 
 dataset = MIMIC3Dataset(
  root="/srv/local/data/physionet.org/files/mimiciii/1.4",

pyhealth/models/test/test_safedrug.py

@@ -3,7 +3,7 @@
 from pyhealth.datasets import MIMIC3Dataset
 from pyhealth.models import SafeDrug
 from pyhealth.tasks import drug_recommendation_mimic3_fn
-from pyhealth.utils import collate_fn_dict
+from pyhealth.datasets.utils import collate_fn_dict
 
 dataset = MIMIC3Dataset(
  root="/srv/local/data/physionet.org/files/mimiciii/1.4",

pyhealth/utils.py

@@ -6,10 +6,6 @@
 import torch
 
 
-def collate_fn_dict(batch):
- return {key: [d[key] for d in batch] for key in batch[0]}
-
-
 def set_seed(seed):
  random.seed(seed)
  np.random.seed(seed)