Forecasting_Future_Conditions

EHR Transformer Decoder

A causal Transformer for learning from longitudinal electronic health record (EHR) event sequences and predicting incident clinical conditions over a five-year horizon.

This repository contains two PyTorch checkpoints:

Research use only. This model is not a medical device or clinical decision-support system. Its outputs must not be used to diagnose, treat, or make decisions about patients.

Model description

The model processes a patient's coded clinical history in chronological order. Each event combines:

• A clinical code embedding
• An event-source/type embedding
• A learned position embedding
• Three continuous time features
• Patient gender and race embeddings

The time features represent days before the prediction anchor, age at the event, and time since the previous event. A causal Transformer learns the sequence representation. The fine-tuned model pools the last non-padding state and produces one probability per target condition.

Supported event sources are conditions, medications, procedures, observations, encounters, care plans, and immunizations.

Checkpoint format

Both files are raw PyTorch state_dict checkpoints saved with torch.save(model.state_dict(), ...). They are not Hugging Face Transformers AutoModel checkpoints and cannot be loaded with AutoModel.from_pretrained() or the hosted inference widget.

Inference requires the model classes and preprocessing pipeline from the source project. It also requires the exact vocabulary, ordered target-condition list, and architecture used during training.

Installation

Clone the source project, then install its dependencies:

git clone <[email protected]:Salma-Jamal/Forecasting_Future_Conditions.git>

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch pandas numpy scikit-learn tqdm huggingface_hub

Download the checkpoints

SalmaJamal/Forecasting_Future_Conditions

Python

from huggingface_hub import hf_hub_download

repo_id = "SalmaJamal/Forecasting_Future_Conditions"

best_model_path = hf_hub_download(
    repo_id=repo_id,
    filename="best_model.pt",
)
pretrained_decoder_path = hf_hub_download(
    repo_id=repo_id,
    filename="pretrained_decoder.pt",
)

print(best_model_path)
print(pretrained_decoder_path)

Command line

hf download SalmaJamal/Forecasting_Future_Conditions \
  best_model.pt pretrained_decoder.pt \
  --local-dir ./checkpoints

Use the pretrained decoder

The source project can rebuild the vocabulary from the training split and use pretrained_decoder.pt to initialize fine-tuning:

python run.py \
  --data-dir ./data \
  --output-dir ./outputs/finetune \
  --skip-pretrain \
  --pretrain-ckpt ./checkpoints/pretrained_decoder.pt

Expected input data

The preprocessing code expects a directory containing:

data/
├── patient_splits.csv
├── target_conditions.csv
├── test_anchors.csv              # optional
├── train_val/
│   ├── patients.csv
│   ├── encounters.csv
│   ├── conditions.csv
│   ├── observations.csv
│   ├── medications.csv
│   ├── procedures.csv
│   ├── immunizations.csv
│   └── careplans.csv
└── test/
    └── ...same table names...

patient_splits.csv requires Id and split columns. target_conditions.csv requires a CODE column. Event tables use PATIENT, CODE, and their source-specific date column. See the source project's README for the complete schema.

Training objective

Training uses two stages:

1. Causal pretraining: next-event code prediction.
2. Multi-label fine-tuning: prediction of target conditions that first occur within five years after the anchor.

The default fine-tuning setup uses class-weighted focal loss, an auxiliary next-event loss, multi-anchor training augmentation, AdamW, a one-cycle learning-rate schedule, gradient clipping, and early stopping on validation mean average precision.