Code

This is the code for the paper "Simulating Radiologists with On-demanding Few-shot Adaptation of Large Language Models". It contains the code for training data selection, on-demanding training and inference, and evaluation. This code also contains a small set of data, which you can use to test the code.

Requirements

The code works with the following environment

python=3.10
pytorch==2.2.2
transformers==4.37.2
deepspeed==0.14.5
nltk==3.8.1

Training Data Selection

Visual feature comparison (VFC)

Visual feature comparison requires a pre-trained vision Transformer. For example, one can download this one and put it under the llm folder. The file structure would be something like.

- llm
|-- clip-vit-large-patch14-336

Run the following commands to obtain the training data.

python TDS/compute_image_features.py --model_path llm/clip-vit-large-patch14-336
python TDS/visual_feature_comparison.py

Pathology matching (PM)

Pathology matching requires a pre-trained pathology matcher. For example, one can download this one and put it under the llm folder. The file structure would be something like.

- llm
|-- llava-v1.6-vicuna-13b

Run the following commands to obtain the training data.

cd pathology_extractor
CUDA_VISIBLE_DEVICES=0 python extract_terms.py --model_path ../llm/llava-v1.6-vicuna-13b
python pathology_matching.py

Merge the data from VFC and PM

Run the following command to merge the data from VFC and PM to obtain the final data.

cd ..
python merge_training_data.py ./data/tmp/vfc ./data/tmp/pm ./data/tds

On Demanding Training and Inference

This step requires a pre-trained multimodal large language model. For example, one can download this one and put it under the llm folder. The file structure would be something like.

- llm
|-- llava-v1.6-vicuna-13b

Run the following commands to train the model without CID.

MODEL_PATH=./llm/llava-v1.6-vicuna-13b
VISION_PATH=./llm/clip-vit-large-patch14-336
INPUT_DIR=./data/tds
INPUT_PATH=./data/annotation.json
IMAGE_DIR=./data/images
RESULTS_OUTPUT_DIR=results/demo
MODEL_OUTPUT_DIR=./output/llava-v1.6-13b-demo/
python run_on_demand.py \
    --gpus 0,1,2,3 \
    --epoch 1 \
    --base_model_path $MODEL_PATH \
    --vision_tower $VISION_PATH \
    --input_dir $INPUT_DIR \
    --input_path $INPUT_PATH \
    --image_dir $IMAGE_DIR \
    --results_output_dir $RESULTS_OUTPUT_DIR \
    --model_output_dir $MODEL_OUTPUT_DIR \
    --deepspeed ./scripts/zero3_offload.json \
    --continue_train

Run the following commands to train the model with CID.

MODEL_PATH=./llm/llava-v1.6-vicuna-13b
VISION_PATH=./llm/clip-vit-large-patch14-336
INPUT_DIR=./data/tds
INPUT_PATH=./data/annotation.json
IMAGE_DIR=./data/images
RESULTS_OUTPUT_DIR=results/demo
MODEL_OUTPUT_DIR=./output/llava-v1.6-13b-demo/
python run_on_demand.py \
    --gpus 0,1,2,3 \
    --epoch 1 \
    --base_model_path $MODEL_PATH \
    --vision_tower $VISION_PATH \
    --input_dir $INPUT_DIR \
    --input_path $INPUT_PATH \
    --image_dir $IMAGE_DIR \
    --results_output_dir $RESULTS_OUTPUT_DIR \
    --model_output_dir $MODEL_OUTPUT_DIR \
    --deepspeed ./scripts/zero3_offload.json \
    --continue_train \
    --use_cid \
    --cid_size 3

Evaluation

Run the following command to get the results on the test data.

python merge_results.py --input_dir results/demo

Notes

All the above commands are available in run.sh.