{"repo": "chaoql/chip2024-xiansheng", "files": {"main.py": {"49d567964305c30780fea6fcd222792596548d4d": {"t": 1785674173.7028763, "n": 13367, "h": "2f263f97cdaa9b6f", "c": "import time\nfrom dotenv import load_dotenv, find_dotenv\nfrom dotenv import dotenv_values\nfrom llama_index.embeddings.huggingface import HuggingFaceEmbedding\nfrom llama_index.core.response_synthesizers.type import ResponseMode\nfrom llama_index.core import Settings\nfrom custom.glmfz import ChatGLM\nfrom custom.prompt import qa_clinical_info_prompt_tmpl_str, qa_core_mechanism_prompt_tmpl_str, \\\n    qa_syndrome_infer_prompt_tmpl_str, qa_clinical_exp_prompt_tmpl_str, new_qa_core_mechanism_prompt_tmpl_str, \\\n    all_clinical_info_prompt_tmpl_str\nfrom llama_index.core import PromptTemplate\nfrom pipeline import rag\nimport pipeline.readData as readData\nimport tools\nimport ensemble\nfrom tqdm import *\nimport os\nimport warnings\nfrom llama_index.llms.dashscope import DashScope, DashScopeGenerationModels\n\nwarnings.filterwarnings('ignore')\nif __name__ == '__main__':\n    _ = load_dotenv(find_dotenv())  # 导入环境\n    config = dotenv_values(\".env\")\n    # 加载嵌入模型\n    Settings.embed_model = HuggingFaceEmbedding(\n        model_name=\"BAAI/bge-large-zh-v1.5\",\n        cache_folder=\"./BAAI/\",\n        embed_batch_size=128,\n        local_files_only=True,  # 仅加载本地模型，不尝试下载\n        device=\"cuda\",\n    )\n    # Settings.embed_model = HuggingFaceEmbedding(\n    #     model_name=\"lier007/xiaobu-embedding-v2\",\n    #     cache_folder=\"BAAI/\",\n    #     embed_batch_size=128,\n    #     local_files_only=True,  # 仅加载本地模型，不尝试下载\n    #     device=\"cuda\",\n    # )\n\n    # 加载大模型\n    # Settings.llm = DashScope(\n    #     model_name=DashScopeGenerationModels.QWEN_TURBO, api_key=config[\"DASHSCOPE_API_KEY\"], max_tokens=1024\n    # )\n    Settings.llm = ChatGLM(\n        api_key=config[\"GLM_KEY\"],\n        model=\"glm-4-plus\",\n        api_base=\"https://open.bigmodel.cn/api/paas/v4/\",\n        is_chat_model=True,\n    )\n\n    # 设置参数\n    with_hyde = False  # 是否采用假设文档\n    persist_dir = \"storeQ\"  # 向量存储地址\n    hybrid_search = False  # 是否采用混合检索\n    with_LLMrerank = True\n    top_k = 5\n    rerank_top_k = 2\n    r_nums = 10  # 循环次数\n    response_mode = ResponseMode.TREE_SUMMARIZE\n    submitPath = \"final_submit/sub_5.txt\"\n    now_best_submitPath = \"final_submit/sub_8.txt\"\n    A_file = \"rdata/B榜.json\"\n    train_file = \"rdata/train.json\"\n    infoList = readData.read_AJson(A_file)\n    if hybrid_search:\n        index1, nodes1 = rag.load_hybrid_data([\"rdata/trainTask1.json\"], \"final_store/hybrid_store01\")\n        # index2, nodes2 = rag.load_data([\"data/bjtrain02ExtraData.json\"], \"store/eval02_2_extra\")\n        index2, nodes2 = rag.load_hybrid_data([\"rdata/trainTask2WOoptions2.json\"], \"final_store/hybrid_store02_1\")\n        index2R, nodes2R = rag.load_hybrid_data([\"rdata/trainTask2WOoptions3.json\"], \"final_store/hybrid_store02_2\")\n        index3, nodes3 = rag.load_hybrid_data(\n            [\"rdata/zntrain3.json\", \"rdata/trainTask3WOoptions.json\", \"rdata/trainExtraDataTask2.json\",\n             \"data/bjtrain02ExtraData.json\"], \"final_store/hybrid_store03_2\")\n        index4, nodes4 = rag.load_json_text_hybrid_data(\n            [\"rdata/trainTask4.json\"], [\"rdata/中医诊断学2.txt\", \"rdata/中医内科学.txt\", \"rdata/中医基础理论.txt\"],\n            \"final_store/hybrid_store04\")\n    else:\n        index1, nodes1 = rag.load_data([\"rdata/trainTask1.json\"], \"final_store/store01\")\n        index2, nodes2 = rag.load_data([\"rdata/trainTask2WOoptions2.json\"], \"final_store/store02_1\")\n        index2R, nodes2R = rag.load_data([\"rdata/trainTask2WOoptions3.json\"], \"final_store/store02_2\")\n        index3, nodes3 = rag.load_data(\n            [\"rdata/zntrain3.json\", \"rdata/trainTask3WOoptions.json\", \"rdata/trainExtraDataTask2.json\",\n             \"rdata/bjtrain02ExtraData.json\"], \"final_store/store03_2\")\n        index4, nodes4 = rag.load_json_text_data(\n            [\"rdata/trainTask4.json\"], [\"rdata/中医诊断学2.txt\", \"rdata/中医内科学.txt\", \"rdata/中医基础理论.txt\"],\n            \"final_store/store04\")\n\n    subResult = []\n    start_time = time.time()\n    print(start_time)\n    resnum = 0\n    if os.path.exists(submitPath):\n        with open(submitPath, 'r', encoding=\"utf-8\", errors=\"ignore\") as file:\n            lines = file.readlines()\n            resnum = len(lines)\n\n    # 读取历史最佳答案，某个病例生成错误后直接替换为历史最佳答案，防止报错打断程序运行\n    # 主要是因为LLM_reranker报错率太高了\n    task1_ans = []\n    task2_ans = []\n    task3_ans = []\n    task4_ans = []\n    with open(now_best_submitPath, 'r', encoding=\"utf-8\", errors=\"ignore\") as file:\n        temp = file.readlines()\n    for i in temp:\n        task1_ans.append(i.split(\"@\")[1].rstrip())\n        task2_ans.append(i.split(\"@\")[2].rstrip())\n        task3_ans.append(i.split(\"@\")[3].rstrip())\n        task4_ans.append(i.split(\"@\")[4].rstrip())\n\n    for i, info in tqdm(enumerate(infoList), total=len(infoList)):\n        if i < resnum:  # 接续生成\n            continue\n\n        global syndrome_answer_str, mechanism_str, syndrome_answer, mechanism_answer\n        case_id = info[\"案例编号\"]\n        clinical_info = info[\"临床资料\"]\n        mechanism_options = info[\"病机选项\"]\n        syndrome_options = info[\"证候选项\"]\n\n        # 直接读取各题当前最佳答案，调试用\n        # core_clinical_info_str = task1_ans[i]\n        # tools.printf(f\"core_clinical_info_str:{core_clinical_info_str}\")\n        #\n        # mechanism_answer = task2_ans[i].split(\";\")\n        # mechanism_str = tools.extract_core_mechanism(mechanism_options, mechanism_answer)\n        # tools.printf(mechanism_str)\n        #\n        # syndrome_answer = task3_ans[i].split(\";\")\n        # syndrome_answer_str = tools.extract_core_mechanism(syndrome_options, syndrome_answer)\n        # tools.printf(syndrome_answer_str)\n        #\n        # task4 = task4_ans[i]\n        # clinical_experience_str = task4.split(\"辨证\")[0].replace(\"临证体会：\", \"\")\n        # diagnosis_str = task4.split(\"辨证\")[1].replace(\"辨证：\", \"\").replace(\"：\", \"\")\n\n        # task1\n        qa_clinical_info_prompt_tmpl_str_temp = qa_clinical_info_prompt_tmpl_str.format(context_str=\"{context_str}\",\n                                                                                        clinical_info=clinical_info)\n        prompt1 = PromptTemplate(qa_clinical_info_prompt_tmpl_str_temp)\n        rag_query_engine = rag.Build_query_engine(with_LLMrerank, rerank_top_k, index1, top_k, response_mode,\n                                                  hybrid_search, nodes1, with_hyde, qa_prompt_tmpl=prompt1)\n        response = rag_query_engine.query(clinical_info)\n\n        core_clinical_info = str(response)\n        tools.printf(f\"core_clinical_info:{core_clinical_info}\")\n\n        all_clinical_info_prompt_tmpl_str_temp = all_clinical_info_prompt_tmpl_str.format(clinical_info=clinical_info)\n        prompt1_1 = PromptTemplate(all_clinical_info_prompt_tmpl_str_temp)\n        rag_query_engine = rag.Build_query_engine(with_LLMrerank, rerank_top_k, index1, top_k, response_mode,\n                                                  hybrid_search, nodes1, with_hyde, qa_prompt_tmpl=prompt1_1)\n        response1_1 = rag_query_engine.query(clinical_info)\n        core_clinical_info1_1 = str(response1_1)\n        core_clinical_info_list = core_clinical_info.split(\";\")\n        print(len(core_clinical_info_list))\n\n        core_clinical_info1_1_list = core_clinical_info1_1.split(\";\")\n        core_clinical_info_result = []\n        core_clinical_info_str = \"\"\n        for info in core_clinical_info_list:\n            if info:\n                core_clinical_info1_1_list.append(info)\n        for info in core_clinical_info1_1_list:\n            if info and info not in core_clinical_info_result:\n                core_clinical_info_result.append(info)\n                core_clinical_info_str += info + \";\"\n        print(len(core_clinical_info_result))\n\n        for j in range(r_nums):\n            # task2\n            try:\n                if j == 0:\n                    qa_core_mechanism_prompt_tmpl_str_temp = qa_core_mechanism_prompt_tmpl_str.format(\n                        context_str=\"{context_str}\",\n                        core_clinical_info=core_clinical_info,\n                        mechanism_options=mechanism_options,\n                    )\n                    prompt2 = PromptTemplate(qa_core_mechanism_prompt_tmpl_str_temp)\n                    rag_query_engine = rag.Build_query_engine(with_LLMrerank, rerank_top_k, index2, top_k,\n                                                              response_mode, hybrid_search, nodes2, with_hyde,\n                                                              qa_prompt_tmpl=prompt2)\n                    response = rag_query_engine.query(core_clinical_info)  #\n                else:\n                    syndrome_answer_str = tools.extract_core_mechanism(syndrome_options, syndrome_answer)\n                    qa_core_mechanism_prompt_tmpl_str_temp = new_qa_core_mechanism_prompt_tmpl_str.format(\n                        context_str=\"{context_str}\",\n                        core_clinical_info=core_clinical_info,\n                        mechanism_options=mechanism_options,\n                        syndrome_answer_str=syndrome_answer_str\n                    )\n                    prompt2 = PromptTemplate(qa_core_mechanism_prompt_tmpl_str_temp)\n                    rag_query_engine = rag.Build_query_engine(with_LLMrerank, rerank_top_k, index2R, top_k,\n                                                              response_mode,\n                                                              hybrid_search, nodes2R, with_hyde, qa_prompt_tmpl=prompt2)\n                    response = rag_query_engine.query(syndrome_answer_str)\n\n                mechanism_answer = str(response)\n                tools.printf(f\"mechanism_answer:{mechanism_answer}\")\n                mechanism_answer = tools.select_answers_parse(mechanism_answer, \"病机答案\")\n                tools.printf(mechanism_answer)\n                mechanism_str = tools.extract_core_mechanism(mechanism_options, mechanism_answer)\n                tools.printf(mechanism_str)\n            except Exception as e:\n                mechanism_answer = task2_ans[i].split(\";\")\n                mechanism_str = tools.extract_core_mechanism(mechanism_options, mechanism_answer)\n                tools.printf(mechanism_str)\n\n            # task3\n            qa_syndrome_infer_prompt_tmpl_str_temp = qa_syndrome_infer_prompt_tmpl_str.format(\n                context_str=\"{context_str}\",\n                core_clinical_info=core_clinical_info,\n                syndrome_options=syndrome_options,\n                mechanism_str=mechanism_str)\n            prompt3 = PromptTemplate(qa_syndrome_infer_prompt_tmpl_str_temp)\n            rag_query_engine = rag.Build_query_engine(with_LLMrerank, rerank_top_k, index3, top_k, response_mode,\n                                                      hybrid_search, nodes3, with_hyde, qa_prompt_tmpl=prompt3)\n            try:\n                response = rag_query_engine.query(mechanism_str + core_clinical_info)\n                for node in response.source_nodes:\n                    tools.printf(node.text)\n                syndrome_answer = str(response)\n                tools.printf(f\"syndrome_answer:{syndrome_answer}\")\n                syndrome_answer = tools.select_answers_parse(syndrome_answer, \"证候答案\")\n                tools.printf(syndrome_answer)\n                syndrome_answer_str = tools.extract_core_mechanism(syndrome_options, syndrome_answer)\n                tools.printf(syndrome_answer_str)\n            except Exception as e:\n                syndrome_answer = task3_ans[i].split(\";\")\n                syndrome_answer_str = tools.extract_core_mechanism(syndrome_options, syndrome_answer)\n                tools.printf(syndrome_answer_str)\n\n        # task4\n        qa_clinical_exp_prompt_tmpl_str_temp = qa_clinical_exp_prompt_tmpl_str.format(context_str=\"{context_str}\",\n                                                                                      clinical_info=clinical_info,\n                                                                                      mechanism_str=mechanism_str)\n        prompt4 = PromptTemplate(qa_clinical_exp_prompt_tmpl_str_temp)\n        rag_query_engine = rag.Build_query_engine(with_LLMrerank, rerank_top_k, index4, top_k, response_mode,\n                                                  hybrid_search, nodes4, with_hyde, qa_prompt_tmpl=prompt4)\n        try:\n            response = rag_query_engine.query(mechanism_str)\n            clinical_experience_str = str(response).replace(\"\\n\", \"\")\n            tools.printf(f\"clinical_experience_str:{clinical_experience_str}\")\n            diagnosis = syndrome_answer_str.split(\";\")\n            diagnosis_str = \"\"\n            for d in diagnosis:\n                if d:\n                    diagnosis_str += d + '，'\n            diagnosis_str = diagnosis_str[:-1]\n            tools.printf(diagnosis_str)\n        except Exception as e:\n            task4 = task4_ans[i]\n            clinical_experience_str = task4.split(\"辨证\")[0].replace(\"临证体会：\", \"\")\n            diagnosis_str = task4.split(\"辨证\")[1].replace(\"辨证：\", \"\").replace(\"：\", \"\")\n\n        resultAll = tools.printInfo(case_id, core_clinical_info_str, mechanism_answer, syndrome_answer,\n                                    clinical_experience_str, diagnosis_str)\n        tools.save_now(resultAll, submitPath)\n        subResult.append(resultAll)\n    end_time = time.time()\n    print(f\"start_time-end_time:{(end_time - start_time) / 60.0}min\")\n    tools.saveTxt(\"submit/subTemp.txt\", subResult)\n\n    ensemble.enseTask1(submitPath=submitPath, llmsubmit_path=submitPath, A_path=A_file)  # 第一问补缺\n"}}}}