Large Language Models for Drug Overdose Prediction from Longitudinal Medical Records

⚠️ Important: Access to the Truven dataset is restricted. It is proprietary data and requires special permissions. Please contact the project authors for more information on how to access it.

---

📁 Project Structure

├── cohort_selection.py                 
├── prompt_design.py                   
├── data_preprocessing_baseline_model.py 
├── fine_tune_data_preparation.py        
├── test_gpt_model.py                  
├── cohorts/                           
├── output/                            
└── README.md

🛠️ Setup Instructions

To get started, create a virtual environment and install the required dependencies.

# Create and activate virtual environment
python -m venv venv
source venv/bin/activate    # For Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

🚀 How to Run the Pipeline

✅ Step 1: Cohort Selection and Data Preprocessing

Run cohort_selection.py to preprocess raw Truven patient visit data and generate case/control cohorts for drug overdose prediction. This step formats the data for model input based on the two prediction windows (7 days and 30 days).

🧾 Arguments

Argument	Type	Description
--cohorts_type	str	Type of cohort to generate (case or control)
--cohorts_file	str	Contain the patient id and corresponding cohort class
--input_data_dir	str	Directory containing raw visit data
--output_data_dir	str	Directory where the preprocessed files will be stored

✅ Example

python cohort_selection.py \
    --cohorts_type case \
    --cohorts_file cohorts/cohorts.csv \
    --input_data_dir truven \
    --output_data_dir truven_output

✏️ Step 2: Prompt Generation for LLMs

Use prompt_design.py to generate textual prompts from the preprocessed data for input to GPT or other LLMs.

✅ Example

python prompt_design.py \
    --max_visit 30 \
    --cohort_type case \
    --prediction_window 1-week \
    --split test \
    --input_dir truven_output/preprocessed_data \
    --output_dir output_prompts

🔧 Step 3: Prepare Data for Fine-Tuning

After creating the cohort files with structured prompts (from Step 2), use them to prepare a fine-tuning dataset for GPT models. This step uses the fine_tune_data_preparation.py script to format the data into a .jsonl file compatible with OpenAI's fine-tuning API.

📜 Usage

Run the script with the following arguments:

Argument	Type	Description
--max_visit	int	Maximum number of visits to include per patient (default: 30)
--source_dir	str	Directory containing the prompt-based cohort files
--target_file	str	Path to save the output fine-tuning file in JSONL format
--prediction_window	str	The prediction window used in the cohort (e.g., 1-week)

✅ Example

python fine_tune_data_preparation.py \
  --max_visit 30 \
  --source_dir OD-Only_optimized_prompt/30/test \
  --target_file formatted_fine_tuning_file.jsonl \
  --prediction_window 1-week

📊 Step 4: Generate Statistical Features for Baseline ML Models

In addition to preparing prompts for LLMs, the cohort files from Step 2 can also be used to extract statistical features for each patient. These features can be used to train and evaluate traditional machine learning models such as XGBoost and Random Forest.

Use the data_preprocessing_baseline_model.py script to generate these features.

📜 Usage

Run the script with the following arguments:

Argument	Type	Description
--max_visit	int	Maximum number of visits to include per patient (default: 30)
--cohort_type	str	Cohort type (case or control)
--prediction_window	str	Prediction window used in the cohort (e.g., 1-week)
--input_dir	str	Directory containing cohort files generated in Step 2
--output_dir	str	Directory to save the extracted feature files

✅ Example

python data_preprocessing_baseline_model.py \
  --max_visit 30 \
  --cohort_type case \
  --prediction_window 1-week \
  --input_dir output_prompts \
  --output_dir output_baseline_features

🧪 Step 5: Test GPT Models (Zero-Shot & Fine-Tuned)

This step evaluates the model's ability to predict drug overdose outcomes using zero-shot GPT models or a fine-tuned GPT model.

Use the test_gpt_model.py script to run the evaluation. Before running the script, you must edit the file to include your OpenAI API key. This is required to access the OpenAI models (e.g., GPT-4, GPT-3.5, etc.).

⚠️ Important: OpenAI API access is not included in this repository. You must obtain your own API key and securely set it inside the script or as an environment variable.

📜 Usage

Run the script with the following arguments:

Argument	Type	Description
--test_directory	str	Path to the test cohort files with prompt-formatted data
--output_file	str	The output file path that contains the prediction for each patient in the test data
--prediction_window	str	Prediction window used in the cohort (e.g., 1-week)
--model_name	str	GPT model to use for evaluation (e.g., gpt-4o, gpt-3.5-turbo, or your fine-tuned model)

✅ Example

python test_gpt_model.py \
  --test_directory OD-Only_optimized_prompt/30/test/case_cohorts/1-week \
  --target_file formatted_fine_tuning_file.jsonl \
  --prediction_window 1-week \
  --model_name gpt-4o