cehrgpt 0.0.1__py3-none-any.whl

cehrgpt/omop/condition_era.py

@@ -0,0 +1,20 @@
+from cehrgpt.omop.omop_argparse import create_omop_argparse
+from cehrgpt.omop.omop_table_builder import OmopTableBuilder
+from cehrgpt.omop.queries.condition_era import CONDITION_ERA_QUERY
+
+CONDITION_ERA = "condition_era"
+
+
+def main(args):
+    OmopTableBuilder.create_omop_query_builder(
+        input_folder=args.input_folder,
+        output_folder=args.output_folder,
+        continue_job=args.continue_job,
+        table_name=CONDITION_ERA,
+        query_template=CONDITION_ERA_QUERY,
+        dependency_list=["condition_occurrence"],
+    ).build()
+
+
+if __name__ == "__main__":
+    main(create_omop_argparse().parse_args())

cehrgpt/omop/observation_period.py

@@ -0,0 +1,43 @@
+import glob
+import os.path
+
+from cehrgpt.omop.omop_argparse import create_omop_argparse
+from cehrgpt.omop.omop_table_builder import OmopTableBuilder
+from cehrgpt.omop.queries.observation_period import (
+    OBSERVATION_PERIOD_QUERY,
+    OBSERVATION_PERIOD_WITH_MEASUREMENT_QUERY,
+)
+
+OBSERVATION_PERIOD = "observation_period"
+
+
+def main(args):
+    include_measurement = (
+        len(glob.glob(os.path.join(args.input_folder, "measurement", "*.parquet"))) > 0
+    )
+    dependency_list = [
+        "person",
+        "visit_occurrence",
+        "condition_occurrence",
+        "procedure_occurrence",
+        "drug_exposure",
+    ]
+    if include_measurement:
+        dependency_list.append("measurement")
+
+    OmopTableBuilder.create_omop_query_builder(
+        input_folder=args.input_folder,
+        output_folder=args.output_folder,
+        continue_job=args.continue_job,
+        table_name=OBSERVATION_PERIOD,
+        query_template=(
+            OBSERVATION_PERIOD_WITH_MEASUREMENT_QUERY
+            if include_measurement
+            else OBSERVATION_PERIOD_QUERY
+        ),
+        dependency_list=dependency_list,
+    ).build()
+
+
+if __name__ == "__main__":
+    main(create_omop_argparse().parse_args())

cehrgpt/omop/omop_argparse.py

@@ -0,0 +1,38 @@
+import argparse
+
+from cehrbert_data.utils.spark_utils import validate_table_names
+
+
+def create_omop_argparse():
+    parser = argparse.ArgumentParser(
+        description="Spark application for creating OMOP table"
+    )
+    parser.add_argument(
+        "--input_folder",
+        dest="input_folder",
+        action="store",
+        help="The path for your input_folder where the raw data is",
+        required=True,
+    )
+    parser.add_argument(
+        "--output_folder",
+        dest="output_folder",
+        action="store",
+        help="The path for your output_folder",
+        required=True,
+    )
+    parser.add_argument(
+        "--continue_job",
+        dest="continue_job",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--domain_table_list",
+        dest="domain_table_list",
+        nargs="+",
+        action="store",
+        help="The list of domain tables you want to download",
+        type=validate_table_names,
+        required=True,
+    )
+    return parser

cehrgpt/omop/omop_table_builder.py

