cehrbert-data 0.0.2__tar.gz → 0.0.4__tar.gz

{cehrbert_data-0.0.2 → cehrbert_data-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cehrbert_data
-Version: 0.0.2
+Version: 0.0.4
 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.4/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.2 → cehrbert_data-0.0.4}/src/cehrbert_data/apps/generate_training_data.py

@@ -8,39 +8,46 @@ 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,
 )
+from cehrbert_data.utils.logging_utils import add_console_logging
 
 
 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 +70,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 +125,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 +213,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 +315,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",
@@ -314,6 +334,9 @@ if __name__ == "__main__":
 
     ARGS = parser.parse_args()
 
+    # Enable logging
+    add_console_logging()
+
     main(
         ARGS.input_folder,
         ARGS.output_folder,
@@ -334,4 +357,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,
     )

{cehrbert_data-0.0.2 → cehrbert_data-0.0.4}/src/cehrbert_data/cohorts/spark_app_base.py

@@ -10,9 +10,9 @@ from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql import functions as F
 from pyspark.sql.window import Window
 
-from cehrbert_data.decorators.patient_event_decorator import AttType
+from cehrbert_data.decorators import AttType
+from cehrbert_data.const.common import VISIT_OCCURRENCE
 from cehrbert_data.utils.spark_utils import (
-    VISIT_OCCURRENCE,
     build_ancestry_table_for,
     create_concept_frequency_data,
     create_hierarchical_sequence_data,
@@ -22,6 +22,7 @@ from cehrbert_data.utils.spark_utils import (
     get_descendant_concept_ids,
     preprocess_domain_table,
 )
+from cehrbert_data.utils.logging_utils import add_console_logging
 
 from ..cohorts.query_builder import ENTRY_COHORT, NEGATIVE_COHORT, QueryBuilder
 
@@ -97,6 +98,7 @@ class BaseCohortBuilder(ABC):
             age_upper_bound: int,
             prior_observation_period: int,
             post_observation_period: int,
+            continue_job: bool = False
     ):
 
         self._query_builder = query_builder
@@ -110,6 +112,7 @@ class BaseCohortBuilder(ABC):
         self._post_observation_period = post_observation_period
         cohort_name = re.sub("[^a-z0-9]+", "_", self._query_builder.get_cohort_name().lower())
         self._output_data_folder = os.path.join(self._output_folder, cohort_name)
+        self._continue_job = continue_job
 
         self.get_logger().info(
             f"query_builder: {query_builder}\n"
@@ -121,6 +124,7 @@ class BaseCohortBuilder(ABC):
             f"age_upper_bound: {age_upper_bound}\n"
             f"prior_observation_period: {prior_observation_period}\n"
             f"post_observation_period: {post_observation_period}\n"
+            f"continue_job: {continue_job}\n"
         )
 
         # Validate the age range, observation_window and prediction_window
@@ -187,6 +191,11 @@ class BaseCohortBuilder(ABC):
 
     def build(self):
         """Build the cohort and write the dataframe as parquet files to _output_data_folder."""
+
+        # Check whether the cohort has been generated
+        if self._continue_job and self.cohort_exists():
+            return self
+
         cohort = self.create_cohort()
 
         cohort = self._apply_observation_period(cohort)
@@ -201,6 +210,13 @@ class BaseCohortBuilder(ABC):
 
         return self
 
+    def cohort_exists(self) -> bool:
+        try:
+            self.load_cohort()
+            return True
+        except Exception:
+            return False
+
     def load_cohort(self):
         return self.spark.read.parquet(self._output_data_folder)
 
@@ -276,7 +292,9 @@ class NestedCohortBuilder:
             exclude_visit_tokens: bool = False,
             is_feature_concept_frequency: bool = False,
             is_roll_up_concept: bool = False,
+            is_drug_roll_up_concept: bool = True,
             include_concept_list: bool = True,
+            refresh_measurement: bool = False,
             is_new_patient_representation: bool = False,
             gpt_patient_sequence: bool = False,
             is_hierarchical_bert: bool = False,
@@ -312,6 +330,7 @@ class NestedCohortBuilder:
         self._classic_bert_seq = classic_bert_seq
         self._is_feature_concept_frequency = is_feature_concept_frequency
         self._is_roll_up_concept = is_roll_up_concept
+        self._is_drug_roll_up_concept = is_drug_roll_up_concept
         self._is_new_patient_representation = is_new_patient_representation
         self._gpt_patient_sequence = gpt_patient_sequence
         self._is_hierarchical_bert = is_hierarchical_bert
