DeepRare/eval.py at main · MAGIC-AI4Med/DeepRare · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import json
import argparse
import re
import numpy as np

from api.interface import Openai_api, gemini_api, deepseek_api, claude_api

# define the LLM handler based on the selected model
class LLM_handler:
    def __init__(self, args):

        if args.model == "openai":
            self.handler = Openai_api(args.openai_apikey, args.openai_model)
        elif args.model == "gemini":
            self.handler = gemini_api(args.gemini_apikey, args.gemini_model)
        elif args.model == "deepseek":
            self.handler = deepseek_api(args.deepseek_apikey, args.deepseek_model)
        elif args.model == "claude":
            self.handler = claude_api(args.claude_apikey, args.claude_model)
        else:
            raise ValueError("Invalid model name.")

def diagnosis_evaluate(predict_diagnosis, golden_diagnosis, handler):
    if predict_diagnosis is None:
        raise Exception("Predict diagnosis is None")

    predict_diagnosis = predict_diagnosis.replace("\n\n\n", "")

    system_prompt = "You are a specialist in the field of rare diseases."
    prompt = 'I will now give you five predicted diseases if the predicted diagnosis is in the standard diagnosis. Please output the predicted rank, otherwise output "No", only output "No" or "1-5" numbers, if the predicted disease has multiple conditions, only output the top rank. Output only "No" or one number, no additional output.'
    prompt += f'Predicted diseases: {predict_diagnosis}\n'
    prompt += f'Standard diagnosis: {golden_diagnosis}\n'

    print("Begin evaluation.....")
    rank = handler.get_completion(system_prompt, prompt)
    rank = rank.replace("\n", "")

    return rank


def main():

    # set up the arguments

    parser = argparse.ArgumentParser()
    parser.add_argument("--results_folder", type=str, default="./result_new/HMS/gpt-4o")
    # model
    parser.add_argument('--model', type=str, default="openai", choices=["openai", "gemini", "deepseek", "claude"])

    # API keys
    parser.add_argument('--openai_apikey', type=str, default='')
    parser.add_argument('--openai_model', type=str, default='gpt-4o', choices=['gpt-4o', 'gpt-4o-mini', 'o1', 'o3-mini', 'o1-mini'])
    parser.add_argument('--gemini_apikey', type=str, default='')
    parser.add_argument('--gemini_model', type=str, default='', choices=['gemini-2.0-pro-exp', 'gemini-2.0-flash-exp', 'gemini-2.0-flash', 'gemini-1.5-pro', 'gemini-1.5-flash-8b', 'gemini-1.5-flash'])
    parser.add_argument('--claude_apikey', type=str, default='')
    parser.add_argument('--claude_model', type=str, default='', choices=['claude-3-7-sonnet-20250219', 'claude-3-7-sonnet-thinking'])
    parser.add_argument('--deepseek_apikey', type=str, default='') # pip install -U 'volcengine-python-sdk[ark]'
    parser.add_argument('--deepseek_model', type=str, default="deepseek-r1-250120")
    parser.add_argument('--uptodate_pwd', type=str, default='')
    parser.add_argument('--uptodate_user', type=str, default='')
    parser.add_argument('--google_api', type=str, default='')
    parser.add_argument('--search_engine_id', type=str, default='')

    args = parser.parse_args()

    results_folder = args.results_folder
    print(f'evaluating {results_folder.split("/")[-2]}, generated by {results_folder.split("/")[-1]}')

    # total number of file
    len_files = len(os.listdir(results_folder))
    print(f"total number of files: {len_files}")

    CNT = 0
    metric = {}
    recall_top_k = []

    # Set up the LLM Model
    handler = LLM_handler(args).handler

    for file in os.listdir(results_folder):
        file = os.path.join(results_folder, file)
        try:
            res = json.load(open(file, "r", encoding="utf-8-sig"))
        except:
            print(file)
            os.remove(file)
            continue

        if "predict_rank" not in res:
            # final_diagnois  first_round_result
            if res["final_diagnois"] is None:
                print(file)
                os.remove(file)
                continue
            predict_rank = diagnosis_evaluate(res["final_diagnois"], res["golden_diagnosis"], handler)
            # predict_rank_1
            res["predict_rank"] = predict_rank
            json.dump(res, open(file, "w", encoding="utf-8-sig"), indent=4, ensure_ascii=False)
        else:
            predict_rank = res["predict_rank"]


        if predict_rank not in ["否", "1", "2", "3", "4", "5", "No"]:
            print(file)
            CNT += 1

        if "否" in predict_rank or "No" in predict_rank:
            recall_top_k.append(11)
        else:
            pattern = r'\b(?:10|[1-9])\b'
            predict_rank = re.findall(pattern, predict_rank)
            if predict_rank[0] not in ["1", "2", "3", "4", "5"]:
                res["predict_rank"] = None
                print(file)
                raise Exception("predict_rank error")
            predict_rank = predict_rank[0]
            recall_top_k.append(int(predict_rank))


    metric['recall_top_1'] = len([i for i in recall_top_k if i <= 1]) / len(recall_top_k)
    metric['recall_top_3'] = len([i for i in recall_top_k if i <= 3]) / len(recall_top_k)
    metric['recall_top_5'] = len([i for i in recall_top_k if i <= 5]) / len(recall_top_k)
    metric['medain_rank'] = np.median(recall_top_k)
    print(metric)
    print("predict_rank error: ", CNT)
    print("Test dataset:", results_folder)


if __name__ == "__main__":
    main()