中文 | English

TCMChat: A Generative Large Language Model for Traditional Chinese Medicine

News

[November 1, 2024] We have fully open-sourced the model weights and training dataset on Huggingface.
[May 17, 2024] Open source model weight on HuggingFace.

Application

Install

git clone https://github.com/ZJUFanLab/TCMChat
cd TCMChat

Create a conda environment

conda create -n baichuan2 python=3.10 -y

First install the dependency package. python environment 3.10+ is recommended.

pip install -r requirements.txt

Weights download

• TCMChat: QA and recommendation of TCM knowledge based on baichuan2-7B-Chat.

Inference

Command line

python cli_infer.py \
--model_name_or_path /your/model/path \
--model_type  chat

Web demo

python gradio_demo.py

We provide an online tool：https://xomics.com.cn/tcmchat

Retrain

Dataset Download

• Pretrain dataset
• SFT dataset
• Benchmark dataset

Note: Before performing pre-training, fine-tuning, and inference, please modify the relevant paths for your model, data, and other related files.

Pre-training

## Slurm cluster
sbatch scripts/pretrain/baichuan2_7b_chat.slurm
## or
bash scripts/pretrain/baichuan2_7b_chat.sh

Fine-tuning

## Slurm cluster
sbatch scripts/sft/baichuan2_7b_chat.slurm
## or
bash scripts/sft/baichuan2_7b_chat.sh

Training details

Please refer to the experimental section of the paper for instructions.

Benchmark evaluation

Choice Question

python evaluation/choices_evaluate/eval.py   --model_path_or_name /your/model/path --model_name  baichuan2-7b-chat --few_shot -sz herb --dev_file_path evaluation/resources/choice/single/tcm-herb_dev.csv --val_file_path evaluation/resources/choice/single/choice_herb_500.csv --log_dir logs/choices

Reading Comprehension

python infers/baichuan_infer.py \
--model_name_or_path /your/model/path / \
--model_type chat \
--save_path /your/save/data/path \
--data_path /your/data/path
##BertScore
python evaluation/question_rouge_bleu.py/question_bert_score.py
## BLEU METEOR
python evaluation/question_rouge_bleu.py/open_question_bleu.py
## ROUGE-x
python evaluation/question_rouge_bleu.py/open_question_rouge.py

Entity Extraction

python infers/baichuan_infer.py \
--model_name_or_path /your/model/path / \
--model_type chat \
--save_path /your/save/data/path \
--data_path /your/data/path

python evaluation/ner_evaluate/tcm_entity_recognition.py

Medical Case Diagnosis

python infers/baichuan_infer.py \
--model_name_or_path /your/model/path / \
--model_type chat \
--save_path /your/save/data/path \
--data_path /your/data/path

python evaluation/acc_evaluate/extract_syndrome.py

Herb or Formula Recommendation

python infers/baichuan_infer.py \
--model_name_or_path /your/model/path / \
--model_type chat \
--save_path /your/save/data/path \
--data_path /your/data/path

python evaluation/recommend_evaluate/mrr_ndcg_p_r.py

ADMET Prediction

Regression

python infers/baichuan_infer.py \
--model_name_or_path /your/model/path / \
--model_type chat \
--save_path /your/save/data/path \
--data_path /your/data/path

python evaluation/admet_evaluate/rmse_mae_mse.py

Classification

python infers/baichuan_infer.py \
--model_name_or_path /your/model/path / \
--model_type chat \
--save_path /your/save/data/path \
--data_path /your/data/path

python evaluation/admet_evaluate/acc_recall_f1.py

Acknowledgments

1. Thanks to technical support from the High-Performance Computing Cluster of Zhejiang University Innovation Center of Yangtze River Delta.
2. Thanks to Baichuan Intelligence for providing open-source large language models