@@ -320,6 +339,7 @@ class NestedCohortBuilder:
         self._is_questionable_outcome_existed = is_questionable_outcome_existed
         self._is_prediction_window_unbounded = is_prediction_window_unbounded
         self._include_concept_list = include_concept_list
+        self._refresh_measurement = refresh_measurement
         self._allow_measurement_only = allow_measurement_only
         self._output_data_folder = os.path.join(
             self._output_folder, re.sub("[^a-z0-9]+", "_", self._cohort_name.lower())
@@ -347,6 +367,7 @@ class NestedCohortBuilder:
             f"allow_measurement_only: {allow_measurement_only}\n"
             f"is_feature_concept_frequency: {is_feature_concept_frequency}\n"
             f"is_roll_up_concept: {is_roll_up_concept}\n"
+            f"is_drug_roll_up_concept: {is_drug_roll_up_concept}\n"
             f"is_new_patient_representation: {is_new_patient_representation}\n"
             f"gpt_patient_sequence: {gpt_patient_sequence}\n"
             f"is_hierarchical_bert: {is_hierarchical_bert}\n"
@@ -355,6 +376,7 @@ class NestedCohortBuilder:
             f"is_remove_index_prediction_starts: {is_remove_index_prediction_starts}\n"
             f"is_prediction_window_unbounded: {is_prediction_window_unbounded}\n"
             f"include_concept_list: {include_concept_list}\n"
+            f"refresh_measurement: {refresh_measurement}\n"
             f"is_observation_window_unbounded: {is_observation_window_unbounded}\n"
             f"is_population_estimation: {is_population_estimation}\n"
             f"att_type: {att_type}\n"
@@ -494,11 +516,21 @@ class NestedCohortBuilder:
             .where(F.col("num_of_concepts") >= self._num_of_concepts)
         )
 
-        # Add time_to_event
-        cohort = cohort.withColumn(
-            "study_end_date",
-            F.coalesce(F.col("outcome_date"), F.date_add(cohort.index_date, self._prediction_window))
-        )
+        if self._is_prediction_window_unbounded:
+            observation_period = self._dependency_dict[OBSERVATION_PERIOD]
+            # Add time_to_event
+            cohort = cohort.join(
+                observation_period.select("person_id", "observation_period_end_date"), "person_id"
+            ).withColumn(
+                "study_end_date",
+                F.coalesce(F.col("outcome_date"), F.col("observation_period_end_date"))
+            ).drop("observation_period_end_date")
+        else:
+            # Add time_to_event
+            cohort = cohort.withColumn(
+                "study_end_date",
+                F.coalesce(F.col("outcome_date"), F.date_add(cohort.index_date, self._prediction_window))
+            )
         cohort = cohort.withColumn("time_to_event", F.datediff("study_end_date", "index_date"))
 
         # if patient_splits is provided, we will
@@ -526,12 +558,14 @@ class NestedCohortBuilder:
         """
         # Extract all ehr records for the patients
         ehr_records = extract_ehr_records(
-            self.spark,
-            self._input_folder,
-            self._ehr_table_list,
-            self._include_visit_type,
-            self._is_roll_up_concept,
-            self._include_concept_list,
+            spark=self.spark,
+            input_folder=self._input_folder,
+            domain_table_list=self._ehr_table_list,
+            include_visit_type=self._include_visit_type,
+            with_diagnosis_rollup=self._is_roll_up_concept,
+            with_drug_rollup=self._is_drug_roll_up_concept,
+            include_concept_list=self._include_concept_list,
+            refresh_measurement=self._refresh_measurement,
         )
 
         # Duplicate the records for cohorts that allow multiple entries
@@ -644,54 +678,31 @@ def create_prediction_cohort(
     :param ehr_table_list:
     :return:
     """
