1 line
14 KiB
JSON
1 line
14 KiB
JSON
{"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"}}}} |