@@ -0,0 +1,86 @@
+from datetime import date
+from typing import Any, Dict, List, Optional
+
+import cehrbert_data.cohorts.spark_app_base
+from cehrbert_data.cohorts.query_builder import QueryBuilder, QuerySpec
+from cehrbert_data.cohorts.spark_app_base import BaseCohortBuilder
+from cehrbert_data.utils.spark_utils import preprocess_domain_table
+
+
+def custom_instantiate_dependencies(spark, input_folder, dependency_list):
+    dependency_dict = dict()
+    for domain_table_name in dependency_list:
+        table = preprocess_domain_table(spark, input_folder, domain_table_name)
+        table.createOrReplaceTempView(domain_table_name)
+        dependency_dict[domain_table_name] = table
+    return dependency_dict
+
+
+# Monkeypatch the function
+cehrbert_data.cohorts.spark_app_base.instantiate_dependencies = (
+    custom_instantiate_dependencies
+)
+
+
+class OmopTableBuilder(BaseCohortBuilder):
+    cohort_required_columns = []
+
+    def __init__(
+        self,
+        query_builder: QueryBuilder,
+        input_folder: str,
+        output_folder: str,
+        continue_job: bool = False,
+    ):
+        super().__init__(
+            query_builder=query_builder,
+            input_folder=input_folder,
+            output_folder=output_folder,
+            date_lower_bound="1900-01-01",
+            date_upper_bound=date.today().strftime("%Y-%m-%d"),
+            age_lower_bound=0,
+            age_upper_bound=200,
+            prior_observation_period=0,
+            post_observation_period=0,
+            continue_job=continue_job,
+        )
+
+    def build(self):
+        # Check whether the cohort has been generated
+        if self._continue_job and self.cohort_exists():
+            return self
+        cohort = self.create_cohort()
+        cohort.write.mode("overwrite").parquet(self._output_data_folder)
+        return self
+
+    @staticmethod
+    def create_omop_query_builder(
+        input_folder: str,
+        output_folder: str,
+        table_name: str,
+        query_template: str,
+        dependency_list: List[str],
+        query_parameters: Optional[Dict[str, Any]] = None,
+        continue_job: bool = False,
+    ):
+
+        if query_parameters is None:
+            query_parameters = dict()
+        query = QuerySpec(
+            table_name=table_name,
+            query_template=query_template,
+            parameters=query_parameters,
+        )
+        table_query_builder = QueryBuilder(
+            cohort_name=table_name,
+            dependency_list=dependency_list,
+            query=query,
+            ancestor_table_specs=[],
+        )
+
+        return OmopTableBuilder(
+            query_builder=table_query_builder,
+            input_folder=input_folder,
+            output_folder=output_folder,
+            continue_job=continue_job,
+        )

cehrgpt/omop/queries/condition_era.py

@@ -0,0 +1,86 @@
+CONDITION_ERA_QUERY = """
+-- Common Table Expression for cteConditionTarget
+WITH cteConditionTarget AS (
+  SELECT
+    co.CONDITION_OCCURRENCE_ID,
+    co.PERSON_ID,
+    co.CONDITION_CONCEPT_ID,
+    co.CONDITION_TYPE_CONCEPT_ID,
+    co.CONDITION_START_DATE,
+    COALESCE(co.CONDITION_END_DATE, DATE_ADD(co.CONDITION_START_DATE, 1)) AS CONDITION_END_DATE
+  FROM CONDITION_OCCURRENCE co
+),
+
+-- Common Table Expression for cteEndDates
+cteEndDates AS (
+  SELECT
+    PERSON_ID,
+    CONDITION_CONCEPT_ID,
+    DATE_ADD(EVENT_DATE, -30) AS END_DATE
+  FROM (
+    SELECT
+      PERSON_ID,
+      CONDITION_CONCEPT_ID,
+      EVENT_DATE,
+      EVENT_TYPE,
+      MAX(START_ORDINAL) OVER (PARTITION BY PERSON_ID, CONDITION_CONCEPT_ID ORDER BY EVENT_DATE, EVENT_TYPE ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS START_ORDINAL,
+      ROW_NUMBER() OVER (PARTITION BY PERSON_ID, CONDITION_CONCEPT_ID ORDER BY EVENT_DATE, EVENT_TYPE) AS OVERALL_ORD
+    FROM (
+      -- Select the start dates, assigning a row number to each
+      SELECT
+        PERSON_ID,
+        CONDITION_CONCEPT_ID,
+        CONDITION_START_DATE AS EVENT_DATE,
+        -1 AS EVENT_TYPE,
+        ROW_NUMBER() OVER (PARTITION BY PERSON_ID, CONDITION_CONCEPT_ID ORDER BY CONDITION_START_DATE) AS START_ORDINAL
+      FROM cteConditionTarget
+
+      UNION ALL
+
+      -- Pad the end dates by 30 to allow a grace period for overlapping ranges
+      SELECT
+        PERSON_ID,
+        CONDITION_CONCEPT_ID,
+        DATE_ADD(CONDITION_END_DATE, 30) AS EVENT_DATE,
+        1 AS EVENT_TYPE,
+        NULL
+      FROM cteConditionTarget
+    ) RAWDATA
+  ) E
+  WHERE (2 * E.START_ORDINAL) - E.OVERALL_ORD = 0
+),
+
+-- Common Table Expression for cteConditionEnds
+cteConditionEnds AS (
+  SELECT
+    c.PERSON_ID,
+    c.CONDITION_CONCEPT_ID,
+    c.CONDITION_TYPE_CONCEPT_ID,
+    c.CONDITION_START_DATE,
+    MIN(e.END_DATE) AS ERA_END_DATE
+  FROM cteConditionTarget c
+  JOIN cteEndDates e ON c.PERSON_ID = e.PERSON_ID AND c.CONDITION_CONCEPT_ID = e.CONDITION_CONCEPT_ID AND e.END_DATE >= c.CONDITION_START_DATE
+  GROUP BY
+    c.PERSON_ID,
+    c.CONDITION_CONCEPT_ID,
+    c.CONDITION_TYPE_CONCEPT_ID,
+    c.CONDITION_START_DATE
+)
+
+SELECT
+  ROW_NUMBER() OVER (ORDER BY CONDITION_CONCEPT_ID) AS condition_era_id,
+  PERSON_ID,
+  CONDITION_CONCEPT_ID,
+  MIN(CONDITION_START_DATE) AS CONDITION_ERA_START_DATE,
+  ERA_END_DATE AS CONDITION_ERA_END_DATE,
+  COUNT(*) AS CONDITION_OCCURRENCE_COUNT
+FROM cteConditionEnds
+GROUP BY
+  PERSON_ID,
+  CONDITION_CONCEPT_ID,
+  CONDITION_TYPE_CONCEPT_ID,
+  ERA_END_DATE
+ORDER BY
+  PERSON_ID,
+  CONDITION_CONCEPT_ID
+"""

