diff --git a/pyproject.toml b/pyproject.toml index f89566ca..a933c89a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,14 +35,13 @@ dependencies = [ "Jinja2==3.1.3", "meds==0.3.3", "meds_reader==0.1.9", - "networkx==3.2.1", + "networkx>=3.2.1", "numpy==1.24.3", "packaging==23.2", "pandas==2.2.0", "peft>=0.10.0", "Pillow==10.3.0", - "pyarrow==15.0.0", - "pydantic==2.6.0", + "pyarrow>=15.0.0", "python-dateutil==2.8.2", "PyYAML==6.0.1", "scikit-learn==1.4.0", diff --git a/sample_data/MIMIC-IV-meds/meds_reader/.done b/sample_data/MIMIC-IV-meds/meds_reader/.done new file mode 100644 index 00000000..e69de29b diff --git a/sample_data/MIMIC-IV-meds/meds_reader/code/data b/sample_data/MIMIC-IV-meds/meds_reader/code/data new file mode 100644 index 00000000..83167a2d Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/code/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/code/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/code/dictionary new file mode 100644 index 00000000..80e50f1b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/code/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/code/zdict b/sample_data/MIMIC-IV-meds/meds_reader/code/zdict new file mode 100644 index 00000000..d1d5680f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/code/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/doses_per_24_hrs/data b/sample_data/MIMIC-IV-meds/meds_reader/doses_per_24_hrs/data new file mode 100644 index 00000000..37a21e9c Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/doses_per_24_hrs/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/doses_per_24_hrs/zdict b/sample_data/MIMIC-IV-meds/meds_reader/doses_per_24_hrs/zdict new file mode 100644 index 00000000..2f175f0e Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/doses_per_24_hrs/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/drg_mortality/data b/sample_data/MIMIC-IV-meds/meds_reader/drg_mortality/data new file mode 100644 index 00000000..822ebcaa Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/drg_mortality/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/drg_mortality/zdict b/sample_data/MIMIC-IV-meds/meds_reader/drg_mortality/zdict new file mode 100644 index 00000000..102e0bbd Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/drg_mortality/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/drg_severity/data b/sample_data/MIMIC-IV-meds/meds_reader/drg_severity/data new file mode 100644 index 00000000..84ccd678 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/drg_severity/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/drg_severity/zdict b/sample_data/MIMIC-IV-meds/meds_reader/drg_severity/zdict new file mode 100644 index 00000000..6947ee82 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/drg_severity/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/emar_id/data b/sample_data/MIMIC-IV-meds/meds_reader/emar_id/data new file mode 100644 index 00000000..9787a8a9 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/emar_id/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/emar_id/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/emar_id/dictionary new file mode 100644 index 00000000..e69de29b diff --git a/sample_data/MIMIC-IV-meds/meds_reader/emar_id/zdict b/sample_data/MIMIC-IV-meds/meds_reader/emar_id/zdict new file mode 100644 index 00000000..e0e7ebe9 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/emar_id/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/emar_seq/data b/sample_data/MIMIC-IV-meds/meds_reader/emar_seq/data new file mode 100644 index 00000000..d5285e61 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/emar_seq/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/emar_seq/zdict b/sample_data/MIMIC-IV-meds/meds_reader/emar_seq/zdict new file mode 100644 index 00000000..0b47f08b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/emar_seq/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/frequency/data b/sample_data/MIMIC-IV-meds/meds_reader/frequency/data new file mode 100644 index 00000000..5b8f3d5e Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/frequency/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/frequency/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/frequency/dictionary new file mode 100644 index 00000000..f4d23567 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/frequency/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/frequency/zdict b/sample_data/MIMIC-IV-meds/meds_reader/frequency/zdict new file mode 100644 index 00000000..58054492 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/frequency/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/hadm_id/data b/sample_data/MIMIC-IV-meds/meds_reader/hadm_id/data new file mode 100644 index 00000000..994d600f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/hadm_id/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/hadm_id/zdict b/sample_data/MIMIC-IV-meds/meds_reader/hadm_id/zdict new file mode 100644 index 00000000..e3c2a683 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/hadm_id/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/icustay_id/data b/sample_data/MIMIC-IV-meds/meds_reader/icustay_id/data new file mode 100644 index 00000000..96b051ae Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/icustay_id/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/icustay_id/zdict b/sample_data/MIMIC-IV-meds/meds_reader/icustay_id/zdict new file mode 100644 index 00000000..4dbdec92 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/icustay_id/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/insurance/data b/sample_data/MIMIC-IV-meds/meds_reader/insurance/data new file mode 100644 index 00000000..41d0341f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/insurance/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/insurance/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/insurance/dictionary new file mode 100644 index 00000000..83bab0f7 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/insurance/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/insurance/zdict b/sample_data/MIMIC-IV-meds/meds_reader/insurance/zdict new file mode 100644 index 00000000..3edba61f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/insurance/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/language/data b/sample_data/MIMIC-IV-meds/meds_reader/language/data new file mode 100644 index 00000000..130c33ad Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/language/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/language/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/language/dictionary new file mode 100644 index 00000000..05dc241d Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/language/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/language/zdict b/sample_data/MIMIC-IV-meds/meds_reader/language/zdict new file mode 100644 index 00000000..1955ae10 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/language/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/link_order_id/data b/sample_data/MIMIC-IV-meds/meds_reader/link_order_id/data new file mode 100644 index 00000000..6f248c58 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/link_order_id/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/link_order_id/zdict b/sample_data/MIMIC-IV-meds/meds_reader/link_order_id/zdict new file mode 100644 index 00000000..0b5aa442 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/link_order_id/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/marital_status/data b/sample_data/MIMIC-IV-meds/meds_reader/marital_status/data new file mode 100644 index 00000000..c77b080f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/marital_status/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/marital_status/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/marital_status/dictionary new file mode 100644 index 00000000..c6f6d96b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/marital_status/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/marital_status/zdict b/sample_data/MIMIC-IV-meds/meds_reader/marital_status/zdict new file mode 100644 index 00000000..14f311ac Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/marital_status/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.length b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.length new file mode 100644 index 00000000..1f63db11 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.length differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.null_map/data b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.null_map/data new file mode 100644 index 00000000..cad614d7 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.null_map/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.null_map/zdict b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.null_map/zdict new file mode 100644 index 00000000..383a0b55 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.null_map/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.properties b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.properties new file mode 100644 index 00000000..d5d27b96 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.properties differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.version b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.version new file mode 100644 index 00000000..0cfbf088 --- /dev/null +++ b/sample_data/MIMIC-IV-meds/meds_reader/meds_reader.version @@ -0,0 +1 @@ +2 diff --git a/sample_data/MIMIC-IV-meds/meds_reader/metadata/.shards.json b/sample_data/MIMIC-IV-meds/meds_reader/metadata/.shards.json new file mode 100644 index 00000000..18fdebdb --- /dev/null +++ b/sample_data/MIMIC-IV-meds/meds_reader/metadata/.shards.json @@ -0,0 +1 @@ +{"train/0": [10022281, 10004235, 10005909, 10018081, 10003046, 10036156, 10020306, 10002930, 10020944, 10018501, 10007795, 10035185, 10021938, 10021312, 10014078, 10006053, 10021118, 10023239, 10007058, 10024043, 10020786, 10005817, 10039997, 10015860, 10038999, 10021487, 10017492, 10020187, 10018845, 10019917, 10008287, 10011398, 10010867, 10037861, 10018328, 10029291, 10022041, 10010471, 10019568, 10005866, 10015931, 10021666, 10037928, 10032725, 10009628, 10002428, 10040025, 10014354, 10007818, 10003400, 10022017, 10007928, 10012552, 10027602, 10005348, 10023117, 10026406, 10008454, 10013049, 10025612, 10039708, 10035631, 10037975, 10009035, 10004422, 10020640, 10027445, 10000032, 10019003, 10006580, 10002495, 10015272, 10001217, 10004720, 10004733, 10031757, 10031404, 10038933, 10016150, 10004457], "tuning/0": [10020740, 10038992, 10009049, 10038081, 10039831, 10019777, 10019172, 10001725, 10016810, 10026255], "held_out/0": [10025463, 10019385, 10014729, 10022880, 10023771, 10012853, 10016742, 10018423, 10029484, 10026354]} \ No newline at end of file diff --git a/sample_data/MIMIC-IV-meds/meds_reader/metadata/codes.parquet b/sample_data/MIMIC-IV-meds/meds_reader/metadata/codes.parquet new file mode 100644 index 00000000..1124274b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/metadata/codes.parquet differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/metadata/dataset.json b/sample_data/MIMIC-IV-meds/meds_reader/metadata/dataset.json new file mode 100644 index 00000000..54477be8 --- /dev/null +++ b/sample_data/MIMIC-IV-meds/meds_reader/metadata/dataset.json @@ -0,0 +1 @@ +{"dataset_name": "MIMIC-IV", "dataset_version": "3.1:0.0.3", "etl_name": "MEDS_transforms", "etl_version": "0.1.1", "meds_version": "0.3.3", "created_at": "2025-02-23T20:51:17.919792"} \ No newline at end of file diff --git a/sample_data/MIMIC-IV-meds/meds_reader/metadata/subject_splits.parquet b/sample_data/MIMIC-IV-meds/meds_reader/metadata/subject_splits.parquet new file mode 100644 index 00000000..ec2ef4b0 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/metadata/subject_splits.parquet differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/numeric_value/data b/sample_data/MIMIC-IV-meds/meds_reader/numeric_value/data new file mode 100644 index 00000000..44e7bb06 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/numeric_value/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/numeric_value/zdict b/sample_data/MIMIC-IV-meds/meds_reader/numeric_value/zdict new file mode 100644 index 00000000..b4018f3d Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/numeric_value/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/order_id/data b/sample_data/MIMIC-IV-meds/meds_reader/order_id/data new file mode 100644 index 00000000..848bb189 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/order_id/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/order_id/zdict b/sample_data/MIMIC-IV-meds/meds_reader/order_id/zdict new file mode 100644 index 00000000..be65f18f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/order_id/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/data b/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/data new file mode 100644 index 00000000..7d7a44c5 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/dictionary new file mode 100644 index 00000000..688044d7 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/zdict b/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/zdict new file mode 100644 index 00000000..2e671ecf Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/ordercategorydescription/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/poe_id/data b/sample_data/MIMIC-IV-meds/meds_reader/poe_id/data new file mode 100644 index 00000000..38412fab Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/poe_id/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/poe_id/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/poe_id/dictionary new file mode 100644 index 00000000..e69de29b diff --git a/sample_data/MIMIC-IV-meds/meds_reader/poe_id/zdict b/sample_data/MIMIC-IV-meds/meds_reader/poe_id/zdict new file mode 100644 index 00000000..5759ea3c Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/poe_id/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/priority/data b/sample_data/MIMIC-IV-meds/meds_reader/priority/data new file mode 100644 index 00000000..b2e8b95b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/priority/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/priority/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/priority/dictionary new file mode 100644 index 00000000..1f1eb5be Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/priority/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/priority/zdict b/sample_data/MIMIC-IV-meds/meds_reader/priority/zdict new file mode 100644 index 00000000..edc04405 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/priority/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/race/data b/sample_data/MIMIC-IV-meds/meds_reader/race/data new file mode 100644 index 00000000..bb181a8f Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/race/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/race/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/race/dictionary new file mode 100644 index 00000000..66680033 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/race/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/race/zdict b/sample_data/MIMIC-IV-meds/meds_reader/race/zdict new file mode 100644 index 00000000..fd3a7d64 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/race/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/route/data b/sample_data/MIMIC-IV-meds/meds_reader/route/data new file mode 100644 index 00000000..52066e0b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/route/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/route/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/route/dictionary new file mode 100644 index 00000000..ad1fea51 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/route/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/route/zdict b/sample_data/MIMIC-IV-meds/meds_reader/route/zdict new file mode 100644 index 00000000..5d6df218 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/route/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/data b/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/data new file mode 100644 index 00000000..717a4f0c Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/dictionary new file mode 100644 index 00000000..1f775c89 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/zdict b/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/zdict new file mode 100644 index 00000000..505566f6 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/statusdescription/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/subject_id b/sample_data/MIMIC-IV-meds/meds_reader/subject_id new file mode 100644 index 00000000..3bc2c336 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/subject_id differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/text_value/data b/sample_data/MIMIC-IV-meds/meds_reader/text_value/data new file mode 100644 index 00000000..7886518b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/text_value/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/text_value/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/text_value/dictionary new file mode 100644 index 00000000..ccce2d32 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/text_value/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/text_value/zdict b/sample_data/MIMIC-IV-meds/meds_reader/text_value/zdict new file mode 100644 index 00000000..b4dce81b Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/text_value/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/time/data b/sample_data/MIMIC-IV-meds/meds_reader/time/data new file mode 100644 index 00000000..8a3848f4 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/time/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/time/zdict b/sample_data/MIMIC-IV-meds/meds_reader/time/zdict new file mode 100644 index 00000000..4732fa23 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/time/zdict differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/unit/data b/sample_data/MIMIC-IV-meds/meds_reader/unit/data new file mode 100644 index 00000000..3fd594e9 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/unit/data differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/unit/dictionary b/sample_data/MIMIC-IV-meds/meds_reader/unit/dictionary new file mode 100644 index 00000000..d4e6921e Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/unit/dictionary differ diff --git a/sample_data/MIMIC-IV-meds/meds_reader/unit/zdict b/sample_data/MIMIC-IV-meds/meds_reader/unit/zdict new file mode 100644 index 00000000..4695cb17 Binary files /dev/null and b/sample_data/MIMIC-IV-meds/meds_reader/unit/zdict differ diff --git a/src/cehrbert/models/hf_models/tokenization_hf_cehrbert.py b/src/cehrbert/models/hf_models/tokenization_hf_cehrbert.py index bb138b42..46a39e0a 100644 --- a/src/cehrbert/models/hf_models/tokenization_hf_cehrbert.py +++ b/src/cehrbert/models/hf_models/tokenization_hf_cehrbert.py @@ -379,11 +379,12 @@ def batched_generator(): ) if data_args.streaming: + first_example = next(iter(dataset)) parts = dataset.map( partial(agg_helper, map_func=map_statistics_partial), batched=True, batch_size=data_args.preprocessing_batch_size, - remove_columns=dataset.column_names, + remove_columns=first_example.keys(), ) else: parts = dataset.map( diff --git a/src/cehrbert/runners/hf_cehrbert_finetune_runner.py b/src/cehrbert/runners/hf_cehrbert_finetune_runner.py index 3e111bb9..39ce1114 100644 --- a/src/cehrbert/runners/hf_cehrbert_finetune_runner.py +++ b/src/cehrbert/runners/hf_cehrbert_finetune_runner.py @@ -95,8 +95,17 @@ def load_finetuned_model(model_args: ModelArguments, model_name_or_path: str) -> def main(): + data_args, model_args, training_args = parse_runner_args() + if data_args.streaming: + # This happens only when streaming is enabled. This is for disabling the warning message + # https://github.com/huggingface/transformers/issues/5486 + os.environ["TOKENIZERS_PARALLELISM"] = "false" + # The iterable dataset doesn't have sharding implemented, so the number of works has to + # be set to 0. Otherwise the trainer will throw an error + training_args.dataloader_num_workers = 0 + tokenizer = load_pretrained_tokenizer(model_args) prepared_ds_path = generate_prepared_ds_path(data_args, model_args, data_folder=data_args.cohort_folder) diff --git a/src/cehrbert/runners/hf_cehrbert_pretrain_runner.py b/src/cehrbert/runners/hf_cehrbert_pretrain_runner.py index 8c190b38..89829935 100644 --- a/src/cehrbert/runners/hf_cehrbert_pretrain_runner.py +++ b/src/cehrbert/runners/hf_cehrbert_pretrain_runner.py @@ -2,8 +2,9 @@ import os from typing import Optional, Union +import torch from datasets import Dataset, DatasetDict, IterableDatasetDict, load_from_disk -from transformers import AutoConfig, Trainer, set_seed +from transformers import Trainer, set_seed from transformers.utils import logging from cehrbert.data_generators.hf_data_generator.hf_dataset import create_cehrbert_pretraining_dataset @@ -95,7 +96,7 @@ def load_and_create_model(model_args: ModelArguments, tokenizer: CehrBertTokeniz model = load_and_create_model(model_args, tokenizer) """ try: - model_config = AutoConfig.from_pretrained(os.path.expanduser(model_args.model_name_or_path)) + model_config = CehrBertConfig.from_pretrained(os.path.expanduser(model_args.model_name_or_path)) except (OSError, ValueError, FileNotFoundError, json.JSONDecodeError) as e: LOG.warning(e) model_config = CehrBertConfig( @@ -262,9 +263,17 @@ def filter_func(examples): if not data_args.streaming: processed_dataset.set_format("pt") + def data_collator(features): + batch = collator(features) + # Convert any float64 tensors to float32 + for key in batch: + if isinstance(batch[key], torch.Tensor) and batch[key].dtype == torch.float64: + batch[key] = batch[key].to(torch.float32) + return batch + trainer = Trainer( model=model, - data_collator=collator, + data_collator=data_collator, train_dataset=processed_dataset["train"], eval_dataset=processed_dataset["validation"], args=training_args, diff --git a/tests/integration_tests/runners/hf_cehrbert_pretrain_runner_meds_streaming_test.py b/tests/integration_tests/runners/hf_cehrbert_pretrain_runner_meds_streaming_test.py new file mode 100644 index 00000000..6f76e49c --- /dev/null +++ b/tests/integration_tests/runners/hf_cehrbert_pretrain_runner_meds_streaming_test.py @@ -0,0 +1,56 @@ +import os +import shutil +import sys +import tempfile +import unittest +from pathlib import Path + +from datasets import disable_caching + +from cehrbert.runners.hf_cehrbert_pretrain_runner import main + +disable_caching() +os.environ["CUDA_VISIBLE_DEVICES"] = "" +os.environ["WANDB_MODE"] = "disabled" +os.environ["TRANSFORMERS_VERBOSITY"] = "info" + + +class HfCehrBertRunnerIntegrationTest(unittest.TestCase): + def setUp(self): + # Get the root folder of the project + root_folder = Path(os.path.abspath(__file__)).parent.parent.parent.parent + data_folder = os.path.join(root_folder, "sample_data", "MIMIC-IV-meds", "meds_reader") + # Create a temporary directory to store model and tokenizer + self.temp_dir = tempfile.mkdtemp() + self.model_folder_path = os.path.join(self.temp_dir, "model") + Path(self.model_folder_path).mkdir(parents=True, exist_ok=True) + self.dataset_prepared_path = os.path.join(self.temp_dir, "dataset_prepared_path") + Path(self.dataset_prepared_path).mkdir(parents=True, exist_ok=True) + sys.argv = [ + "hf_cehrbert_pretraining_runner.py", + "--model_name_or_path", + self.model_folder_path, + "--tokenizer_name_or_path", + self.model_folder_path, + "--output_dir", + self.model_folder_path, + "--data_folder", + data_folder, + "--dataset_prepared_path", + self.dataset_prepared_path, + "--max_steps", + "10", + "--streaming", + "--is_data_in_meds", + ] + + def tearDown(self): + # Remove the temporary directory + shutil.rmtree(self.temp_dir) + + def test_train_model(self): + main() + + +if __name__ == "__main__": + unittest.main()