⚕️Automated Medical Coding on MIMIC-III and MIMIC-IV: A Critical Review and Replicability Study

Official source code repository for the SIGIR 2023 paper Automated Medical Coding on MIMIC-III and MIMIC-IV: A Critical Review and Replicability Study

@inproceedings{edinAutomatedMedicalCoding2023,
  address = {Taipei, Taiwan},
  title = {Automated {Medical} {Coding} on {MIMIC}-{III} and {MIMIC}-{IV}: {A} {Critical} {Review} and {Replicability} {Study}},
  isbn = {978-1-4503-9408-6},
  shorttitle = {Automated {Medical} {Coding} on {MIMIC}-{III} and {MIMIC}-{IV}},
  doi = {10.1145/3539618.3591918},
  booktitle = {Proceedings of the 46th {International} {ACM} {SIGIR} {Conference} on {Research} and {Development} in {Information} {Retrieval}},
  publisher = {ACM Press},
  author = {Edin, Joakim and Junge, Alexander and Havtorn, Jakob D. and Borgholt, Lasse and Maistro, Maria and Ruotsalo, Tuukka and Maaløe, Lars},
  year = {2023}
}

Introduction

Automatic medical coding is the task of automatically assigning diagnosis and procedure codes based on discharge summaries from electronic health records. This repository contains the code used in the paper Automated medical coding on MIMIC-III and MIMIC-IV: A Critical Review and Replicability Study. The repository contains code for training and evaluating medical coding models and new splits for MIMIC-III and the newly released MIMIC-IV. The following models have been implemented:

Model	Paper	Original Code
CNN	Explainable Prediction of Medical Codes from Clinical Text	link
Bi-GRU	Explainable Prediction of Medical Codes from Clinical Text	link
CAML	Explainable Prediction of Medical Codes from Clinical Text	link
MultiResCNN	ICD Coding from Clinical Text Using Multi-Filter Residual Convolutional Neural Network	link
LAAT	A Label Attention Model for ICD Coding from Clinical Text	link
PLM-ICD	PLM-ICD: Automatic ICD Coding with Pretrained Language Models	link

The splits are found in files/data. The splits are described in the paper.

How to reproduce results

Setup Conda environment

Create a conda environement conda create -n coding python=3.10
Install the packages pip install . -e

Prepare MIMIC-III

This code has been developed on MIMIC-III v1.4.

Download the MIMIC-III data into your preferred location path/to/mimiciii. Please note that you need to complete training to acces the data. The training is free, but takes a couple of hours. - link to data access
Open the file src/settings.py
Change the variable DOWNLOAD_DIRECTORY_MIMICIII to the path of your downloaded data path/to/mimiciii
If you want to use the MIMIC-III full and MIMIC-III 50 from the Explainable Prediction of Medical Codes from Clinical Text you need to run python prepare_data/prepare_mimiciii_mullenbach.py
If you want to use MIMIC-III clean from our paper you need to run python prepare_data/prepare_mimiciii.py

Prepare MIMIC-IV

This code has been developed on MIMIC-IV and MIMIC-IV v2.2.

Download MIMIC-IV and MIMIC-IV-NOTE into your preferred location path/to/mimiciv and path/to/mimiciv-note. Please note that you need to complete training to acces the data. The training is free, but takes a couple of hours. - link to data access
Open the file src/settings.py
Change the variable DOWNLOAD_DIRECTORY_MIMICIV to the path of your downloaded data path/to/mimiciv
Change the variable DOWNLOAD_DIRECTORY_MIMICIV_NOTE to the path of your downloaded data path/to/mimiciv-note
Run python prepare_data/prepare_mimiciv.py

Before running experiments

Create a weights and biases account. It is possible to run the experiments without wandb.
Download the model checkpoints and unzip it. Please note that these model weights can't be used commercially due to the MIMIC License.
If you want to train PLM-ICD, you need to download RoBERTa-base-PM-M3-Voc, unzip it and change the model_path parameter in configs/model/plm_icd.yaml and configs/text_transform /huggingface.yaml to the path of the download.

Running experiments

Training

You can run any experiment found in configs/experiment. Here are some examples:

Train PLM-ICD on MIMIC-III clean on GPU 0: python main.py experiment=mimiciii_clean/plm_icd gpu=0
Train CAML on MIMIC-III full on GPU 6: python main.py experiment=mimiciii_full/caml gpu=6
Train LAAT on MIMIC-IV ICD-9 full on GPU 6: python main.py experiment=mimiciv_icd9/laat gpu=6
Train LAAT on MIMIC-IV ICD-9 full on GPU 6 without weights and biases: python main.py experiment=mimiciv_icd9/laat gpu=6 callbacks=no_wandb trainer.print_metrics=true

Evaluation

If you just want to evaluate the models using the provided model_checkpoints you need to do set trainer.epochs=0 and provide the path to the models checkpoint load_model=path/to/model_checkpoint. Make sure you the correct model-checkpoint with the correct configs.

Example: Evaluate PLM-ICD on MIMIC-IV ICD-10 on GPU 1: python main.py experiment=mimiciv_icd10/plm_icd gpu=1 load_model=path/to/model_checkpoints/mimiciv_icd10/plm_icd epochs=0

Overview of the repository

configs

We use Hydra for configurations. The condigs for every experiment is found in configs/experiments. Furthermore, the configuration for the sweeps are found in configs/sweeps. We used Weights and Biases Sweeps for most of our experiments.

files

This is where the images and data is stored.

notebooks

The directory only contains one notebook used for the code analysis. The notebook is not aimed to be used by others, but is included for others to validate our data analysis.

prepare_data

The directory contains all the code for preparing the datasets and generating splits.

reports

This is the code used to generate the plots and tables used in the paper. The code uses the Weights and Biases API to fetch the experiment results. The code is not usable by others, but was included for the possibility to validate our figures and tables.

src

This is were the code for running the experiments is found.

tests

The directory contains the unit tests

⚠️ Known issues

LAAT and PLM-ICD are unstable. The loss will sometimes diverge during training. The issue seems to be overflow in the softmax function in the label-wise attention. Using batch norm or layer norm before the softmax function might solve the issue. We did not try to fix the issue as we didn't want to change the original method during our reproducibility.
The code was only tested on a server with 128 GB RAM. A user with 32 GB RAM reported issues fitting MIMIC-IV into memory.

Acknowledgement

Thank you Sotiris Lamprinidis for providing an efficient implementation of our multi-label stratification algorithm and some data preprocessing helper functions.

JoakimEdin / medical-coding-reproducibility