cehrgpt/omop/queries/observation_period.py

@@ -0,0 +1,135 @@
+OBSERVATION_PERIOD_QUERY = """
+select
+    person_id,
+    observation_period_start_date,
+    case
+        when observation_period_end_date >= add_months(current_date(), -12) then current_date()
+        else observation_period_end_date
+    end as observation_period_end_date,
+    period_type_concept_id
+from (
+     SELECT person_id,
+            MIN(observation_period_start_date) AS observation_period_start_date,
+            MAX(observation_period_end_date)   AS observation_period_end_date,
+            44814725                           as period_type_concept_id
+     FROM (
+              SELECT pt.person_id             AS person_id,
+                     MIN(vt.visit_start_date) AS observation_period_start_date,
+                     MAX(vt.visit_end_date)   AS observation_period_end_date
+              FROM person as pt
+                       JOIN visit_occurrence as vt ON pt.person_id = vt.person_id
+              WHERE YEAR(vt.visit_start_date) >= 1985
+                AND vt.visit_start_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT pt.person_id                 AS person_id,
+                     MIN(co.condition_start_date) AS observation_period_start_date,
+                     MAX(co.condition_start_date) AS observation_period_end_date
+              FROM person as pt
+                       JOIN condition_occurrence as co ON pt.person_id = co.person_id
+              WHERE YEAR(co.condition_start_date) >= 1985
+                AND co.condition_start_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT pt.person_id           AS person_id,
+                     MIN(po.procedure_date) AS observation_period_start_date,
+                     MAX(po.procedure_date) AS observation_period_end_date
+              FROM person as pt
+                       JOIN procedure_occurrence as po ON pt.person_id = po.person_id
+              WHERE YEAR(po.procedure_date) >= 1985
+                AND po.procedure_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT pt.person_id                     AS person_id,
+                     MIN(de.drug_exposure_start_date) AS observation_period_start_date,
+                     MAX(de.drug_exposure_start_date) AS observation_period_end_date
+              FROM person as pt
+                       JOIN drug_exposure as de ON pt.person_id = de.person_id
+              WHERE YEAR(de.drug_exposure_start_date) >= 1985
+                AND de.drug_exposure_start_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+          ) as x
+     WHERE x.observation_period_end_date IS NOT NULL
+     GROUP BY person_id
+ ) as z
+"""
+
+OBSERVATION_PERIOD_WITH_MEASUREMENT_QUERY = """
+select
+    person_id,
+    observation_period_start_date,
+    case
+        when observation_period_end_date >= add_months(current_date(), -12) then current_date()
+        else observation_period_end_date
+    end as observation_period_end_date,
+    period_type_concept_id
+from (
+     SELECT person_id,
+            MIN(observation_period_start_date) AS observation_period_start_date,
+            MAX(observation_period_end_date)   AS observation_period_end_date,
+            44814725                           as period_type_concept_id
+     FROM (
+              SELECT pt.person_id             AS person_id,
+                     MIN(vt.visit_start_date) AS observation_period_start_date,
+                     MAX(vt.visit_end_date)   AS observation_period_end_date
+              FROM person as pt
+                       JOIN visit_occurrence as vt ON pt.person_id = vt.person_id
+              WHERE YEAR(vt.visit_start_date) >= 1985
+                AND vt.visit_start_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT pt.person_id                 AS person_id,
+                     MIN(co.condition_start_date) AS observation_period_start_date,
+                     MAX(co.condition_start_date) AS observation_period_end_date
+              FROM person as pt
+                       JOIN condition_occurrence as co ON pt.person_id = co.person_id
+              WHERE YEAR(co.condition_start_date) >= 1985
+                AND co.condition_start_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT pt.person_id           AS person_id,
+                     MIN(po.procedure_date) AS observation_period_start_date,
+                     MAX(po.procedure_date) AS observation_period_end_date
+              FROM person as pt
+                       JOIN procedure_occurrence as po ON pt.person_id = po.person_id
+              WHERE YEAR(po.procedure_date) >= 1985
+                AND po.procedure_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT pt.person_id                     AS person_id,
+                     MIN(de.drug_exposure_start_date) AS observation_period_start_date,
+                     MAX(de.drug_exposure_start_date) AS observation_period_end_date
+              FROM person as pt
+                       JOIN drug_exposure as de ON pt.person_id = de.person_id
+              WHERE YEAR(de.drug_exposure_start_date) >= 1985
+                AND de.drug_exposure_start_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+
+              UNION
+
+              SELECT
+                pt.person_id AS person_id,
+                MIN(m.measurement_date) AS observation_period_start_date,
+                MAX(m.measurement_date) AS observation_period_end_date
+              FROM person as pt
+                JOIN measurement as m ON pt.person_id = m.person_id
+              WHERE YEAR(m.measurement_date) >= 1985
+                AND m.measurement_date <= current_date() --set lower and upper bound to ignore spurious dates
+              GROUP BY pt.person_id
+          ) as x
+     WHERE x.observation_period_end_date IS NOT NULL
+     GROUP BY person_id
+ ) as z
+"""

