cehrgpt 0.0.2__tar.gz → 0.1.1__tar.gz

{cehrgpt-0.0.2 → cehrgpt-0.1.1}/.gitignore

@@ -4,6 +4,9 @@
 venv*
 dist/*
 
+*png
+*json
+
 *ipynb_checkpoints/
 *h5
 *logs

{cehrgpt-0.0.2/src/cehrgpt.egg-info → cehrgpt-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: cehrgpt
-Version: 0.0.2
+Version: 0.1.1
 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,13 +12,15 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: cehrbert==1.3.8
+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.40.0
+Requires-Dist: transformers==4.44.1
 Requires-Dist: tokenizers==0.19.0
 Requires-Dist: peft==0.10.0
-Requires-Dist: trl==0.11.4
+Requires-Dist: lightgbm
+Requires-Dist: polars
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
@@ -29,14 +31,15 @@ Requires-Dist: hypothesis; extra == "dev"
 Requires-Dist: black; extra == "dev"
 Provides-Extra: flash-attn
 Requires-Dist: flash_attn; extra == "flash-attn"
+Dynamic: license-file
 
 # 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-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.

{cehrgpt-0.0.2 → cehrgpt-0.1.1}/README.md

@@ -2,9 +2,9 @@
 
 [![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.
@@ -77,4 +77,4 @@ sh scripts/omop_pipeline.sh \
   author={Natarajan, K and others},
   journal={arXiv preprint arXiv:2402.04400},
   year={2024}
-}
+}

{cehrgpt-0.0.2 → cehrgpt-0.1.1}/pyproject.toml

@@ -28,13 +28,15 @@ classifiers = [
 ]
 
 dependencies = [
-    "cehrbert==1.3.8",
+    "cehrbert==1.4.5",
+    "cehrbert_data==0.0.11",
     "openai==1.54.3",
     "optuna==4.0.0",
-    "transformers==4.40.0",
+    "transformers==4.44.1",
     "tokenizers==0.19.0",
     "peft==0.10.0",
-    "trl==0.11.4",
+    "lightgbm",
+    "polars",
 ]
 
 [tool.setuptools_scm]

{cehrgpt-0.0.2 → cehrgpt-0.1.1}/scripts/level_three_evaluation.sh

@@ -29,7 +29,8 @@ python -u -m cehrbert_data.prediction_cohorts.cad_cabg_cohort \
     -dl 1985-01-01 -du 2023-12-31 \
     -l 18 -u 100 -ow 360 -ps 0 -pw 360 -f \
     --att_type cehr_bert \
-    --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+    --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+    --is_remove_index_prediction_starts
 
 # Run Predictions on CAD CABG
 echo "Run predictions on cad_cabg"
@@ -56,9 +57,10 @@ python -u -m cehrbert_data.prediction_cohorts.hf_readmission \
   -c hf_readmission_bow \
   -i "$OMOP_FOLDER" \
   -o "$OMOP_FOLDER/cohorts/hf_readmission" \
-  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 360 -ps 0 -pw 30 -f \
+  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 360 -ps 1 -pw 30 -f \
   --att_type cehr_bert \
-  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+  --is_remove_index_prediction_starts
 
 # Run predictions on HF Readmission
 echo "Run predictions on hf_readmission"
@@ -85,9 +87,10 @@ python -u -m cehrbert_data.prediction_cohorts.copd_readmission \
   -c copd_readmission_bow \
   -i "$OMOP_FOLDER" \
   -o "$OMOP_FOLDER/cohorts/copd_readmission" \
-  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 720 -ps 0 -pw 360 -f \
+  -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 360 -ps 1 -pw 30 -f \
   --att_type cehr_bert \
-  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+  --is_remove_index_prediction_starts
 
 # Run predictions on COPD Readmission
 echo "Run predictions on copd_readmission"
@@ -145,7 +148,8 @@ python -u -m cehrbert_data.prediction_cohorts.afib_ischemic_stroke \
   -o "$OMOP_FOLDER/cohorts/afib_ischemic_stroke" \
   -dl 1985-01-01 -du 2023-12-31 -l 18 -u 100 -ow 720 -ps 0 -pw 360 -f \
   --att_type cehr_bert \
-  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv
+  --ehr_table_list condition_occurrence procedure_occurrence drug_exposure -iv \
+  --is_remove_index_prediction_starts
 
 # Run predictions on AFIB Ischemic Stroke
 echo "Run predictions on afib_ischemic_stroke"

cehrgpt-0.1.1/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.1/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.0.2 → cehrgpt-0.1.1}/src/cehrgpt/data/hf_cehrgpt_dataset.py

@@ -1,9 +1,10 @@
-from typing import Union
+from typing import Optional, Union
 
 from cehrbert.data_generators.hf_data_generator.hf_dataset import (
     FINETUNING_COLUMNS,
     apply_cehrbert_dataset_mapping,
 )
+from cehrbert.data_generators.hf_data_generator.meds_utils import CacheFileCollector
 from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments
 from datasets import Dataset, DatasetDict
 
@@ -22,6 +23,7 @@ CEHRGPT_COLUMNS = [
     "num_of_visits",
     "values",
     "value_indicators",
+    "epoch_times",
 ]
 
 TRANSFORMER_COLUMNS = ["input_ids"]
@@ -31,16 +33,25 @@ def create_cehrgpt_pretraining_dataset(
     dataset: Union[Dataset, DatasetDict],
     cehrgpt_tokenizer: CehrGptTokenizer,
     data_args: DataTrainingArguments,
-) -> Dataset:
+    cache_file_collector: Optional[CacheFileCollector] = None,
+) -> Union[Dataset, DatasetDict]:
     required_columns = TRANSFORMER_COLUMNS + CEHRGPT_COLUMNS
+    # TODO: temp solution, this column is mixed typed and causes an issue when transforming the data
+    if not data_args.streaming:
+        if isinstance(dataset, DatasetDict):
+            all_columns = dataset["train"].column_names
+        else:
+            all_columns = dataset.column_names
+        if "visit_concept_ids" in all_columns:
+            dataset.remove_columns(["visit_concept_ids"])
     dataset = apply_cehrbert_dataset_mapping(
         dataset,
         HFCehrGptTokenizationMapping(cehrgpt_tokenizer),
         num_proc=data_args.preprocessing_num_workers,
         batch_size=data_args.preprocessing_batch_size,
         streaming=data_args.streaming,
+        cache_file_collector=cache_file_collector,
     )
-
     if not data_args.streaming:
         if isinstance(dataset, DatasetDict):
             all_columns = dataset["train"].column_names
@@ -56,8 +67,17 @@ def create_cehrgpt_finetuning_dataset(
     dataset: Union[Dataset, DatasetDict],
     cehrgpt_tokenizer: CehrGptTokenizer,
     data_args: DataTrainingArguments,
-) -> Dataset:
+    cache_file_collector: Optional[CacheFileCollector] = None,
+) -> Union[Dataset, DatasetDict]:
     required_columns = TRANSFORMER_COLUMNS + CEHRGPT_COLUMNS + FINETUNING_COLUMNS
+    # TODO: temp solution, this column is mixed typed and causes an issue when transforming the data
+    if not data_args.streaming:
+        if isinstance(dataset, DatasetDict):
+            all_columns = dataset["train"].column_names
+        else:
+            all_columns = dataset.column_names
+        if "visit_concept_ids" in all_columns:
+            dataset.remove_columns(["visit_concept_ids"])
     mapping_functions = [
         HFFineTuningMapping(cehrgpt_tokenizer),
     ]
@@ -68,6 +88,7 @@ def create_cehrgpt_finetuning_dataset(
             num_proc=data_args.preprocessing_num_workers,
             batch_size=data_args.preprocessing_batch_size,
             streaming=data_args.streaming,
+            cache_file_collector=cache_file_collector,
         )
 
     if not data_args.streaming: