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:
| File | Contents |
|---|---|
pretrained_decoder.pt |
Decoder pretrained to predict the next clinical event code |
best_model.pt |
Full decoder and multi-label classification head selected by validation mean average precision |
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 <git@github.com: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:
- Causal pretraining: next-event code prediction.
- 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.