cehrgpt/omop/sample_omop_tables.py

@@ -0,0 +1,71 @@
+import argparse
+import os
+
+from pyspark.sql import SparkSession
+
+# Define timestamp column for filtering based on the folder name
+omop_tables = [
+    "person",
+    "death",
+    "visit_occurrence",
+    "condition_occurrence",
+    "procedure_occurrence",
+    "drug_exposure",
+    "measurement",
+    "observation",
+    "observation_period",
+    "condition_era",
+    "drug_era",
+]
+
+
+# Main function to process the folders and upload tables
+def main(args):
+    spark = (
+        SparkSession.builder.appName("Sample OMOP Tables")
+        .config("spark.sql.legacy.parquet.int96RebaseModeInRead", "CORRECTED")
+        .config("spark.sql.legacy.parquet.int96RebaseModeInWrite", "CORRECTED")
+        .config("spark.sql.legacy.parquet.datetimeRebaseModeInRead", "CORRECTED")
+        .config("spark.sql.legacy.parquet.datetimeRebaseModeInWrite", "CORRECTED")
+        .getOrCreate()
+    )
+    patient_sample = spark.read.parquet(args.person_sample)
+    for omop_table in omop_tables:
+        omop_table_input = os.path.join(args.omop_folder, omop_table)
+        omop_table_output = os.path.join(args.output_folder, omop_table)
+        # If the input does not exist, let's skip
+        if not os.path.exists(omop_table_input):
+            continue
+        # If the output already exists, and overwrite is not set to True, let's skip it
+        if os.path.exists(omop_table_output) and not args.overwrite:
+            continue
+        omop_dataframe = spark.read.parquet(omop_table_input)
+        sub_omop_dataframe = omop_dataframe.join(
+            patient_sample.select("person_id"), "person_id"
+        )
+        sub_omop_dataframe.write.mode("overwrite" if args.overwrite else None).parquet(
+            omop_table_output
+        )
+
+
+# Argument parsing moved under __name__ == "__main__"
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Arguments for uploading OMOP tables")
+    parser.add_argument(
+        "--person_sample",
+        required=True,
+    )
+    parser.add_argument(
+        "--omop_folder",
+        required=True,
+    )
+    parser.add_argument(
+        "--output_folder",
+        required=True,
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+    )
+    # Call the main function with parsed arguments
+    main(parser.parse_args())

