This work builds on top of the datasets released with https://github.com/glee4810/EHRSQL. Use the instructions mentioned in the link to download the database and the datasets.

Set up conda environment using (python>=3.10):

pip install -r requirements.txt

Two models are used in this work - Mistral and Llama 2. Download the models from the link below: Mistral7B : https://huggingface.co/mistralai/Mistral-7B-v0.1 Llama27B : https://huggingface.co/meta-llama/Llama-2-7b

Both the bash scripts start a FSDP process. Set the --nnodes(number of nodes) and --nproc-per-node(number of GPUs) in run.sh and run_multi.sh as per the compute available. The training to reproduce the results can be performed using the following command:

Note: Pass the path of the model, the learning rate, the dataset path, output directory, model name and dataset name as parameters to the script. The training set for both the datasets are assumed to be named 'train_mimic.json' , 'train_eicu.json'. Change the name in the bash command according to the file name.

bash run.sh https://huggingface.co/mistralai/Mistral-7B-v0.1 1e-5 train_mimic.json ./mistral_main_prompt_mimic/ mistral mimic

bash run.sh https://huggingface.co/mistralai/Mistral-7B-v0.1 1e-5 train_eicu.json ./mistral_main_prompt_eicu/ mistral eicu

bash run.sh https://huggingface.co/meta-llama/Llama-2-7b 1e-5 train_mimic.json ./llama2_main_prompt_mimic/ llama2 mimic

bash run.sh https://huggingface.co/meta-llama/Llama-2-7b 1e-5 train_eicu.json ./llama2_main_prompt_eicu/ llama2 eicu

The same steps from the previous section have to be followed for multi funciton training.

bash run_multi.sh https://huggingface.co/mistralai/Mistral-7B-v0.1 1e-5 train_multistep.json ./mistral_multitool_prompt_mimic/ mistral mimic

Go to the inference folder for instructions to perform inference on the development set.

For evalauting the code generated by the finetuned models we follow the same method that is used in EHRSQL. Refer to the link and follow the same procedure for evaluation, or follow the steps below. Install the func-timeout if you have not.

pip install func-timeout

Run the following. The generated file will be in the same directory as the json input file

./eval/run_abstain_entropy.sh {json file from the inference step} {formatted output filename}

Prepare the database in the database directory and run the following to get the score. The {score output file} variable is not used, the score will be printed on terminal.

./eval/run_ehrsql_eval.sh {original validation json} {formatted output filename} {score output file} {dataset name}