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

{cehrbert_data-0.0.4 → cehrbert_data-0.0.5}/.github/workflows/python-build.yml

@@ -68,7 +68,7 @@ jobs:
         name: python-package-distributions
         path: dist/
     - name: Sign the dists with Sigstore
-      uses: sigstore/gh-action-sigstore-python@v2.1.1
+      uses: sigstore/gh-action-sigstore-python@v3.0.0
       with:
         inputs: >-
           ./dist/*.tar.gz
@@ -90,4 +90,4 @@ jobs:
       run: >-
         gh release upload
         '${{ github.ref_name }}' dist/**
-        --repo '${{ github.repository }}'
+        --repo '${{ github.repository }}'

{cehrbert_data-0.0.4 → cehrbert_data-0.0.5}/.github/workflows/tests.yml

@@ -23,6 +23,11 @@ jobs:
         uses: actions/setup-python@v3
         with:
           python-version: "3.10"
+      - name: Set up Java 11
+        uses: actions/setup-java@v3
+        with:
+          java-version: "11" # specify the Java version here
+          distribution: "temurin" # or use 'adopt' or 'zulu', depending on your preference
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

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

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

@@ -41,13 +41,15 @@ def main(
         apply_age_filter: bool,
         include_death: bool,
         att_type: AttType,
+        inpatient_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
+        refresh_measurement: bool = False,
+        aggregate_by_hour: bool = True,
 ):
     spark = SparkSession.builder.appName("Generate CEHR-BERT Training Data").getOrCreate()
 
@@ -71,6 +73,7 @@ def main(
         f"use_age_group: {use_age_group}\n"
         f"with_drug_rollup: {with_drug_rollup}\n"
         f"refresh_measurement: {refresh_measurement}\n"
+        f"aggregate_by_hour: {aggregate_by_hour}\n"
     )
 
     domain_tables = []
@@ -128,7 +131,8 @@ def main(
         processed_measurement = get_measurement_table(
             spark,
             input_folder,
-            refresh=refresh_measurement
+            refresh=refresh_measurement,
+            aggregate_by_hour=aggregate_by_hour,
         )
         if patient_events:
             # Union all measurement records together with other domain records
@@ -165,9 +169,12 @@ def main(
             patient_demographic=person if gpt_patient_sequence else None,
             death=death,
             att_type=att_type,
+            inpatient_att_type=inpatient_att_type,
             exclude_demographic=exclude_demographic,
             use_age_group=use_age_group,
             include_inpatient_hour_token=include_inpatient_hour_token,
+            spark=spark,
+            persistence_folder=output_folder,
         )
     else:
         sequence_data = create_sequence_data(
@@ -325,12 +332,23 @@ if __name__ == "__main__":
         dest="refresh_measurement",
         action="store_true"
     )
+    parser.add_argument(
+        "--aggregate_by_hour",
+        dest="aggregate_by_hour",
+        action="store_true"
+    )
     parser.add_argument(
         "--att_type",
         dest="att_type",
         action="store",
         choices=[e.value for e in AttType],
     )
+    parser.add_argument(
+        "--inpatient_att_type",
+        dest="inpatient_att_type",
+        action="store",
+        choices=[e.value for e in AttType],
+    )
 
     ARGS = parser.parse_args()
 
@@ -352,10 +370,12 @@ if __name__ == "__main__":
         ARGS.apply_age_filter,
         ARGS.include_death,
         AttType(ARGS.att_type),
+        AttType(ARGS.inpatient_att_type),
         exclude_demographic=ARGS.exclude_demographic,
         use_age_group=ARGS.use_age_group,
         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,
+        aggregate_by_hour=ARGS.aggregate_by_hour,
     )

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

@@ -295,6 +295,7 @@ class NestedCohortBuilder:
             is_drug_roll_up_concept: bool = True,
             include_concept_list: bool = True,
             refresh_measurement: bool = False,
+            aggregate_by_hour: bool = True,
             is_new_patient_representation: bool = False,
             gpt_patient_sequence: bool = False,
             is_hierarchical_bert: bool = False,
@@ -306,9 +307,12 @@ class NestedCohortBuilder:
             is_observation_window_unbounded: bool = False,
             is_population_estimation: bool = False,
             att_type: AttType = AttType.CEHR_BERT,
+            inpatient_att_type: AttType = AttType.MIX,
             exclude_demographic: bool = True,
             use_age_group: bool = False,
             single_contribution: bool = False,
+            exclude_features: bool = True,
+            meds_format: bool = False,
     ):
         self._cohort_name = cohort_name
         self._input_folder = input_folder
@@ -340,15 +344,19 @@ class NestedCohortBuilder:
         self._is_prediction_window_unbounded = is_prediction_window_unbounded
         self._include_concept_list = include_concept_list
         self._refresh_measurement = refresh_measurement
+        self._aggregate_by_hour = aggregate_by_hour
         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())
         )
         self._is_population_estimation = is_population_estimation
         self._att_type = att_type
+        self._inpatient_att_type = inpatient_att_type
         self._exclude_demographic = exclude_demographic
         self._use_age_group = use_age_group
         self._single_contribution = single_contribution
+        self._exclude_features = exclude_features
+        self._meds_format = meds_format
 
         self.get_logger().info(
             f"cohort_name: {cohort_name}\n"
@@ -377,12 +385,16 @@ class NestedCohortBuilder:
             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"aggregate_by_hour: {aggregate_by_hour}\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"
+            f"inpatient_att_type: {inpatient_att_type}\n"
             f"exclude_demographic: {exclude_demographic}\n"
             f"use_age_group: {use_age_group}\n"
             f"single_contribution: {single_contribution}\n"
+            f"extract_features: {exclude_features}\n"
+            f"meds_format: {meds_format}\n"
         )
 
         self.spark = SparkSession.builder.appName(f"Generate {self._cohort_name}").getOrCreate()
@@ -508,13 +520,30 @@ class NestedCohortBuilder:
             )
             cohort = cohort.withColumn("record_rank", record_rank).where("record_rank == 1").drop("record_rank")
 
-        ehr_records_for_cohorts = self.extract_ehr_records_for_cohort(cohort)
-        # ehr_records_for_cohorts.show()
-        cohort = (
-            cohort.join(ehr_records_for_cohorts, ["person_id", "cohort_member_id"])
-            .where(F.col("num_of_visits") >= self._num_of_visits)
-            .where(F.col("num_of_concepts") >= self._num_of_concepts)
-        )
+        if not self._exclude_features:
+            ehr_records_for_cohorts = self.extract_ehr_records_for_cohort(cohort)
+            cohort = (
+                cohort.join(ehr_records_for_cohorts, ["person_id", "cohort_member_id"])
+                .where(F.col("num_of_visits") >= self._num_of_visits)
+                .where(F.col("num_of_concepts") >= self._num_of_concepts)
+            )
+
+        person_id_column = "person_id"
+        index_date_column = "index_date"
+        if self._meds_format:
+            cohort = cohort.withColumnRenamed(
+                "person_id", "subject_id"
+            ).withColumnRenamed(
+                "index_date", "prediction_time"
+            ).withColumnRenamed(
+                "label", "boolean_value"
+            ).withColumn(
+                "prediction_time", F.to_timestamp("prediction_time")
+            ).withColumn(
+                "boolean_value", F.col("boolean_value").cast("boolean")
+            )
+            person_id_column = "subject_id"
+            index_date_column = "prediction_time"
 
         if self._is_prediction_window_unbounded:
             observation_period = self._dependency_dict[OBSERVATION_PERIOD]
@@ -529,9 +558,14 @@ class NestedCohortBuilder:
             # 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))
+                F.coalesce(
+                    F.col("outcome_date"),
+                    F.date_add(
+                        cohort[index_date_column], self._prediction_window
+                    )
+                )
             )
-        cohort = cohort.withColumn("time_to_event", F.datediff("study_end_date", "index_date"))
+        cohort = cohort.withColumn("time_to_event", F.datediff("study_end_date", index_date_column))
 
         # if patient_splits is provided, we will
         if self._patient_splits_folder:
@@ -547,7 +581,7 @@ class NestedCohortBuilder:
             cohort.where('split="test"').write.mode("overwrite").parquet(os.path.join(self._output_data_folder, "test"))
             shutil.rmtree(os.path.join(self._output_data_folder, "temp"))
         else:
-            cohort.orderBy("person_id", "cohort_member_id").write.mode("overwrite").parquet(self._output_data_folder)
+            cohort.orderBy(person_id_column, index_date_column).write.mode("overwrite").parquet(self._output_data_folder)
 
     def extract_ehr_records_for_cohort(self, cohort: DataFrame):
         """
@@ -566,47 +600,54 @@ class NestedCohortBuilder:
             with_drug_rollup=self._is_drug_roll_up_concept,
             include_concept_list=self._include_concept_list,
             refresh_measurement=self._refresh_measurement,
+            aggregate_by_hour=self._aggregate_by_hour,
         )
 
         # Duplicate the records for cohorts that allow multiple entries
-        ehr_records = ehr_records.join(cohort, "person_id").select(
-            [ehr_records[field_name] for field_name in ehr_records.schema.fieldNames()] + ["cohort_member_id"]
-        )
+        ehr_records = ehr_records.alias("ehr").join(
+            cohort.alias("cohort"), F.col("ehr.person_id") == F.col("cohort.person_id")
+        ).select(
+            [F.col("ehr." + col) for col in ehr_records.columns] + [F.col("cohort.cohort_member_id")]
+        ).selectExpr("*")
 
         # Only allow the data records that occurred between the index date and the prediction window
         if self._is_population_estimation:
             if self._is_prediction_window_unbounded:
-                record_window_filter = ehr_records["date"] <= F.current_date()
+                record_window_filter = F.col("ehr.date") <= F.current_date()
             else:
-                record_window_filter = ehr_records["date"] <= F.date_add(cohort["index_date"], self._prediction_window)
+                record_window_filter = F.col("ehr.date") <= F.date_add(
+                    F.col("cohort.index_date"), self._prediction_window
+                )
         else:
             # For patient level prediction, we remove all records post index date
             if self._is_observation_post_index:
-                record_window_filter = ehr_records["date"].between(
-                    cohort["index_date"],
-                    F.date_add(cohort["index_date"], self._observation_window),
+                record_window_filter = F.col("ehr.date").between(
+                    F.col("cohort.index_date"),
+                    F.date_add(F.col("cohort.index_date"), self._observation_window),
                 )
             else:
                 if self._is_observation_window_unbounded:
-                    record_window_filter = ehr_records["date"] <= F.date_sub(
-                        cohort["index_date"], self._hold_off_window
+                    record_window_filter = F.col("ehr.date") <= F.date_sub(
+                        F.col("cohort.index_date"), self._hold_off_window
                     )
                 else:
-                    record_window_filter = ehr_records["date"].between(
+                    record_window_filter = F.col("ehr.date").between(
                         F.date_sub(
-                            cohort["index_date"],
+                            F.col("cohort.index_date"),
                             self._observation_window + self._hold_off_window,
                         ),
-                        F.date_sub(cohort["index_date"], self._hold_off_window),
+                        F.date_sub(F.col("cohort.index_date"), self._hold_off_window),
                     )
 
-        cohort_ehr_records = (
-            ehr_records.join(
-                cohort,
-                (ehr_records.person_id == cohort.person_id) & (ehr_records.cohort_member_id == cohort.cohort_member_id),
-            )
-            .where(record_window_filter)
-            .select([ehr_records[field_name] for field_name in ehr_records.schema.fieldNames()])
+        # Somehow the dataframe join does not work without using the alias
+        cohort_ehr_records = ehr_records.alias("ehr").join(
+            cohort.alias("cohort"),
+            (F.col("ehr.person_id") == F.col("cohort.person_id")) &
+            (F.col("ehr.cohort_member_id") == F.col("cohort.cohort_member_id")),
+        ).where(
+            record_window_filter
+        ).select(
+            [F.col("ehr." + field_name) for field_name in ehr_records.schema.fieldNames()]
         )
 
         if self._is_hierarchical_bert:
@@ -647,6 +688,7 @@ class NestedCohortBuilder:
                 exclude_visit_tokens=self._exclude_visit_tokens,
                 patient_demographic=(patient_demographic if self._gpt_patient_sequence else None),
                 att_type=self._att_type,
+                inpatient_att_type=self._inpatient_att_type,
                 exclude_demographic=self._exclude_demographic,
                 use_age_group=self._use_age_group,
             )
@@ -748,6 +790,7 @@ def create_prediction_cohort(
         is_drug_roll_up_concept=spark_args.is_drug_roll_up_concept,
         include_concept_list=spark_args.include_concept_list,
         refresh_measurement=spark_args.refresh_measurement,
+        aggregate_by_hour=spark_args.aggregate_by_hour,
         is_new_patient_representation=spark_args.is_new_patient_representation,
         gpt_patient_sequence=spark_args.gpt_patient_sequence,
         is_hierarchical_bert=spark_args.is_hierarchical_bert,
@@ -763,7 +806,10 @@ def create_prediction_cohort(
         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),
+        inpatient_att_type=AttType(spark_args.inpatient_att_type),
         exclude_demographic=spark_args.exclude_demographic,
         use_age_group=spark_args.use_age_group,
         single_contribution=spark_args.single_contribution,
+        exclude_features=spark_args.exclude_features,
+        meds_format=spark_args.meds_format,
     ).build()

{cehrbert_data-0.0.4 → cehrbert_data-0.0.5}/src/cehrbert_data/decorators/artificial_time_token_decorator.py

@@ -1,14 +1,11 @@
-from pyspark.sql import DataFrame, functions as F, types as T, Window as W
+import os.path
+
+from pyspark.sql import SparkSession, DataFrame, functions as F, types as T, Window as W
 
 from ..const.common import NA
 from ..const.artificial_tokens import VS_TOKEN, VE_TOKEN
 from .patient_event_decorator_base import (
-    PatientEventDecorator, AttType,
-    time_day_token,
-    time_week_token,
-    time_month_token,
-    time_mix_token,
-    time_token_func
+    PatientEventDecorator, AttType, get_att_function
 )
 from .token_priority import (
     ATT_TOKEN_PRIORITY,
@@ -16,6 +13,7 @@ from .token_priority import (
     VISIT_TYPE_TOKEN_PRIORITY,
     DISCHARGE_TOKEN_PRIORITY,
     VE_TOKEN_PRIORITY,
+    FIRST_VISIT_HOUR_TOKEN_PRIORITY,
     get_inpatient_token_priority,
     get_inpatient_att_token_priority
 )
@@ -28,20 +26,28 @@ class AttEventDecorator(PatientEventDecorator):
             include_visit_type,
             exclude_visit_tokens,
             att_type: AttType,
+            inpatient_att_type: AttType,
             include_inpatient_hour_token: bool = False,
+            spark: SparkSession = None,
+            persistence_folder: str = None,
     ):
         self._visit_occurrence = visit_occurrence
         self._include_visit_type = include_visit_type
         self._exclude_visit_tokens = exclude_visit_tokens
         self._att_type = att_type
+        self._inpatient_att_type = inpatient_att_type
         self._include_inpatient_hour_token = include_inpatient_hour_token
+        super().__init__(spark=spark, persistence_folder=persistence_folder)
+
+    def get_name(self):
+        return "att_events"
 
     def _decorate(self, patient_events: DataFrame):
         if self._att_type == AttType.NONE:
             return patient_events
 
         # visits should the following columns (person_id,
-        # visit_concept_id, visit_start_date, visit_occurrence_id, domain, concept_value)
+        # visit_concept_id, visit_start_date, visit_occurrence_id, domain)
         cohort_member_person_pair = patient_events.select("person_id", "cohort_member_id").distinct()
         valid_visit_ids = patient_events.groupby(
             "cohort_member_id",
@@ -55,9 +61,21 @@ class AttEventDecorator(PatientEventDecorator):
             F.max("concept_order").alias("max_concept_order"),
         )
 
+        # The visit records are joined to the cohort members (there could be multiple entries for the same patient)
+        # if multiple entries are present, we duplicate the visit records for those. If the visit_occurrence dataframe
+        # contains visits for each cohort member, then we need to add cohort_member_id to the joined expression as well.
+        if "cohort_member_id" in self._visit_occurrence.columns:
+            joined_expr = ["person_id", "cohort_member_id"]
+        else:
+            joined_expr = ["person_id"]
+
         visit_occurrence = (
-            self._visit_occurrence.select(
+            self._visit_occurrence.join(
+                cohort_member_person_pair,
+                joined_expr
+            ).select(
                 "person_id",
+                "cohort_member_id",
                 F.col("visit_start_date").cast(T.DateType()).alias("date"),
                 F.col("visit_start_date").cast(T.DateType()).alias("visit_start_date"),
                 F.col("visit_start_datetime").cast(T.TimestampType()).alias("visit_start_datetime"),
@@ -65,14 +83,18 @@ class AttEventDecorator(PatientEventDecorator):
                 "visit_concept_id",
                 "visit_occurrence_id",
                 F.lit("visit").alias("domain"),
-                F.lit(0.0).alias("concept_value"),
+                F.lit(None).cast("float").alias("number_as_value"),
+                F.lit(None).cast("string").alias("concept_as_value"),
+                F.lit(0).alias("is_numeric_type"),
                 F.lit(0).alias("concept_value_mask"),
                 F.lit(0).alias("mlm_skip_value"),
                 "age",
                 "discharged_to_concept_id",
             )
-            .join(valid_visit_ids, "visit_occurrence_id")
-            .join(cohort_member_person_pair, ["person_id", "cohort_member_id"])
+            .join(
+                valid_visit_ids,
+                ["visit_occurrence_id", "cohort_member_id"]
+            )
         )
 
         # We assume outpatient visits end on the same day, therefore we start visit_end_date to visit_start_date due
@@ -89,7 +111,10 @@ class AttEventDecorator(PatientEventDecorator):
         visit_occurrence = visit_occurrence.withColumn("date_in_week", weeks_since_epoch_udf)
 
         # Cache visit for faster processing
-        visit_occurrence.cache()
+        visit_occurrence = self.try_persist_data(
+            visit_occurrence,
+            os.path.join(self.get_name(), "visit_occurrence_temp"),
+        )
 
         visits = visit_occurrence.drop("discharged_to_concept_id")
 
@@ -134,18 +159,7 @@ class AttEventDecorator(PatientEventDecorator):
         )
 
         # Udf for calculating the time token
-        if self._att_type == AttType.DAY:
-            att_func = time_day_token
-        elif self._att_type == AttType.WEEK:
-            att_func = time_week_token
-        elif self._att_type == AttType.MONTH:
-            att_func = time_month_token
-        elif self._att_type == AttType.MIX:
-            att_func = time_mix_token
-        else:
-            att_func = time_token_func
-
-        time_token_udf = F.udf(att_func, T.StringType())
+        time_token_udf = F.udf(get_att_function(self._att_type), T.StringType())
 
         att_tokens = (
             visits.withColumn("datetime", F.to_timestamp("date"))
@@ -191,6 +205,12 @@ class AttEventDecorator(PatientEventDecorator):
 
         artificial_tokens = artificial_tokens.drop("visit_end_date")
 
+        # Try persisting artificial events
+        artificial_tokens = self.try_persist_data(
+            artificial_tokens,
+            os.path.join(self.get_name(), "artificial_tokens"),
+        )
+
         # Retrieving the events that are ONLY linked to inpatient visits
         inpatient_visits = (
             visit_occurrence
@@ -201,6 +221,7 @@ class AttEventDecorator(PatientEventDecorator):
             inpatient_visits, ["visit_occurrence_id", "cohort_member_id"]
         )
 
+        inpatient_time_token_udf = F.udf(get_att_function(self._inpatient_att_type), T.StringType())
         # Fill in the visit_end_date if null (because some visits are still ongoing at the time of data extraction)
         # Bound the event dates within visit_start_date and visit_end_date
         # Generate a span rank to indicate the position of the group of events
@@ -245,6 +266,11 @@ class AttEventDecorator(PatientEventDecorator):
         # Add discharge events to the inpatient visits
         inpatient_events = inpatient_events.unionByName(discharge_events)
 
+        # Try persisting the inpatient events for fasting processing
+        inpatient_events = self.try_persist_data(
+            inpatient_events, os.path.join(self.get_name(), "inpatient_events")
+        )
+
         # Get the prev days_since_epoch
         inpatient_prev_date_udf = F.lag("date").over(
             W.partitionBy("cohort_member_id", "visit_occurrence_id").orderBy("concept_order")
@@ -264,16 +290,42 @@ class AttEventDecorator(PatientEventDecorator):
             )
             inpatient_att_token = F.when(
                 F.col("hour_delta") < 24, F.concat(F.lit("i-H"), F.col("hour_delta"))
-            ).otherwise(F.concat(F.lit("i-"), time_token_udf("time_delta")))
+            ).otherwise(F.concat(F.lit("i-"), inpatient_time_token_udf("time_delta")))
+
+            # We need to insert an ATT token between midnight and the visit start datetime
+            first_inpatient_hour_delta_udf = (
+                F.floor((F.unix_timestamp("visit_start_datetime") - F.unix_timestamp(
+                    F.col("visit_start_datetime").cast("date"))) / 3600)
+            )
+
+            first_hour_tokens = (
+                visits.where(F.col("visit_concept_id").isin([9201, 262, 8971, 8920]))
+                .withColumn("hour_delta", first_inpatient_hour_delta_udf)
+                .where(F.col("hour_delta") > 0)
+                .withColumn("date", F.col("visit_start_date"))
+                .withColumn("datetime", F.to_timestamp("date"))
+                .withColumn("standard_concept_id", F.concat(F.lit("i-H"), F.col("hour_delta")))
+                .withColumn("visit_concept_order", F.col("min_visit_concept_order"))
+                .withColumn("concept_order", F.lit(0))
+                .withColumn("priority", F.lit(FIRST_VISIT_HOUR_TOKEN_PRIORITY))
+                .withColumn("unit", F.lit(NA))
+                .withColumn("event_group_id", F.lit(NA))
+                .drop("min_visit_concept_order", "max_visit_concept_order")
+                .drop("min_concept_order", "max_concept_order")
+                .drop("hour_delta", "visit_end_date")
+            )
+
             # Create ATT tokens within the inpatient visits
             inpatient_att_events = (
                 inpatient_events.withColumn(
+                    "time_stamp_hour", F.hour("datetime")
+                ).withColumn(
                     "is_span_boundary",
                     F.row_number().over(
-                        W.partitionBy("cohort_member_id", "visit_occurrence_id", "concept_order").orderBy("priority")
+                        W.partitionBy("cohort_member_id", "visit_occurrence_id")
+                        .orderBy("priority", "date", "time_stamp_hour")
                     ),
                 )
-                .where(F.col("is_span_boundary") == 1)
                 .withColumn("prev_date", inpatient_prev_date_udf)
                 .withColumn("time_delta", inpatient_time_delta_udf)
                 .withColumn("prev_datetime", inpatient_prev_datetime_udf)
@@ -284,12 +336,17 @@ class AttEventDecorator(PatientEventDecorator):
                 .withColumn("visit_concept_order", F.col("visit_concept_order"))
                 .withColumn("priority", get_inpatient_att_token_priority())
                 .withColumn("concept_value_mask", F.lit(0))
-                .withColumn("concept_value", F.lit(0.0))
+                .withColumn("number_as_value", F.lit(None).cast("float"))
+                .withColumn("concept_as_value", F.lit(None).cast("string"))
+                .withColumn("is_numeric_type", F.lit(0))
                 .withColumn("unit", F.lit(NA))
                 .withColumn("event_group_id", F.lit(NA))
                 .drop("prev_date", "time_delta", "is_span_boundary")
-                .drop("prev_datetime", "hour_delta")
+                .drop("prev_datetime", "hour_delta", "time_stamp_hour")
             )
+
+            # Insert the first hour tokens between the visit type and first medical event
+            inpatient_att_events = inpatient_att_events.unionByName(first_hour_tokens)
         else:
             # Create ATT tokens within the inpatient visits
             inpatient_att_events = (
@@ -306,17 +363,24 @@ class AttEventDecorator(PatientEventDecorator):
                 .where(F.col("prev_date").isNotNull())
                 .withColumn(
                     "standard_concept_id",
-                    F.concat(F.lit("i-"), time_token_udf("time_delta")),
+                    F.concat(F.lit("i-"), inpatient_time_token_udf("time_delta")),
                 )
                 .withColumn("visit_concept_order", F.col("visit_concept_order"))
                 .withColumn("priority", get_inpatient_att_token_priority())
                 .withColumn("concept_value_mask", F.lit(0))
-                .withColumn("concept_value", F.lit(0.0))
+                .withColumn("number_as_value", F.lit(None).cast("float"))
+                .withColumn("concept_as_value", F.lit(None).cast("string"))
+                .withColumn("is_numeric_type", F.lit(0))
                 .withColumn("unit", F.lit(NA))
                 .withColumn("event_group_id", F.lit(NA))
                 .drop("prev_date", "time_delta", "is_span_boundary")
             )
 
+        # Try persisting the inpatient att events
+        inpatient_att_events = self.try_persist_data(
+            inpatient_att_events, os.path.join(self.get_name(), "inpatient_att_events")
+        )
+
         self.validate(inpatient_events)
         self.validate(inpatient_att_events)
 
@@ -326,6 +390,10 @@ class AttEventDecorator(PatientEventDecorator):
             ["visit_occurrence_id", "cohort_member_id"],
             how="left_anti",
         )
+        # Try persisting the other events
+        other_events = self.try_persist_data(
+            other_events, os.path.join(self.get_name(), "other_events")
+        )
 
         patient_events = inpatient_events.unionByName(inpatient_att_events).unionByName(other_events)