cehrgpt/runners/gpt_runner_util.py

@@ -0,0 +1,99 @@
+import dataclasses
+import os
+import sys
+from typing import Tuple
+
+from cehrbert.runners.hf_runner_argument_dataclass import (
+    DataTrainingArguments,
+    ModelArguments,
+)
+from transformers import HfArgumentParser, TrainingArguments
+from transformers.utils import logging
+from trl.trainer.dpo_config import DPOConfig
+
+from cehrgpt.runners.hf_gpt_runner_argument_dataclass import CehrGPTArguments
+
+LOG = logging.get_logger("transformers")
+
+
+def parse_dynamic_arguments(
+    argument_classes: Tuple[dataclasses.dataclass, ...] = (
+        DataTrainingArguments,
+        ModelArguments,
+        TrainingArguments,
+    )
+) -> Tuple:
+    """
+    Parses command-line arguments with extended flexibility, allowing for the inclusion of custom argument classes.
+
+    This function utilizes `HfArgumentParser` to parse arguments from command line input, JSON, or YAML files.
+    By default, it expects `ModelArguments`, `DataTrainingArguments`, and `TrainingArguments`, but it can be extended
+    with additional argument classes through the `argument_classes` parameter, making it suitable
+    for various custom setups.
+
+    Parameters:
+        argument_classes (Tuple[Type]): A tuple of argument classes to be parsed. Defaults to
+        `(ModelArguments, DataTrainingArguments, TrainingArguments)`. Additional argument classes can be specified
+        for greater flexibility in configuration.
+
+    Returns:
+        Tuple: A tuple of parsed arguments, one for each argument class provided. The order of the returned tuple
+        matches the order of the `argument_classes` parameter.
+
+    Raises:
+        FileNotFoundError: If the specified JSON or YAML file does not exist.
+        json.JSONDecodeError: If there is an error parsing a JSON file.
+        yaml.YAMLError: If there is an error parsing a YAML file.
+        Exception: For other issues that occur during argument parsing.
+
+    Example usage:
+        - Command-line: `python training_script.py --model_name_or_path bert-base-uncased --do_train`
+        - JSON file: `python training_script.py config.json`
+        - YAML file: `python training_script.py config.yaml`
+
+    Flexibility:
+        The function can be customized to include new argument classes as needed:
+
+        Example with a custom argument class:
+            ```python
+            class CustomArguments:
+                # Define custom arguments here
+                pass
+
+
+            custom_args = parse_extended_args(
+                (ModelArguments, DataTrainingArguments, TrainingArguments, CustomArguments)
+            )
+            ```
+        This example demonstrates how to include additional argument classes
+        beyond the defaults for a more tailored setup.
+    """
+    parser = HfArgumentParser(argument_classes)
+
+    # Check if input is a JSON or YAML file
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        args = parser.parse_json_file(json_file=os.path.expanduser(sys.argv[1]))
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        args = parser.parse_yaml_file(yaml_file=os.path.expanduser(sys.argv[1]))
+    else:
+        args = parser.parse_args_into_dataclasses()
+
+    return tuple(args)
+
+
+def parse_runner_args() -> (
+    Tuple[CehrGPTArguments, DataTrainingArguments, ModelArguments, TrainingArguments]
+):
+    cehrgpt_args, data_args, model_args, training_args = parse_dynamic_arguments(
+        (CehrGPTArguments, DataTrainingArguments, ModelArguments, TrainingArguments)
+    )
+    return cehrgpt_args, data_args, model_args, training_args
+
+
+def parse_dpo_runner_args() -> (
+    Tuple[CehrGPTArguments, DataTrainingArguments, ModelArguments, DPOConfig]
+):
+    cehrgpt_args, data_args, model_args, dpo_config = parse_dynamic_arguments(
+        (CehrGPTArguments, DataTrainingArguments, ModelArguments, DPOConfig)
+    )
+    return cehrgpt_args, data_args, model_args, dpo_config