cehrbert-data 0.0.1__tar.gz → 0.0.3__tar.gz

cehrbert_data-0.0.3/.github/workflows/python-build.yml

@@ -0,0 +1,93 @@
+name: Publish Python 🐍 distribution 📦 to PyPI and TestPyPI
+
+on: push
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.x"
+    - name: Install pypa/build
+      run: >-
+        python3 -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    if: startsWith(github.ref, 'refs/tags/')  # only publish to PyPI on tag pushes
+    needs:
+    - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/cehrbert-data  # Replace <package-name> with your PyPI project name
+    permissions:
+      id-token: write  # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    name: >-
+      Sign the Python 🐍 distribution 📦 with Sigstore
+      and upload them to GitHub Release
+    needs:
+    - publish-to-pypi
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write  # IMPORTANT: mandatory for making GitHub Releases
+      id-token: write  # IMPORTANT: mandatory for sigstore
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Sign the dists with Sigstore
+      uses: sigstore/gh-action-sigstore-python@v2.1.1
+      with:
+        inputs: >-
+          ./dist/*.tar.gz
+          ./dist/*.whl
+    - name: Create GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release create
+        '${{ github.ref_name }}'
+        --repo '${{ github.repository }}'
+        --notes ""
+    - name: Upload artifact signatures to GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      # Upload to GitHub Release using the `gh` CLI.
+      # `dist/` contains the built packages, and the
+      # sigstore-produced signatures and certificates.
+      run: >-
+        gh release upload
+        '${{ github.ref_name }}' dist/**
+        --repo '${{ github.repository }}'

{cehrbert_data-0.0.1 → cehrbert_data-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cehrbert_data
-Version: 0.0.1
+Version: 0.0.3
 Summary: The Spark ETL tools for generating the CEHR-BERT and CEHR-GPT pre-training and finetuning data
 Author-email: Chao Pang <chaopang229@gmail.com>, Xinzhuo Jiang <xj2193@cumc.columbia.edu>, Krishna Kalluri <kk3326@cumc.columbia.edu>, Nishanth Parameshwar Pavinkurve <np2689@cumc.columbia.edu>, Karthik Natarajan <kn2174@cumc.columbia.edu>
 License: MIT License

cehrbert_data-0.0.3/src/cehrbert_data/apps/generate_required_labs.py

@@ -0,0 +1,140 @@
+import argparse
+import os
+
+from pyspark.sql import SparkSession, DataFrame
+from pyspark.sql import functions as F
+from pyspark.sql.window import Window
+
+from cehrbert_data.const.common import (
+    CONCEPT,
+    MEASUREMENT,
+    REQUIRED_MEASUREMENT,
+    NUMERIC_MEASUREMENT_STATS
+)
+from cehrbert_data.utils.spark_utils import preprocess_domain_table
+from cehrbert_data.queries.measurement_queries import (
+    LAB_PREVALENCE_QUERY,
+    MEASUREMENT_UNIT_STATS_QUERY
+)
+
+
+def main(
+        input_folder,
+        output_folder,
+        num_of_numeric_labs,
+        num_of_categorical_labs,
+        min_num_of_patients
+):
+    spark = SparkSession.builder.appName("Generate required labs").getOrCreate()
+
+    # Load measurement as a dataframe in pyspark
+    measurement = preprocess_domain_table(spark, input_folder, MEASUREMENT)
+    # Load concept as a dataframe in pyspark
+    concept = preprocess_domain_table(spark, input_folder, CONCEPT)
+    # Create the local measurement view
+    measurement.createOrReplaceTempView(MEASUREMENT)
+    # Create the local concept view
+    concept.createOrReplaceTempView(CONCEPT)
+    # Create the
+    required_lab_dataframe = generate_required_labs(
+        spark, num_of_numeric_labs, num_of_categorical_labs, min_num_of_patients
+    )
+    required_lab_dataframe.write.mode("overwrite").parquet(
+        os.path.join(output_folder, REQUIRED_MEASUREMENT)
+    )
+    # Reload the dataframe from the disk
+    required_lab_dataframe = spark.read.parquet(
+        os.path.join(output_folder, REQUIRED_MEASUREMENT)
+    )
+    required_lab_dataframe.createOrReplaceTempView(REQUIRED_MEASUREMENT)
+    numeric_measurement_stats_dataframe = spark.sql(MEASUREMENT_UNIT_STATS_QUERY)
+    numeric_measurement_stats_dataframe.write.mode("overwrite").parquet(
+        os.path.join(output_folder, NUMERIC_MEASUREMENT_STATS)
+    )
+
+
+def generate_required_labs(
+        spark: SparkSession,
+        num_of_numeric_labs: int,
+        num_of_categorical_labs: int,
+        min_num_of_patients: int
+) -> DataFrame:
+
+    prevalent_labs = spark.sql(LAB_PREVALENCE_QUERY)
+    prevalent_labs = prevalent_labs.where(F.col("person_count") >= min_num_of_patients)
+    # Cache the dataframe for faster computation in the below transformations
+    prevalent_labs.cache()
+    prevalent_numeric_labs = (
+        prevalent_labs.withColumn("is_numeric", F.col("numeric_percentage") >= 0.5)
+        .where("is_numeric")
+        .withColumn("rn", F.row_number().over(Window.orderBy(F.desc("freq"))))
+        .where(F.col("rn") <= num_of_numeric_labs)
+        .drop("rn")
+    )
+    prevalent_categorical_labs = (
+        prevalent_labs.withColumn("is_categorical", F.col("categorical_percentage") >= 0.5)
+        .where("is_categorical")
+        .withColumn("is_numeric", ~F.col("is_categorical"))
+        .withColumn("rn", F.row_number().over(Window.orderBy(F.desc("freq"))))
+        .where(F.col("rn") <= num_of_categorical_labs)
+        .drop("is_categorical")
+        .drop("rn")
+    )
+    return prevalent_numeric_labs.unionAll(prevalent_categorical_labs)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Arguments for generate " "required labs to be included")
+    parser.add_argument(
+        "-i",
+        "--input_folder",
+        dest="input_folder",
+        action="store",
+        help="The path for your input_folder where the raw data is",
+        required=True,
+    )
+    parser.add_argument(
+        "-o",
+        "--output_folder",
+        dest="output_folder",
+        action="store",
+        help="The path for your output_folder",
+        required=True,
+    )
+    parser.add_argument(
+        "--num_of_numeric_labs",
+        dest="num_of_numeric_labs",
+        action="store",
+        type=int,
+        default=100,
+        help="The top most prevalent numeric labs to be included",
+        required=False,
+    )
+    parser.add_argument(
+        "--num_of_categorical_labs",
+        dest="num_of_categorical_labs",
+        action="store",
+        type=int,
+        default=100,
+        help="The top most prevalent categorical labs to be included",
+        required=False,
+    )
+    parser.add_argument(
+        "--min_num_of_patients",
+        dest="min_num_of_patients",
+        action="store",
+        type=int,
+        default=0,
+        help="Min no.of patients linked to concepts to be included",
+        required=False,
+    )
+
+    ARGS = parser.parse_args()
+
+    main(
+        ARGS.input_folder,
+        ARGS.output_folder,
+        ARGS.num_of_numeric_labs,
+        ARGS.num_of_categorical_labs,
+        ARGS.min_num_of_patients
+    )

{cehrbert_data-0.0.1 → cehrbert_data-0.0.3}/src/cehrbert_data/apps/generate_training_data.py

@@ -8,39 +8,45 @@ from pyspark.sql import SparkSession
 from pyspark.sql import functions as F
 from pyspark.sql.window import Window
 
-from cehrbert_data.const.common import DEATH, MEASUREMENT, PERSON, REQUIRED_MEASUREMENT, VISIT_OCCURRENCE
-from cehrbert_data.decorators.patient_event_decorator import AttType
+from cehrbert_data.const.common import (
+    PERSON,
+    VISIT_OCCURRENCE,
+    DEATH,
+    MEASUREMENT
+)
+from cehrbert_data.decorators import AttType
 from cehrbert_data.utils.spark_utils import (
     create_sequence_data,
     create_sequence_data_with_att,
     join_domain_tables,
     preprocess_domain_table,
-    process_measurement,
+    get_measurement_table,
     validate_table_names,
 )
 
 
 def main(
-    input_folder,
-    output_folder,
-    domain_table_list,
-    date_filter,
-    include_visit_type,
-    is_new_patient_representation,
-    exclude_visit_tokens,
-    is_classic_bert,
-    include_prolonged_stay,
-    include_concept_list: bool,
-    gpt_patient_sequence: bool,
-    apply_age_filter: bool,
-    include_death: bool,
-    att_type: AttType,
-    include_sequence_information_content: bool = False,
-    exclude_demographic: bool = False,
-    use_age_group: bool = False,
-    with_drug_rollup: bool = True,
-    include_inpatient_hour_token: bool = False,
-    continue_from_events: bool = False,
+        input_folder,
+        output_folder,
+        domain_table_list,
+        date_filter,
+        include_visit_type,
+        is_new_patient_representation,
+        exclude_visit_tokens,
+        is_classic_bert,
+        include_prolonged_stay,
+        include_concept_list: bool,
+        gpt_patient_sequence: bool,
+        apply_age_filter: bool,
+        include_death: bool,
+        att_type: AttType,
+        include_sequence_information_content: bool = False,
+        exclude_demographic: bool = False,
+        use_age_group: bool = False,
+        with_drug_rollup: bool = True,
+        include_inpatient_hour_token: bool = False,
+        continue_from_events: bool = False,
+        refresh_measurement: bool = False
 ):
     spark = SparkSession.builder.appName("Generate CEHR-BERT Training Data").getOrCreate()
 
@@ -63,6 +69,7 @@ def main(
         f"exclude_demographic: {exclude_demographic}\n"
         f"use_age_group: {use_age_group}\n"
         f"with_drug_rollup: {with_drug_rollup}\n"
+        f"refresh_measurement: {refresh_measurement}\n"
     )
 
     domain_tables = []
@@ -117,17 +124,16 @@ def main(
 
     # Process the measurement table if exists
     if MEASUREMENT in domain_table_list:
-        measurement = preprocess_domain_table(spark, input_folder, MEASUREMENT)
-        required_measurement = preprocess_domain_table(spark, input_folder, REQUIRED_MEASUREMENT)
-        # The select is necessary to make sure the order of the columns is the same as the
-        # original dataframe, otherwise the union might use the wrong columns
-        scaled_measurement = process_measurement(spark, measurement, required_measurement, output_folder)
-
+        processed_measurement = get_measurement_table(
+            spark,
+            input_folder,
+            refresh=refresh_measurement
+        )
         if patient_events:
             # Union all measurement records together with other domain records
-            patient_events = patient_events.unionByName(scaled_measurement)
+            patient_events = patient_events.unionByName(processed_measurement)
         else:
-            patient_events = scaled_measurement
+            patient_events = processed_measurement
 
     patient_events = (
         patient_events.join(visit_occurrence_person, "visit_occurrence_id")
@@ -206,17 +212,21 @@ def main(
     patient_splits_folder = os.path.join(input_folder, "patient_splits")
     if os.path.exists(patient_splits_folder):
         patient_splits = spark.read.parquet(patient_splits_folder)
-        sequence_data.join(patient_splits, "person_id").write.mode("overwrite").parquet(
-            os.path.join(output_folder, "patient_sequence", "temp")
+        temp_folder = os.path.join(output_folder, "patient_sequence", "temp")
+        sequence_data.join(
+            patient_splits.select("person_id", "split"),
+            "person_id"
+        ).write.mode("overwrite").parquet(
+            temp_folder
         )
-        sequence_data = spark.read.parquet(os.path.join(output_folder, "patient_sequence", "temp"))
+        sequence_data = spark.read.parquet(temp_folder)
         sequence_data.where('split="train"').write.mode("overwrite").parquet(
             os.path.join(output_folder, "patient_sequence/train")
         )
         sequence_data.where('split="test"').write.mode("overwrite").parquet(
             os.path.join(output_folder, "patient_sequence/test")
         )
-        shutil.rmtree(os.path.join(output_folder, "patient_sequence", "temp"))
+        shutil.rmtree(temp_folder)
     else:
         sequence_data.write.mode("overwrite").parquet(os.path.join(output_folder, "patient_sequence"))
 
@@ -304,7 +314,16 @@ if __name__ == "__main__":
         dest="include_inpatient_hour_token",
         action="store_true",
     )
-    parser.add_argument("--continue_from_events", dest="continue_from_events", action="store_true")
+    parser.add_argument(
+        "--continue_from_events",
+        dest="continue_from_events",
+        action="store_true"
+    )
+    parser.add_argument(
+        "--refresh_measurement",
+        dest="refresh_measurement",
+        action="store_true"
+    )
     parser.add_argument(
         "--att_type",
         dest="att_type",
@@ -334,4 +353,5 @@ if __name__ == "__main__":
         with_drug_rollup=ARGS.with_drug_rollup,
         include_inpatient_hour_token=ARGS.include_inpatient_hour_token,
         continue_from_events=ARGS.continue_from_events,
+        refresh_measurement=ARGS.refresh_measurement,
     )