-    cohort_name = spark_args.cohort_name
-    input_folder = spark_args.input_folder
-    output_folder = spark_args.output_folder
-    date_lower_bound = spark_args.date_lower_bound
-    date_upper_bound = spark_args.date_upper_bound
-    age_lower_bound = spark_args.age_lower_bound
-    age_upper_bound = spark_args.age_upper_bound
-    observation_window = spark_args.observation_window
-    prediction_start_days = spark_args.prediction_start_days
-    prediction_window = spark_args.prediction_window
-    hold_off_window = spark_args.hold_off_window
-    num_of_visits = spark_args.num_of_visits
-    num_of_concepts = spark_args.num_of_concepts
-    include_visit_type = spark_args.include_visit_type
-    exclude_visit_tokens = spark_args.exclude_visit_tokens
-    is_feature_concept_frequency = spark_args.is_feature_concept_frequency
-    is_roll_up_concept = spark_args.is_roll_up_concept
-    is_window_post_index = spark_args.is_window_post_index
-    is_new_patient_representation = spark_args.is_new_patient_representation
-    is_hierarchical_bert = spark_args.is_hierarchical_bert
-    classic_bert_seq = spark_args.classic_bert_seq
-    is_first_time_outcome = spark_args.is_first_time_outcome
-    is_prediction_window_unbounded = spark_args.is_prediction_window_unbounded
-    is_observation_window_unbounded = spark_args.is_observation_window_unbounded
-    # If the outcome negative query exists, that means we need to remove those questionable
-    # outcomes from the target cohort
-    is_questionable_outcome_existed = outcome_query_builder.get_negative_query() is not None
-
-    # Do we want to remove those records whose outcome occur between index_date and the start of
-    # the prediction window
-    is_remove_index_prediction_starts = spark_args.is_remove_index_prediction_starts
+    # Add logging to spark application output when enable_logging is set to True
+    if spark_args.enable_logging:
+        add_console_logging()
 
     # Toggle the prior/post observation_period depending on the is_window_post_index flag
