DiagRL-7B

For implementation detailes, please refer to DiagRL github repository

Introduction

This repository contains the 7B version of DiagRL Model trained through retrieval-augmented reinforcement learning for medical diagnosis. Our model is post-trained from Qwen2.5-7B-Instruct, keeping general model ability while achieving better retrieval-reasoning abilities in medical domain.

The model can be loaded and directly inferenced as a general-purpose medical LLM as:

Basic Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "QiaoyuZheng/DiagRL-7B"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "A patient suffers from fever, pain, feeling very tired, paleness in face, frequent infections, easy infections and bruising, bleeding with no clear cause, such as in the nose or gums and shortness of breath."
messages = [
    {"role": "system", "content": "You are DiagRL, created by SJTU. You are a helpful agent on diagnosis."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(response)

Advanced Usage

When assisted by customized retrieval corpus, the model can interact with the environment through LangChain framework

Detailed usage is comming soon ...

Benchmark

We show the main results of DiagRL and compare it to vanilla models and other baselines. For more details, please refer to our paper.

Table: Main diagnosis performance.
We calculate top-1 and top-5 accuracy among common and rare disease diagnosis datasets and compare our \ModelName{} with other representative models. “Env” means we allow the model to use our proposed environment as assistance. All results are shown in percentage.