# MedClarify: An information-seeking language model for medical diagnosis with case-specific follow-up questions

## Table of Contents

1. [Installation](#installation)
2. [Quick Start](#quick-start)
3. [Command Options](#command-options)
4. [Results & Analysis](#results--analysis)
5. [Datasets & Preprocessing](#data-preprocess)

---

## Installation

- Prerequisites: Python 3.12+

- Install dependencies
```
pip install -r requirements.txt
```

- Set up your API keys in `.env` file
````
DEEPINFRA_API_KEY=[your deepinfra api key here]
OPENAI_API_KEY=[your openai api key here]
ICD_CLIENT_ID=[your icd client id here]
ICD_CLIENT_SECRET=[your icd client secret here]
````

More information about how to create an account for ICD-11 API can be found here: https://icd.who.int/icdapi

---

## Quick Start

The fastest way to run experiments:

```bash
# Navigate
cd diagnosis_experiment

# Run a baseline experiment
python main.py --baseline

# Run a custom experiment
python main.py \
  --medclarify
  --dataset mediq \
  --model deepseek-ai/DeepSeek-R1-0528
```

All the results will be saved under the directory `logs`.

---

## Command Options

| Argument | Description | Options |
|-----------|-------------|-----------|
| `--single_shot_incomplete` | Run single-shot experiment on incomplete cases | - |
| `--single_shot_complete` | Run single-shot experiment on complete cases | - |
| `--baseline` | Run naive multi-turn baseline experiment | - |
| `--medclarify` | Run MedClarify experiment | - |
| `--dataset` | Dataset source | default: `nejm`; Other options: `mediq`, `medqa` |
| `--model` | LLM model | default: `meta-llama/Llama-3.3-70B-Instruct-Turbo`; Other options: `gpt-5.1-2025-11-13`, `deepseek-ai/DeepSeek-R1-0528` |
| `--mask_type` | Medical feature category to mask | default: `None`; Other options: `Symptoms`, `Imaging`, `Laboratory`, `Past_Medical_History`, `Physical_Examination`, `Social_History` |
| `--max_num_questions` | Max. interactive questions | default: `10` |
| `--top_k` | Top-k diagnoses to consider | dafault: `5` |

#### Experiments

##### Replication of Figure 4
```
python main.py --single_shot_complete --dataset nejm
python main.py --single_shot_complete --dataset mediq
python main.py --single_shot_complete --dataset medqa
```
```
python main.py --single_shot_incomplete --mask_type Imaging
python main.py --single_shot_incomplete --mask_type Laboratory
python main.py --single_shot_incomplete --mask_type Symptoms
python main.py --single_shot_incomplete --mask_type Past_Medical_History
python main.py --single_shot_incomplete --mask_type Physical_Examination
python main.py --single_shot_incomplete --mask_type Social_History
```
- `single_shot_complete`- a direct diagnosis without further interactive questioning with patient agent on complete cases.

- `single_shot_incomplete` - a direct diagnosis without further interactive questioning with patient agent on incomplete cases; this experiment requires to choose an option (i.e., `Symptoms`, `Imaging`, `Laboratory`, `Past_Medical_History`, `Physical_Examination`, `Social_History`) for `--mask_type` argument. `dataset` argument for choosing different datasets (i.e., `nejm`, `mediq`, `medqa`)

##### Replication of Figure 5, Extended Figure 6
```
python main.py --baseline --dataset nejm --mask_type Imaging
python main.py --baseline --dataset nejm --mask_type Laboratory
python main.py --baseline --dataset nejm --mask_type Symptoms
python main.py --baseline --dataset nejm --mask_type Past_Medical_History
python main.py --baseline --dataset nejm --mask_type Physical_Examination
python main.py --baseline --dataset nejm --mask_type Social_History
```
```
python main.py --medclarify --dataset nejm --mask_type Imaging
python main.py --medclarify --dataset nejm --mask_type Laboratory
python main.py --medclarify --dataset nejm --mask_type Symptoms
python main.py --medclarify --dataset nejm --mask_type Past_Medical_History
python main.py --medclarify --dataset nejm --mask_type Physical_Examination
python main.py --medclarify --dataset nejm --mask_type Social_History
```
- `baseline`- implement a multi-turn diagnostic by LLM asking the naive questions.

- `medclarify` - implement a multi-turn diagnostic by using a optimized questioning strategy and confidence calibration through Bayesian updating.

##### Replication of Figure 6, 8, Extended Figure 5, 7, 8, 9
```
python main.py --baseline --dataset nejm
python main.py --baseline --dataset mediq
python main.py --baseline --dataset medqa
```
```
python main.py --medclarify --dataset nejm
python main.py --medclarify --dataset mediq
python main.py --medclarify --dataset medqa
```

##### Replication of Figure 7, Extended Figure 10
```
python main.py --baseline --dataset nejm --model deepseek-ai/DeepSeek-R1-0528
python main.py --baseline --dataset nejm --model gpt-5.1-2025-11-13
```
```
python main.py --medclarify --dataset nejm --model deepseek-ai/DeepSeek-R1-0528
python main.py --medclarify --dataset nejm --model gpt-5.1-2025-11-13
```
- `dataset` argument for choosing different datasets (i.e., `nejm`, `mediq`, `medqa`)

## Results & Analysis

- Run the results in [Experiments](#experiments), the logs will be saved under `logs`directory.
- The result Figures can be generated using: `notebook/Result_Analysis.ipynb`
- Replace the path of the results (i.e., `MedClarify_(all-features_masking)_nejm_2026-02-12_22-38-23/medclarify_(all-features_masking)_nejm_case_summaries_2026-02-12_22-38-23.jsonl`) as instructed in the beginning of the notebook.
- All the diagrams will save automatically under the directory `img`

## Datasets & Preprocessing

- Datasets used in the experiment are under the directory `dataset`
  - Raw datasets: `nejm.jsonl`, `mediq.jsonl`, and `medqa.jsonl` 
  - Processed datasets: `structural_nejm.jsonl`, `structural_mediq.jsonl`

- The structural datasets can be automatically generated through running the `notebook/Data_Processing.ipynb`