-    prior_observation_period = 0 if is_window_post_index else observation_window + hold_off_window
-    post_observation_period = observation_window + hold_off_window if is_window_post_index else 0
+    prior_observation_period = (
+        0 if spark_args.is_window_post_index else spark_args.observation_window + spark_args.hold_off_window
+    )
+    post_observation_period = (
+        spark_args.observation_window + spark_args.hold_off_window if spark_args.is_window_post_index else 0
+    )
 
     # Generate the target cohort
     target_cohort = (
         BaseCohortBuilder(
             query_builder=target_query_builder,
-            input_folder=input_folder,
-            output_folder=output_folder,
-            date_lower_bound=date_lower_bound,
-            date_upper_bound=date_upper_bound,
-            age_lower_bound=age_lower_bound,
-            age_upper_bound=age_upper_bound,
+            input_folder=spark_args.input_folder,
+            output_folder=spark_args.output_folder,
+            date_lower_bound=spark_args.date_lower_bound,
+            date_upper_bound=spark_args.date_upper_bound,
+            age_lower_bound=spark_args.age_lower_bound,
+            age_upper_bound=spark_args.age_upper_bound,
             prior_observation_period=prior_observation_period,
             post_observation_period=post_observation_period,
+            continue_job=spark_args.continue_job
         )
         .build()
         .load_cohort()
@@ -701,12 +712,12 @@ def create_prediction_cohort(
     outcome_cohort = (
         BaseCohortBuilder(
             query_builder=outcome_query_builder,
-            input_folder=input_folder,
-            output_folder=output_folder,
-            date_lower_bound=date_lower_bound,
-            date_upper_bound=date_upper_bound,
-            age_lower_bound=age_lower_bound,
-            age_upper_bound=age_upper_bound,
+            input_folder=spark_args.input_folder,
+            output_folder=spark_args.output_folder,
+            date_lower_bound=spark_args.date_lower_bound,
+            date_upper_bound=spark_args.date_upper_bound,
+            age_lower_bound=spark_args.age_lower_bound,
+            age_upper_bound=spark_args.age_upper_bound,
             prior_observation_period=0,
             post_observation_period=0,
         )
@@ -715,35 +726,41 @@ def create_prediction_cohort(
     )
 
     NestedCohortBuilder(
-        cohort_name=cohort_name,
-        input_folder=input_folder,
-        output_folder=output_folder,
+        cohort_name=spark_args.cohort_name,
+        input_folder=spark_args.input_folder,
+        output_folder=spark_args.output_folder,
         patient_splits_folder=spark_args.patient_splits_folder,
         target_cohort=target_cohort,
         outcome_cohort=outcome_cohort,
         ehr_table_list=ehr_table_list,
-        observation_window=observation_window,
-        hold_off_window=hold_off_window,
-        prediction_start_days=prediction_start_days,
-        prediction_window=prediction_window,
-        num_of_visits=num_of_visits,
-        num_of_concepts=num_of_concepts,
-        is_window_post_index=is_window_post_index,
-        include_visit_type=include_visit_type,
-        exclude_visit_tokens=exclude_visit_tokens,
+        observation_window=spark_args.observation_window,
+        hold_off_window=spark_args.hold_off_window,
+        prediction_start_days=spark_args.prediction_start_days,
+        prediction_window=spark_args.prediction_window,
+        num_of_visits=spark_args.num_of_visits,
+        num_of_concepts=spark_args.num_of_concepts,
+        is_window_post_index=spark_args.is_window_post_index,
+        include_visit_type=spark_args.include_visit_type,
+        exclude_visit_tokens=spark_args.exclude_visit_tokens,
         allow_measurement_only=spark_args.allow_measurement_only,
-        is_feature_concept_frequency=is_feature_concept_frequency,
-        is_roll_up_concept=is_roll_up_concept,
+        is_feature_concept_frequency=spark_args.is_feature_concept_frequency,
+        is_roll_up_concept=spark_args.is_roll_up_concept,
+        is_drug_roll_up_concept=spark_args.is_drug_roll_up_concept,
         include_concept_list=spark_args.include_concept_list,
-        is_new_patient_representation=is_new_patient_representation,
+        refresh_measurement=spark_args.refresh_measurement,
+        is_new_patient_representation=spark_args.is_new_patient_representation,
         gpt_patient_sequence=spark_args.gpt_patient_sequence,
-        is_hierarchical_bert=is_hierarchical_bert,
-        classic_bert_seq=classic_bert_seq,
-        is_first_time_outcome=is_first_time_outcome,
-        is_questionable_outcome_existed=is_questionable_outcome_existed,
-        is_prediction_window_unbounded=is_prediction_window_unbounded,
-        is_remove_index_prediction_starts=is_remove_index_prediction_starts,
-        is_observation_window_unbounded=is_observation_window_unbounded,
+        is_hierarchical_bert=spark_args.is_hierarchical_bert,
+        classic_bert_seq=spark_args.classic_bert_seq,
+        is_first_time_outcome=spark_args.is_first_time_outcome,
+        # If the outcome negative query exists, that means we need to remove those questionable
+        # outcomes from the target cohort
+        is_questionable_outcome_existed=outcome_query_builder.get_negative_query() is not None,
+        is_prediction_window_unbounded=spark_args.is_prediction_window_unbounded,
+        # Do we want to remove those records whose outcome occur between index_date
+        # and the start of the prediction window
+        is_remove_index_prediction_starts=spark_args.is_remove_index_prediction_starts,
+        is_observation_window_unbounded=spark_args.is_observation_window_unbounded,
         is_population_estimation=spark_args.is_population_estimation,
         att_type=AttType(spark_args.att_type),
         exclude_demographic=spark_args.exclude_demographic,

cehrbert_data-0.0.4/src/cehrbert_data/const/artificial_tokens.py

@@ -0,0 +1,3 @@
+VS_TOKEN = "[VS]"
+VE_TOKEN = "[VE]"
+DEATH_TOKEN = "[DEATH]"

{cehrbert_data-0.0.2 → cehrbert_data-0.0.4}/src/cehrbert_data/const/common.py

@@ -6,6 +6,7 @@ DRUG_EXPOSURE = "drug_exposure"
 DEVICE_EXPOSURE = "device_exposure"
 OBSERVATION = "observation"
 MEASUREMENT = "measurement"
+PROCESSED_MEASUREMENT = "processed_measurement"
 CATEGORICAL_MEASUREMENT = "categorical_measurement"
 OBSERVATION_PERIOD = "observation_period"
 DEATH = "death"
@@ -23,6 +24,10 @@ CDM_TABLES = [
     DEATH,
 ]
 REQUIRED_MEASUREMENT = "required_measurement"
+NUMERIC_MEASUREMENT_STATS = "numeric_measurement_stats"
 UNKNOWN_CONCEPT = "[UNKNOWN]"
+NA = "N/A"
 CONCEPT = "concept"
 CONCEPT_ANCESTOR = "concept_ancestor"
+MEASUREMENT_QUESTION_PREFIX = "1-Question:"
+MEASUREMENT_ANSWER_PREFIX = "2-Answer:"

cehrbert_data-0.0.4/src/cehrbert_data/decorators/__init__.py

@@ -0,0 +1,5 @@
+from .artificial_time_token_decorator import AttEventDecorator
+from .death_event_decorator import DeathEventDecorator
+from .clinical_event_decorator import ClinicalEventDecorator
+from .demographic_event_decorator import DemographicEventDecorator
+from .patient_event_decorator_base import time_token_func, get_att_function, AttType