CoD, Towards an Interpretable Medical Agent using Chain of Diagnosis

📃 Paper • 🖥️ Demo

📚 CoD Data • 📚 Disease Database • 📝 DxBench(Benchmark)

🤗 DiagnosisGPT-34B • 🤗 DiagnosisGPT-6B

✨ Updates

• [2024/07/19]: We launched our diagnostic LLMs, DiagnosisGPT-6B and DiagnosisGPT-34B, supporting automatic diagnosis for 9,604 diseases.
• [2024/07/18]: We released our data, including CoD data, the disease database, and the DxBench benchmark.
• [2024/07/18]: We introduced the Chain-of-Diagnosis (CoD) to improve the trustworthiness and interpretability of medical diagnostics.

⚡ Introduction

We propose the Chain-of-Diagnosis to improve interpretability in medical diagnostics for Large Language Models (LLMs). CoD features include:

1. Transforming the opaque decision-making process into a five-step diagnostic chain that reflects a physician’s thought process.
2. Producing a confidence distribution, where higher confidence suggests greater certainty in diagnosing a specific disease. CoD formalizes the diagnostic process as a process of reducing diagnostic certainty entropy.

👩‍⚕️ Model

We developed DiagnosisGPT, an LLM using the CoD method for automated diagnosis of 9,604 diseases.

	Backbone	Checkpoint
DiagnosisGPT-7B	Yi-6B	HF Link
DiagnosisGPT-34B	Yi-34B	HF Link

• Command Line Interface

Chat via the command line:

python cod_cli.py --model_dir path-to-diagnosisgpt-vision-model

• Model Inference

Inference using our ChatBot:

query = 'I have a sore throat and seem to have a fever.'

from cod_cli import DiagnosisChatbot
bot = DiagnosisChatbot(path-to-diagnosisgpt-vision-model, confidence_threshold=0.5)

output, history, confidence_distribution = bot.inference(query)
print(output) # Prints the output of the model

📚 Data

Data	Description	Download Link
CoD Data	(1) Data of 50K synthetic patients constructed from medical encyclopedia. (2) 50K Chain-of-Diagnosis training data.	Link
Disease Database	Database of 9,604 diseases with symptoms and treatments.	Link
DxBench	A large-scale disease diagnosis benchmark built on real patient-doctor dialogue, covering 15 departments and 461 diseases.	Link

🏋️‍♂️ Training

1. Prepare the Training Data You can download the CoD data and run the following code for preprocessing:

python train/datapre.py --model_path 01-ai/Yi-34B --data_dir CoD_en.json

2. Training Then, you can train the model by:

cd train
bash fine_tune.sh

🧐 Evaluation

Run the following code for evaluation.

# Evaluate DxBench 
accelerate launch eval.py --model_path path_to_the_model --data_path ./evaluation/DxBench_zh.json --max_inquiry 0 --threshold 0.5
# Evaluate Dxy Dataset
accelerate launch eval.py --model_path path_to_the_model --data_path ./evaluation/Dxy_zh.json --max_inquiry 0 --threshold 0.5
# Evaluate Muzhi Dataset
accelerate launch eval.py --model_path path_to_the_model --data_path ./evaluation/Muzhi_zh.json --max_inquiry 0 --threshold 0.5

Citation

@misc{chen2024codinterpretablemedicalagent,
      title={CoD, Towards an Interpretable Medical Agent using Chain of Diagnosis}, 
      author={Junying Chen and Chi Gui and Anningzhe Gao and Ke Ji and Xidong Wang and Xiang Wan and Benyou Wang},
      year={2024},
      eprint={2407.13301},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2407.13301}, 
}