cehrgpt 0.1.0__tar.gz → 0.1.2__tar.gz

{cehrgpt-0.1.0 → cehrgpt-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cehrgpt
-Version: 0.1.0
+Version: 0.1.2
 Summary: CEHR-GPT: Generating Electronic Health Records with Chronological Patient Timelines
 Author-email: Chao Pang <chaopang229@gmail.com>, Xinzhuo Jiang <xj2193@cumc.columbia.edu>, Krishna Kalluri <kk3326@cumc.columbia.edu>, Elise Minto <em3697@cumc.columbia.edu>, Jason Patterson <jp3477@cumc.columbia.edu>, Nishanth Parameshwar Pavinkurve <np2689@cumc.columbia.edu>, Karthik Natarajan <kn2174@cumc.columbia.edu>
 License: MIT License
@@ -12,14 +12,15 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: cehrbert==1.4.1
-Requires-Dist: cehrbert_data==0.0.7
+Requires-Dist: cehrbert==1.4.5
+Requires-Dist: cehrbert_data==0.0.11
 Requires-Dist: openai==1.54.3
 Requires-Dist: optuna==4.0.0
-Requires-Dist: transformers==4.44.0
+Requires-Dist: transformers==4.44.1
 Requires-Dist: tokenizers==0.19.0
 Requires-Dist: peft==0.10.0
 Requires-Dist: lightgbm
+Requires-Dist: polars
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
@@ -36,9 +37,9 @@ Dynamic: license-file
 
 [![PyPI - Version](https://img.shields.io/pypi/v/cehrgpt)](https://pypi.org/project/cehrgpt/)
 ![Python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)
-[![tests](https://github.com/knatarajan-lab/cehrgpt-public/actions/workflows/tests.yaml/badge.svg)](https://github.com/knatarajan-lab/cehrgpt-public/actions/workflows/tests.yml)
-[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/knatarajan-lab/cehrgpt-public/blob/main/LICENSE)
-[![contributors](https://img.shields.io/github/contributors/knatarajan-lab/cehrgpt-public.svg)](https://github.com/knatarajan-lab/cehrgpt-public/graphs/contributors)
+[![tests](https://github.com/knatarajan-lab/cehrgpt/actions/workflows/tests.yaml/badge.svg)](https://github.com/knatarajan-lab/cehrgpt/actions/workflows/tests.yaml)
+[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/knatarajan-lab/cehrgpt/blob/main/LICENSE)
+[![contributors](https://img.shields.io/github/contributors/knatarajan-lab/cehrgpt.svg)](https://github.com/knatarajan-lab/cehrgpt/graphs/contributors)
 
 ## Description
 CEHRGPT is a synthetic data generation model developed to handle structured electronic health records (EHR) with enhanced privacy and reliability. It leverages state-of-the-art natural language processing techniques to create realistic, anonymized patient data that can be used for research and development without compromising patient privacy.
@@ -104,6 +105,100 @@ sh scripts/omop_pipeline.sh \
   $OMOP_VOCAB_DIR
 ```
 
+# MEDS Support
+
+This section demonstrates how to pretrain CEHR-GPT using MIMIC-IV data in the MEDS (Medical Event Data Standard) format.
+
+## Prerequisites
+
+Set up the required environment variables before beginning:
+
+```bash
+export CEHR_GPT_MODEL_DIR=""    # Path to CEHR-GPT model directory
+export MEDS_DIR=""              # Path to MEDS data directory
+export MEDS_READER_DIR=""       # Path to MEDS reader output directory
+```
+
+## Step 1: Create MIMIC MEDS Data
+
+Transform your MIMIC files into MEDS format by following the instructions in the [MEDS_transforms](https://github.com/mmcdermott/MEDS_transforms/) repository.
+
+## Step 2: Create the MEDS Reader
+
+Convert the MEDS data for use with CEHR-GPT:
+
+```bash
+meds_reader_convert $MEDS_DIR $MEDS_READER_DIR --num_threads 10
+```
+
+## Step 3: Pretrain CEHR-GPT
+
+Run the pretraining process using the prepared MEDS data:
+
+```bash
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $CEHR_GPT_MODEL_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --do_train true --seed 42 \
+  --dataloader_num_workers 16 --dataloader_prefetch_factor 8 \
+  --hidden_size 768 --num_hidden_layers 14 --max_position_embeddings 8192 \
+  --evaluation_strategy epoch --save_strategy epoch \
+  --sample_packing --max_tokens_per_batch 16384 \
+  --warmup_steps 500 --weight_decay 0.01 \
+  --num_train_epochs 50 --learning_rate 0.0002 \
+  --use_early_stopping --early_stopping_threshold 0.001 \
+  --is_data_in_meds --inpatient_att_function_type day \
+  --att_function_type day --include_inpatient_hour_token \
+  --include_auxiliary_token --include_demographic_prompt \
+  --meds_to_cehrbert_conversion_type "MedsToBertMimic4"
+```
+
+## Step 4: Generate MEDS Trajectories
+
+### Environment Setup for Trajectory Generation
+
+Configure additional environment variables for trajectory generation with task labels (`subject_id`, `prediction_time`, `boolean_value` [optional]):
+
+```bash
+# MEDS_LABEL_COHORT_DIR must contain a set of parquet files
+export MEDS_LABEL_COHORT_DIR=""     # Path to cohort labels directory
+export MEDS_TRAJECTORY_DIR=""       # Path for trajectory output
+```
+
+### Generate Trajectories
+
+Create synthetic patient trajectories using the trained model:
+
+> **Important:** The total sequence length (`generation_input_length` + `generation_max_new_tokens`) cannot exceed the `max_position_embeddings` value (8192) defined during pretraining.
+
+```bash
+python -u -m cehrgpt.generation.cehrgpt_conditional_generation \
+  --cohort_folder $MEDS_LABEL_COHORT_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $MEDS_TRAJECTORY_DIR \
+  --per_device_eval_batch_size 16 \
+  --num_of_trajectories_per_sample 2 \
+  --generation_input_length 4096 \
+  --generation_max_new_tokens 4096 \
+  --is_data_in_meds \
+  --att_function_type day --inpatient_att_function_type day \
+  --meds_to_cehrbert_conversion_type MedsToBertMimic4 \
+  --include_auxiliary_token --include_demographic_prompt \
+  --include_inpatient_hour_token
+```
+
+### Parameters Explanation
+
+- `generation_input_length`: Controls the length of input context for generation
+- `generation_max_new_tokens`: Maximum number of new tokens to generate
+- `num_of_trajectories_per_sample`: Number of trajectories to generate per patient sample
+
 ## Citation
 ```
 @article{cehrgpt2024,

cehrgpt-0.1.2/README.md

@@ -0,0 +1,174 @@
+# CEHRGPT
+
+[![PyPI - Version](https://img.shields.io/pypi/v/cehrgpt)](https://pypi.org/project/cehrgpt/)
+![Python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)
+[![tests](https://github.com/knatarajan-lab/cehrgpt/actions/workflows/tests.yaml/badge.svg)](https://github.com/knatarajan-lab/cehrgpt/actions/workflows/tests.yaml)
+[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/knatarajan-lab/cehrgpt/blob/main/LICENSE)
+[![contributors](https://img.shields.io/github/contributors/knatarajan-lab/cehrgpt.svg)](https://github.com/knatarajan-lab/cehrgpt/graphs/contributors)
+
+## Description
+CEHRGPT is a synthetic data generation model developed to handle structured electronic health records (EHR) with enhanced privacy and reliability. It leverages state-of-the-art natural language processing techniques to create realistic, anonymized patient data that can be used for research and development without compromising patient privacy.
+
+## Features
+- **Synthetic Patient Data Generation**: Generates comprehensive patient profiles including demographics, medical history, treatment courses, and outcomes.
+- **Privacy-Preserving**: Implements techniques to ensure the generated data does not reveal identifiable information.
+- **Compatibility with OMOP**: Fully compatible with the OMOP common data model, allowing seamless integration with existing healthcare data systems.
+- **Extensible**: Designed to be adaptable to new datasets and different EHR systems.
+
+## Installation
+To install CEHRGPT, clone this repository and install the required dependencies.
+
+```bash
+git clone https://github.com/knatarajan-lab/cehrgpt.git
+cd cehrgpt
+pip install .
+```
+
+## Pretrain
+Pretrain cehrgpt using the Hugging Face trainer, the parameters can be found in the sample configuration yaml
+```bash
+mkdir test_results
+# This is NOT required when streaming is set to true
+mkdir test_dataset_prepared
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner sample_configs/cehrgpt_pretrain_sample_config.yaml
+```
+
+## Generate synthetic sequences
+Generate synthetic sequences using the trained model
+```bash
+export TRANSFORMERS_VERBOSITY=info
+export CUDA_VISIBLE_DEVICES="0"
+python -u -m cehrgpt.generation.generate_batch_hf_gpt_sequence \
+  --model_folder test_results \
+  --tokenizer_folder test_results \
+  --output_folder test_results \
+  --num_of_patients 128 \
+  --batch_size 32 \
+  --buffer_size 128 \
+  --context_window 1024 \
+  --sampling_strategy TopPStrategy \
+  --top_p 1.0 --temperature 1.0 --repetition_penalty 1.0 \
+  --epsilon_cutoff 0.00 \
+  --demographic_data_path sample_data/pretrain
+```
+
+## Convert synthetic sequences to OMOP
+```bash
+# omop converter requires the OHDSI vocabulary
+export OMOP_VOCAB_DIR = ""
+# the omop derived tables need to be built using pyspark
+export SPARK_WORKER_INSTANCES="1"
+export SPARK_WORKER_CORES="8"
+export SPARK_EXECUTOR_CORES="2"
+export SPARK_DRIVER_MEMORY="2g"
+export SPARK_EXECUTOR_MEMORY="2g"
+
+# Convert the sequences, create the omop derived tables
+sh scripts/omop_pipeline.sh \
+  test_results/top_p10000/generated_sequences/ \
+  test_results/top_p10000/restored_omop/ \
+  $OMOP_VOCAB_DIR
+```
+
+# MEDS Support
+
+This section demonstrates how to pretrain CEHR-GPT using MIMIC-IV data in the MEDS (Medical Event Data Standard) format.
+
+## Prerequisites
+
+Set up the required environment variables before beginning:
+
+```bash
+export CEHR_GPT_MODEL_DIR=""    # Path to CEHR-GPT model directory
+export MEDS_DIR=""              # Path to MEDS data directory
+export MEDS_READER_DIR=""       # Path to MEDS reader output directory
+```
+
+## Step 1: Create MIMIC MEDS Data
+
+Transform your MIMIC files into MEDS format by following the instructions in the [MEDS_transforms](https://github.com/mmcdermott/MEDS_transforms/) repository.
+
+## Step 2: Create the MEDS Reader
+
+Convert the MEDS data for use with CEHR-GPT:
+
+```bash
+meds_reader_convert $MEDS_DIR $MEDS_READER_DIR --num_threads 10
+```
+
+## Step 3: Pretrain CEHR-GPT
+
+Run the pretraining process using the prepared MEDS data:
+
+```bash
+python -u -m cehrgpt.runners.hf_cehrgpt_pretrain_runner \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $CEHR_GPT_MODEL_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --do_train true --seed 42 \
+  --dataloader_num_workers 16 --dataloader_prefetch_factor 8 \
+  --hidden_size 768 --num_hidden_layers 14 --max_position_embeddings 8192 \
+  --evaluation_strategy epoch --save_strategy epoch \
+  --sample_packing --max_tokens_per_batch 16384 \
+  --warmup_steps 500 --weight_decay 0.01 \
+  --num_train_epochs 50 --learning_rate 0.0002 \
+  --use_early_stopping --early_stopping_threshold 0.001 \
+  --is_data_in_meds --inpatient_att_function_type day \
+  --att_function_type day --include_inpatient_hour_token \
+  --include_auxiliary_token --include_demographic_prompt \
+  --meds_to_cehrbert_conversion_type "MedsToBertMimic4"
+```
+
+## Step 4: Generate MEDS Trajectories
+
+### Environment Setup for Trajectory Generation
+
+Configure additional environment variables for trajectory generation with task labels (`subject_id`, `prediction_time`, `boolean_value` [optional]):
+
+```bash
+# MEDS_LABEL_COHORT_DIR must contain a set of parquet files
+export MEDS_LABEL_COHORT_DIR=""     # Path to cohort labels directory
+export MEDS_TRAJECTORY_DIR=""       # Path for trajectory output
+```
+
+### Generate Trajectories
+
+Create synthetic patient trajectories using the trained model:
+
+> **Important:** The total sequence length (`generation_input_length` + `generation_max_new_tokens`) cannot exceed the `max_position_embeddings` value (8192) defined during pretraining.
+
+```bash
+python -u -m cehrgpt.generation.cehrgpt_conditional_generation \
+  --cohort_folder $MEDS_LABEL_COHORT_DIR \
+  --data_folder $MEDS_READER_DIR \
+  --dataset_prepared_path "$CEHR_GPT_MODEL_DIR/dataset_prepared" \
+  --model_name_or_path $CEHR_GPT_MODEL_DIR \
+  --tokenizer_name_or_path $CEHR_GPT_MODEL_DIR \
+  --output_dir $MEDS_TRAJECTORY_DIR \
+  --per_device_eval_batch_size 16 \
+  --num_of_trajectories_per_sample 2 \
+  --generation_input_length 4096 \
+  --generation_max_new_tokens 4096 \
+  --is_data_in_meds \
+  --att_function_type day --inpatient_att_function_type day \
+  --meds_to_cehrbert_conversion_type MedsToBertMimic4 \
+  --include_auxiliary_token --include_demographic_prompt \
+  --include_inpatient_hour_token
+```
+
+### Parameters Explanation
+
+- `generation_input_length`: Controls the length of input context for generation
+- `generation_max_new_tokens`: Maximum number of new tokens to generate
+- `num_of_trajectories_per_sample`: Number of trajectories to generate per patient sample
+
+## Citation
+```
+@article{cehrgpt2024,
+  title={CEHRGPT: Synthetic Data Generation for Electronic Health Records},
+  author={Natarajan, K and others},
+  journal={arXiv preprint arXiv:2402.04400},
+  year={2024}
+}

{cehrgpt-0.1.0 → cehrgpt-0.1.2}/pyproject.toml

@@ -28,14 +28,15 @@ classifiers = [
 ]
 
 dependencies = [
-    "cehrbert==1.4.1",
-    "cehrbert_data==0.0.7",
+    "cehrbert==1.4.5",
+    "cehrbert_data==0.0.11",
     "openai==1.54.3",
     "optuna==4.0.0",
-    "transformers==4.44.0",
+    "transformers==4.44.1",
     "tokenizers==0.19.0",
     "peft==0.10.0",
     "lightgbm",
+    "polars",
 ]
 
 [tool.setuptools_scm]

cehrgpt-0.1.2/scripts/run_linear_prob.sh

@@ -0,0 +1,260 @@
+#!/bin/sh
+
+# Function to display usage information
+usage() {
+    echo "Usage: $0 [options]"
+    echo ""
+    echo "Options:"
+    echo "  --base_dir=DIR                 Base directory containing cohorts (required)"
+    echo "  --dataset_prepared_path=PATH   Path to prepared dataset (required)"
+    echo "  --model_path=PATH              Path to pre-trained model and tokenizer (required)"
+    echo "  --preprocessing_workers=NUM    Number of preprocessing workers (required)"
+    echo "  --batch_size=NUM               Batch size for evaluation (required)"
+    echo "  --output_dir=DIR               Output directory for results (required)"
+    echo "  --model_name=NAME              Name for the model output directory (default: cehrgpt_model)"
+    echo "  --max_tokens_per_batch=NUM     Maximum tokens per batch (default: 16384)"
+    echo "  --torch_type=TYPE              Torch data type (default: float32)"
+    echo "  --disable_sample_packing       Disable sample packing (enabled by default)"
+    echo ""
+    echo "Example:"
+    echo "  $0 --base_dir=/path/to/cohorts --dataset_prepared_path=/path/to/dataset_prepared \\"
+    echo "     --model_path=/path/to/model --preprocessing_workers=16 --batch_size=64 \\"
+    echo "     --output_dir=/path/to/outputs --model_name=my_model --torch_type=float16"
+    exit 1
+}
+
+# Default values
+MODEL_NAME="cehrgpt_model"
+MAX_TOKENS_PER_BATCH="16384"
+TORCH_TYPE="bfloat16"
+DISABLE_SAMPLE_PACKING="false"
+
+# Parse command line arguments
+for arg in "$@"; do
+    case $arg in
+        --base_dir=*)
+            BASE_DIR="${arg#*=}"
+            ;;
+        --dataset_prepared_path=*)
+            DATASET_PREPARED_PATH="${arg#*=}"
+            ;;
+        --model_path=*)
+            MODEL_PATH="${arg#*=}"
+            ;;
+        --preprocessing_workers=*)
+            PREPROCESSING_WORKERS="${arg#*=}"
+            ;;
+        --batch_size=*)
+            BATCH_SIZE="${arg#*=}"
+            ;;
+        --output_dir=*)
+            OUTPUT_DIR="${arg#*=}"
+            ;;
+        --model_name=*)
+            MODEL_NAME="${arg#*=}"
+            ;;
+        --max_tokens_per_batch=*)
+            MAX_TOKENS_PER_BATCH="${arg#*=}"
+            ;;
+        --torch_type=*)
+            TORCH_TYPE="${arg#*=}"
+            ;;
+        --disable_sample_packing)
+            DISABLE_SAMPLE_PACKING="true"
+            ;;
+        --help|-h)
+            usage
+            ;;
+        *)
+            echo "Error: Unknown option: $arg"
+            usage
+            ;;
+    esac
+done
+
+# Check for required arguments
+if [ -z "$BASE_DIR" ] || [ -z "$DATASET_PREPARED_PATH" ] || [ -z "$MODEL_PATH" ] || [ -z "$PREPROCESSING_WORKERS" ] || [ -z "$BATCH_SIZE" ] || [ -z "$OUTPUT_DIR" ]; then
+    echo "Error: Missing required arguments"
+    usage
+fi
+
+# Validate arguments
+if [ ! -d "$BASE_DIR" ]; then
+    echo "Error: Base directory does not exist: $BASE_DIR"
+    exit 1
+fi
+
+if [ ! -d "$DATASET_PREPARED_PATH" ]; then
+    echo "Error: Dataset prepared path does not exist: $DATASET_PREPARED_PATH"
+    exit 1
+fi
+
+if [ ! -d "$MODEL_PATH" ]; then
+    echo "Error: Model path does not exist: $MODEL_PATH"
+    exit 1
+fi
+
+# Create output directory if it doesn't exist
+mkdir -p "$OUTPUT_DIR"
+
+# Check if preprocessing workers is a number
+if ! [ "$PREPROCESSING_WORKERS" -eq "$PREPROCESSING_WORKERS" ] 2>/dev/null; then
+    echo "Error: Preprocessing workers must be a number: $PREPROCESSING_WORKERS"
+    exit 1
+fi
+
+# Check if batch size is a number
+if ! [ "$BATCH_SIZE" -eq "$BATCH_SIZE" ] 2>/dev/null; then
+    echo "Error: Batch size must be a number: $BATCH_SIZE"
+    exit 1
+fi
+
+# Check if max tokens per batch is a number
+if ! [ "$MAX_TOKENS_PER_BATCH" -eq "$MAX_TOKENS_PER_BATCH" ] 2>/dev/null; then
+    echo "Error: Max tokens per batch must be a number: $MAX_TOKENS_PER_BATCH"
+    exit 1
+fi
+
+# Validate torch_type (common PyTorch data types)
+case "$TORCH_TYPE" in
+    float16|float32|float64|bfloat16|int8|int16|int32|int64)
+        ;;
+    *)
+        echo "Error: Invalid torch_type. Supported types: float16, float32, float64, bfloat16, int8, int16, int32, int64"
+        exit 1
+        ;;
+esac
+
+# Validate disable_sample_packing is boolean-like
+if [ "$DISABLE_SAMPLE_PACKING" != "true" ] && [ "$DISABLE_SAMPLE_PACKING" != "false" ]; then
+    echo "Error: disable_sample_packing must be 'true' or 'false': $DISABLE_SAMPLE_PACKING"
+    exit 1
+fi
+
+# Log file setup
+LOG_DIR="$BASE_DIR/logs"
+mkdir -p "$LOG_DIR"
+TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
+MAIN_LOG="$LOG_DIR/run_${TIMESTAMP}.log"
+
+# Log function
+log() {
+    message="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
+    echo "$message" | tee -a "$MAIN_LOG"
+}
+
+# Main execution
+log "Starting feature extraction and model training process"
+log "Configuration:"
+log "  --base_dir=$BASE_DIR"
+log "  --dataset_prepared_path=$DATASET_PREPARED_PATH"
+log "  --model_path=$MODEL_PATH"
+log "  --preprocessing_workers=$PREPROCESSING_WORKERS"
+log "  --batch_size=$BATCH_SIZE"
+log "  --output_dir=$OUTPUT_DIR"
+log "  --model_name=$MODEL_NAME"
+log "  --max_tokens_per_batch=$MAX_TOKENS_PER_BATCH"
+log "  --torch_type=$TORCH_TYPE"
+log "  --disable_sample_packing=$DISABLE_SAMPLE_PACKING"
+
+# Find valid cohorts and write to a temp file
+TEMP_COHORT_LIST="$LOG_DIR/cohort_list_${TIMESTAMP}.txt"
+> "$TEMP_COHORT_LIST" # Clear the file
+
+# Find all valid cohorts (directories with train and test subdirectories)
+for cohort_dir in "$BASE_DIR"/*; do
+    if [ -d "$cohort_dir" ] && [ -d "$cohort_dir/train" ] && [ -d "$cohort_dir/test" ]; then
+        cohort_name=$(basename "$cohort_dir")
+        echo "$cohort_name" >> "$TEMP_COHORT_LIST"
+    fi
+done
+
+# Check if any valid cohorts were found
+if [ ! -s "$TEMP_COHORT_LIST" ]; then
+    log "ERROR: No valid cohorts found in $BASE_DIR"
+    rm -f "$TEMP_COHORT_LIST"
+    exit 1
+fi
+
+# Display all cohorts that will be processed
+cohort_count=$(wc -l < "$TEMP_COHORT_LIST")
+log "Found $cohort_count cohorts to process:"
+while read -r cohort; do
+    log "- $cohort"
+done < "$TEMP_COHORT_LIST"
+
+# Process each cohort sequentially
+while read -r cohort_name; do
+    cohort_dir="$OUTPUT_DIR/$cohort_name"
+    output_dir="$cohort_dir/$MODEL_NAME"
+
+    log "===================================================="
+    log "Processing cohort: $cohort_name"
+    log "===================================================="
+
+    cohort_log="$LOG_DIR/${cohort_name}_${TIMESTAMP}.log"
+
+    # Create output directory if it doesn't exist
+    mkdir -p "$output_dir"
+
+    # Prepare command for feature extraction
+    FEATURE_CMD="python -u -m cehrgpt.tools.linear_prob.compute_cehrgpt_features \
+        --data_folder \"$BASE_DIR/$cohort_name/train/\" \
+        --test_data_folder \"$BASE_DIR/$cohort_name/test/\" \
+        --dataset_prepared_path \"$DATASET_PREPARED_PATH\" \
+        --model_name_or_path \"$MODEL_PATH\" \
+        --tokenizer_name_or_path \"$MODEL_PATH\" \
+        --output_dir \"$output_dir\" \
+        --preprocessing_num_workers \"$PREPROCESSING_WORKERS\" \
+        --per_device_eval_batch_size \"$BATCH_SIZE\" \
+        --max_tokens_per_batch \"$MAX_TOKENS_PER_BATCH\" \
+        --torch_type \"$TORCH_TYPE\""
+
+    # Add sample packing flag if not disabled
+    if [ "$DISABLE_SAMPLE_PACKING" = "false" ]; then
+        FEATURE_CMD="$FEATURE_CMD --sample_packing"
+    fi
+
+    # Step 1: Feature extraction
+    log "Starting feature extraction for $cohort_name..."
+    log "Command: $FEATURE_CMD"
+
+    eval "$FEATURE_CMD > \"$cohort_log\" 2>&1"
+
+    feature_extraction_status=$?
+    if [ $feature_extraction_status -ne 0 ]; then
+        log "ERROR: Feature extraction failed for $cohort_name. Check $cohort_log for details."
+        continue
+    fi
+
+    # Step 2: Model training
+    log "Starting model training for $cohort_name..."
+    log "Command: python -u -m cehrgpt.tools.linear_prob.train_with_cehrgpt_features --features_data_dir $output_dir --output_dir $output_dir"
+
+    python -u -m cehrgpt.tools.linear_prob.train_with_cehrgpt_features \
+        --features_data_dir "$output_dir" \
+        --output_dir "$output_dir" \
+        >> "$cohort_log" 2>&1
+
+    echo "Running meds-evaluation for logistic regression for $TASK_NAME..."
+    meds-evaluation-cli predictions_path="$output_dir/logistic/test_predictions" \
+      output_dir="$output_dir/logistic/"
+
+        # Check if the second command succeeded
+    if [ $? -ne 0 ]; then
+        echo "Error: Running meds-evaluation failed for logistic regression for task $TASK_NAME"
+    fi
+
+    model_training_status=$?
+    if [ $model_training_status -ne 0 ]; then
+        log "ERROR: Model training failed for $cohort_name. Check $cohort_log for details."
+        continue
+    fi
+
+    log "Successfully completed processing for $cohort_name"
+done < "$TEMP_COHORT_LIST"
+
+# Clean up
+rm -f "$TEMP_COHORT_LIST"
+
+log "All processing completed"

cehrgpt-0.1.2/src/cehrgpt/analysis/irregularity.py

@@ -0,0 +1,36 @@
+import os
+
+import polars as pl
+
+from cehrgpt.gpt_utils import extract_time_interval_in_days, is_att_token
+
+
+def main(args):
+    dataset = pl.read_parquet(os.path.join(args.input_dir, "*.parquet"))
+    time_token_frequency_df = (
+        dataset.select(pl.col("concept_ids").explode().alias("concept_id"))
+        .filter(pl.col("concept_id").map_elements(is_att_token))
+        .with_columns(
+            pl.col("concept_id")
+            .map_elements(extract_time_interval_in_days)
+            .alias("time_interval")
+        )
+    )
+    results = time_token_frequency_df.select(
+        pl.mean("time_interval").alias("mean"), pl.std("time_interval").alias("std")
+    ).to_dicts()[0]
+    print(results)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="EHR Irregularity analysis")
+    parser.add_argument(
+        "--input_dir",
+        dest="input_dir",
+        action="store",
+        help="The path for where the input data folder",
+        required=True,
+    )
+    main(parser.parse_args())

{cehrgpt-0.1.0 → cehrgpt-0.1.2}/src/cehrgpt/data/hf_cehrgpt_dataset.py

@@ -23,6 +23,7 @@ CEHRGPT_COLUMNS = [
     "num_of_visits",
     "values",
     "value_indicators",
+    "epoch_times",
 ]
 
 TRANSFORMER_COLUMNS = ["input_ids"]