1 line
31 KiB
JSON
1 line
31 KiB
JSON
{"repo": "bubble-bubble77/ieea", "files": {"LLMReason/A_DBP_part3.py": {"86cd89f4f9d5f7dfb275157c5bc3d5b4df748516": {"t": 1785768903.8773036, "n": 30103, "h": "ccaaafd77bdafbe3", "c": " \nfrom DBP_retrieval import *\nfrom prompt import *\nfrom openai import OpenAI\nimport os\nimport pytz\nfrom datetime import datetime\nfrom Prepare_EMB.DBP15K.util import *\nfrom Prepare_EMB.DBP15K.dataset import DBPDataset\nfrom Prepare_EMB.DBP15K.model import ST_Encoder_Module, Loss_Module\nfrom torch.optim.lr_scheduler import ExponentialLR\nfrom tqdm import trange\nimport time\n\ntimezone = pytz.timezone('Asia/Shanghai')\ncurrent_date = datetime.now(timezone).strftime(\"%m%d_%H%M\")\n\nclass DBP_Part_3():\n def __init__(self, dataset, side=None, equal_rule= None, align_rule=None, rule3=None):\n self.dataset = dataset\n print('-------------')\n print(current_date)\n print(self.dataset)\n self.side = side\n self.re = Retrivel(dataset)\n self.en1_list, self.en2_list = self.ent1_ent2_load()\n self.seed_dict = self.seed_load()\n\n self.EntN_Attr = self.re.original_attr\n self.Id_Attr = self.re.id_Attr\n self.Id_Name = self.re.id_o_Name\n self.index1, self.index2 = self.re.index_1, self.re.index_2\n self.Name_id = self.Name_Id()\n self.Id_rel_dic = self.Rel_list_to_dict()\n self.trip_1_dic, self.trip_2_dic = self.Triple_load()\n self.trip_dic = self.trip_1_dic | self.trip_2_dic\n\n self.LLM_load()\n # self.LLM_type='qwen3-max'\n self.LLM_type = 'deepseek-v3.2-exp'\n # self.LLM_type = 'deepseek-v3.1'\n\n self.equal = equal_rule\n self.align = align_rule\n self.rule3 = rule3\n\n # rule1\n if self.equal is not None:\n self.euql_ent_dict = self.equal_ent_load()\n print(\"equal_ent has loaded..\")\n\n # rule2\n if self.align is not None:\n self.align_ent_dict = self.complete_aligned_neighbor()\n print(\"align_neighbor has loaded..\")\n with open('../Results/DBP.txt', 'a') as f:\n f.write(\"\\n\\n\")\n f.write(f\"{current_date}\\n\")\n f.write(f\"align neighbor num of {self.dataset}, side:{self.side}: {len(self.align_ent_dict)}\\n\")\n f.write(f\"align neighbor ratio of {self.dataset}, side:{self.side}: {len(self.align_ent_dict)/len(self.en1_list)}\\n\")\n\n # rule3\n if self.equal is not None and self.align is not None:\n self.equal_AND_align = 0\n for ent in self.align_ent_dict.keys():\n if ent in self.euql_ent_dict.keys() and self.euql_ent_dict[ent] != []:\n self.equal_AND_align += 1\n with open('../Results/DBP.txt', 'a') as f:\n f.write(f\"equal AND align num of {self.dataset}, side:{self.side}: {self.equal_AND_align}\\n\")\n f.write(f\"equal AND align ratio of {self.dataset}, side:{self.side}: {self.equal_AND_align/len(self.en1_list)}\\n\")\n \n\n self.device = 'cuda:0'\n if self.equal is not None and self.align is not None and self.rule3 is not None:\n self.structure_learn()\n else:\n self.es_score = self.sim_score_load()\n\n\n def seed_load(self):\n filepath = f'../data/DBP15K/DBP_1/'\n filename = os.path.join(filepath, self.dataset, 'ill_ent_ids')\n seed_pair = {}\n with open(filename, 'r') as f:\n for line in f:\n try:\n e1, e2 = line.split()\n # e1和e2之间可以互查\n seed_pair[int(e1)] = int(e2)\n seed_pair[int(e2)] = int(e1)\n except:\n continue\n return seed_pair \n \n def Name_Id(self):\n Name_id = {}\n with open(f'../data//DBP15K/DBP_1/{self.dataset}/ent_ids_1', 'r') as f:\n for line in f:\n id_1, ol_name = line.strip().split('\\t')\n Name_id[ol_name.split('/')[-1]] = int(id_1)\n\n with open(f'../data//DBP15K/DBP_1/{self.dataset}/ent_ids_2', 'r') as f:\n for line in f:\n id_2, ol_name = line.strip().split('\\t')\n Name_id[ol_name.split('/')[-1]] = int(id_2)\n return Name_id\n \n def Triple_load(self):\n triple1_path = f\"../data/DBP15K/DBP_1/{self.dataset}/triples_1\"\n triple2_path = f\"../data/DBP15K/DBP_1/{self.dataset}/triples_2\"\n\n def trip_load(path, index_dict):\n trip_dic = {}\n with open(path,'r') as f:\n for line in f:\n ent1, _, ent2 = line.strip().split('\\t')\n if ent1 not in trip_dic.keys():\n trip_dic[int(ent1)] = [int(ent2)]\n else:\n trip_dic[int(ent1)].append(int(ent2))\n return trip_dic\n \n trip_1_dic = trip_load(triple1_path, self.index1)\n trip_2_dic = trip_load(triple2_path, self.index2)\n return trip_1_dic, trip_2_dic\n \n def Rel_list_to_dict(self):\n s_d, t_d = self.dataset.split('_')\n rel_path_1 = f\"../data/DBP15K/dbpdata/DBP15k/{self.dataset}/{s_d}_rel_triples\"\n rel_path_2 = f\"../data/DBP15K/dbpdata/DBP15k/{self.dataset}/{t_d}_rel_triples\"\n\n def rel_ent2_load(path):\n Id_rel_dic = {}\n with open(path, 'r') as f:\n for line in f:\n line = line.strip()\n ent1 = line.split('\\t')[0].split('/')[-1]\n rel = line.split('\\t')[1].split('/')[-1]\n ent2 = line.split('\\t')[2].split('/')[-1] \n\n ent1_id = self.Name_id[ent1]\n if ent1_id not in Id_rel_dic.keys():\n Id_rel_dic[ent1_id] = {ent2: rel} # {ent2_name:rel_Name}\n else:\n Id_rel_dic[ent1_id][ent2] = rel\n return Id_rel_dic\n \n s_d_rel_list = rel_ent2_load(rel_path_1) #[ent_id:{ent_name: rel_name}]\n t_d_rel_list = rel_ent2_load(rel_path_2)\n\n rel_list = s_d_rel_list | t_d_rel_list\n\n return rel_list\n\n\n def complete_aligned_neighbor(self):\n def rule2(ent_list):\n rule2_dict = {}\n for ent1 in ent_list:\n candi_tuple_list = []\n # A->[rel->B--B'->[(r1, b1)(r2, b2)(r3, b3)...]]\n # ent1的邻居\n if ent1 in self.trip_dic.keys():\n neighbor_of_ent = self.trip_dic[ent1]\n else:\n neighbor_of_ent = []\n has_aligned_dict = {}\n has_aligned_set = set()\n for nei in neighbor_of_ent: \n if nei in self.seed_dict.keys():\n rel_propm_list = [] \n aligned_nei = self.seed_dict[nei]\n has_aligned_dict[nei] = aligned_nei\n if aligned_nei in self.trip_dic.keys():\n probli_neigbor = self.trip_dic[aligned_nei] # aligned_nei的邻居\n has_aligned_set.add(nei)\n for pn in probli_neigbor:\n pn_name = self.Id_Name[pn]\n if aligned_nei in self.Id_rel_dic.keys() and pn_name in self.Id_rel_dic[aligned_nei].keys():\n rel_name = self.Id_rel_dic[aligned_nei][pn_name]\n rel_propm_list.append((rel_name, pn, pn_name))\n if len(rel_propm_list):\n candi_tuple_list.append((rel_name, nei, aligned_nei, rel_propm_list)) \n if len(candi_tuple_list):\n rule2_dict[ent1] = candi_tuple_list\n return rule2_dict\n\n e1_complete = rule2(self.en1_list)\n ent_complete = e1_complete\n return ent_complete\n\n\n def ent1_ent2_load(self):\n Imcomplete_data_save_path = f'../data/DBP/IC_data/'\n Ic_e1_path = os.path.join(Imcomplete_data_save_path, f'{self.dataset}_{self.side}_top_Ic_e1.txt')\n Ic_e2_path = os.path.join(Imcomplete_data_save_path, f'{self.dataset}_{self.side}_Ic_e2.txt')\n\n ent1_list = []\n with open(Ic_e1_path, 'r') as f:\n for line in f:\n ent1_list.append(int(line.strip()))\n print(f\"ent1 读取完成，共 {len(ent1_list)} 条记录\")\n\n ent2_list = []\n with open(Ic_e2_path, 'r') as f:\n for line in f:\n ent2_list.append(int(line.strip()))\n print(f\"ent2 读取完成，共 {len(ent2_list)} 条记录\")\n\n return ent1_list, ent2_list\n\n def structure_learn(self):\n print('Graph retrieval....')\n\n\n self.G_dataset = DBPDataset('../data/DBP15K/DBP_1/{}'.format(self.dataset),\n device=self.device)\n self.train_pair = self.G_dataset.train_pair\n\n train_pair = [(ent1, self.seed_dict[ent1]) for ent1 in self.en1_list]\n import random\n random.shuffle(train_pair) \n sample_size = int(len(train_pair) * 0.2)\n random_sample = train_pair[:sample_size] \n self.train_pair.extend(random_sample)\n\n self.train_pair = np.array(self.train_pair)\n self.train_pair = torch.tensor(self.train_pair).to(self.device)\n\n self.test_pair = [(ent1, self.seed_dict[ent1]) for ent1 in self.en1_list]\n self.test_pair = np.array(self.test_pair)\n self.test_pair = torch.tensor(self.test_pair)\n\n self.ent_adj, self.rel_adj, self.node_size, self.rel_size, \\\n self.adj_list, self.r_index, self.r_val, self.triple_size, self.adj = self.G_dataset.reconstruct_search(None,\n None,\n self.G_dataset.kg1,\n self.G_dataset.kg2,\n new=True)\n self.structure_encoder = ST_Encoder_Module(\n node_hidden=128,\n rel_hidden=128,\n node_size=self.node_size,\n rel_size=self.rel_size,\n device=self.device,\n dropout_rate=0.1,\n depth=4).to(self.device)\n\n self.loss_model = Loss_Module(node_size=self.node_size, gamma=2).to(self.device)\n self.lr = 0.005\n self.optimizer = torch.optim.RMSprop(self.structure_encoder.parameters(), lr=self.lr) # 3-9\n self.scheduler = ExponentialLR(self.optimizer, gamma=0.999)\n # train_epoch = 10\n train_epoch = 50\n batch_size = 1024\n for epoch in trange(train_epoch):\n self.structure_encoder.train()\n for i in range(0, len(self.train_pair), batch_size):\n batch_pair = self.train_pair[i:i+batch_size]\n if len(batch_pair) == 0:\n continue\n feature_list = self.structure_encoder(\n self.ent_adj, self.rel_adj, self.node_size,\n self.rel_size, self.adj_list, self.r_index, self.r_val,\n self.triple_size, mask=None)\n loss = self.loss_model(batch_pair, feature_list[0], weight=True)\n loss.backward()\n self.optimizer.step()\n self.optimizer.zero_grad()\n self.scheduler.step()\n self.structure_encoder.eval()\n with torch.no_grad():\n feature_list = self.structure_encoder(\n self.ent_adj.to(self.device), self.rel_adj.to(self.device),\n self.node_size, self.rel_size,\n self.adj_list.to(self.device),\n self.r_index.to(self.device), self.r_val.to(self.device),\n self.triple_size,\n mask=None)\n self.out_feature = feature_list[0].cpu()\n\n def topk_ST_e1(self, k=10):\n gid1 = np.array(self.en1_list)\n gid2 = np.array(self.en2_list)\n \n self.out_feature = self.out_feature / (torch.linalg.norm(self.out_feature, dim=-1, keepdim=True) + 1e-5)\n index_a = torch.LongTensor(gid1)\n index_b = torch.LongTensor(gid2)\n\n Lvec = self.out_feature[index_a]\n Rvec = self.out_feature[index_b]\n Lvec = Lvec / (torch.linalg.norm(Lvec, dim=-1, keepdim=True) + 1e-5)\n Rvec = Rvec / (torch.linalg.norm(Rvec, dim=-1, keepdim=True) + 1e-5)\n A_sim = torch.mm(Lvec, Rvec.t())\n self.topk_st_sim_matrix = A_sim.numpy()\n x_topk_indices = np.argsort(-self.topk_st_sim_matrix, axis=1)[:, :k]\n return x_topk_indices\n \n def sim_score_load(self):\n folder_path = '../data/DBP15K/DBP_fuse'\n\n if self.side is not None:\n file_name = os.path.join(folder_path,\n f'Fuse_{self.dataset}_{self.side}.npy')\n else:\n file_name = os.path.join(folder_path, f'Fuse_{self.dataset}.npy')\n moda_matrix = np.load(file_name)\n m, _ = moda_matrix.shape\n es_score = (moda_matrix - moda_matrix.min()) / (moda_matrix.max() - moda_matrix.min())\n return es_score\n \n def topk_e1(self, k=10):\n ent2_list = [item - self.es_score.shape[0] for item in self.en2_list]\n new_esscore_matrix = self.es_score[np.ix_(self.en1_list, ent2_list)]\n topk_indices = np.argsort(-new_esscore_matrix, axis=1)[:, :k]\n return topk_indices\n \n def equal_ent_load(self):\n equal_rel_path = f'../data/DBP15K/Equal_rel/{self.dataset}_{self.side}.txt'\n equal_dict = {}\n if os.path.exists(equal_rel_path):\n with open(equal_rel_path, 'r') as f:\n for line in f:\n parts = line.strip().split('\\t')\n if len(parts) >= 2:\n ent = parts[0]\n equal_ents = parts[1:] \n equal_dict[int(ent)] = equal_ents\n else:\n print('...')\n self.rel_equal()\n self.equal_ent_load()\n return equal_dict\n \n\n def rel_equal(self):\n global total_tokens, has_equal_rel_num, avg_time\n avg_time = 0\n total_tokens = 0\n has_equal_rel_num = 0\n\n def ent_euqal_entity(ent_list):\n global total_tokens, has_equal_rel_num, avg_time\n ent_equa_list = {}\n for ent1 in ent_list:\n if ent1 in self.Id_rel_dic.keys():\n rel_dic = self.Id_rel_dic[ent1]\n else:\n continue\n rel_if_prompt = rel_if_equal_en(rel_dic, ent1, self.Id_Name[ent1])\n start_time = time.time()\n try:\n completion = self.LLM.chat.completions.create(\n model=self.LLM_type,\n messages=rel_if_prompt)\n llm_response = completion.choices[0].message.content\n LLM_results = self.parse_llm_response(llm_response)\n total_tokens += int(completion.usage.total_tokens)\n\n except Exception as e:\n print(f\"LLM推理出错: {e}\")\n LLM_results = []\n end_time = time.time()\n avg_time += (end_time - start_time)\n\n if LLM_results: # 非空\n has_equal_rel_num += 1\n Id_list = []\n for name in LLM_results:\n if name in self.Name_id.keys():\n Id_list.append(self.Name_id[name]) \n ent_equa_list[ent1] = Id_list\n else:\n ent_equa_list[ent1] = []\n\n return ent_equa_list\n \n en1_equal_list = ent_euqal_entity(self.en1_list)\n # en2_equal_list = ent_euqal_entity(self.en2_list)\n\n # save\n equal_rel_path = f'../data/DBP15K/Equal_rel/{self.dataset}_{self.side}.txt'\n with open(equal_rel_path, 'w') as f:\n for ent, equal_ent_list in en1_equal_list.items():\n line = str(ent) + '\\t' + '\\t'.join(str(item) for item in equal_ent_list) + '\\n'\n f.write(line)\n \n # for ent, equal_ent_list in en2_equal_list.items():\n # line = str(ent) + '\\t' + '\\t'.join(str(item) for item in equal_ent_list) + '\\n'\n # f.write(line)\n\n # 记录推理时间；推理token平均数；真正有等价实体的比例\n record_path = f'/home/wluyao/VEA/Results/DBP.txt'\n with open(record_path, 'a') as f:\n f.write(\"\\n\\n\")\n f.write(f\"{current_date}\\n\")\n f.write(f\"Equal relation generation of {self.dataset}, side:{self.side}\\n\")\n f.write(f\"Total entities :{len(self.en1_list)}\\n\")\n f.write(f\"Has equal relation entities :{has_equal_rel_num}, ratio is {has_equal_rel_num/(len(self.en1_list)):.2f}\\n\")\n f.write(f\"average token: {total_tokens/(len(self.en1_list)):.2f}\\n\")\n f.write(f\"average time: {avg_time/(len(self.en1_list)):.2f}\\n\")\n\n # exit(0)\n\n\n def LLM_load(self):\n self.LLM = OpenAI(\n api_key=\"sk-yourapikeyhere\",\n base_url=\"https://www.dmxapi.cn/v1\"\n )\n \n def parse_llm_response(self, response):\n import re\n import ast\n response = response.strip()\n try:\n list_pattern = r'\\[.*\\]'\n match = re.search(list_pattern, response)\n if match:\n list_str = match.group()\n result_list = ast.literal_eval(list_str)\n if isinstance(result_list, list):\n return result_list\n except:\n pass\n try:\n entity_pattern = r'[\"\\']([^\"\\']+)[\"\\']'\n entities = re.findall(entity_pattern, response)\n return entities\n except:\n return []\n\n \n # run\n def inference(self):\n hard_count = 0\n hit1 = 0\n hit10 = 0\n prompt_tokens = 0\n completion_tokens = 0\n total_tokens = 0\n avg_time = 0\n llm_ratio = len(self.en1_list)\n for epoch in range(0,1):\n if not len(self.en1_list):\n break\n else:\n for i, entity in enumerate(tqdm(self.en1_list)):\n ent1 = entity\n if self.equal is not None and self.align is not None and self.rule3 is not None:\n topk_indices = self.topk_ST_e1(10)\n else:\n topk_indices = self.topk_e1(10)\n top_ent2_list = [self.en2_list[int(it)] for it in topk_indices[i]]\n top5_list = top_ent2_list[:5]\n align_prompt = self.rule1_2_prompt_gen_en(ent1, top5_list, self.equal, self.align)\n messages = [\n {\"role\": \"user\", \"content\": align_prompt}\n ]\n start_time = time.time()\n try:\n completion = self.LLM.chat.completions.create(\n model=self.LLM_type,\n messages=messages)\n llm_response = completion.choices[0].message.content\n LLM_results = self.parse_llm_response(llm_response)\n\n prompt_tokens += int(completion.usage.prompt_tokens)\n completion_tokens += int(completion.usage.completion_tokens)\n total_tokens += int(completion.usage.total_tokens)\n\n except Exception as e:\n print(f\"LLM推理出错: {e}\")\n LLM_results = []\n\n end_time = time.time()\n avg_time += (end_time - start_time)\n if len(LLM_results):\n try:\n hard_count += 1\n if int(LLM_results[0]) == self.seed_dict[ent1]:\n hit1 += 1\n if self.seed_dict[ent1] in top_ent2_list:\n hit10 += 1\n if int(LLM_results[0]) in self.en2_list:\n self.en2_list.remove(int(LLM_results[0]))\n self.en1_list.remove(ent1)\n except:\n continue\n LLM_right_num = hit1\n LLM_num = llm_ratio-len(self.en1_list)\n\n if len(self.en1_list):\n llm_ratio = 1 - (len(self.en1_list) / llm_ratio)\n if self.equal is not None and self.align is not None and self.rule3 is not None:\n topk_indices = self.topk_ST_e1(10)\n else:\n topk_indices = self.topk_e1(10)\n for i in range(len(topk_indices)):\n hard_count += 1\n ent1 = self.en1_list[i]\n top_ent2_list = [self.en2_list[int(it)] for it in topk_indices[i]]\n ent2 = self.seed_dict[ent1]\n if ent2 == top_ent2_list[0]:\n hit1 += 1\n if ent2 in top_ent2_list:\n hit10 += 1\n\n print(f\"hit@1:{hit1},\")\n print(f\"hit@10:{hit10}, \")\n print(f\"hard_count:{hard_count} \")\n\n print(f\"prompt_tokens:{prompt_tokens/hard_count}\")\n print(f\"completion_tokens:{completion_tokens/hard_count}\")\n print(f\"total_tokens:{total_tokens/hard_count}\")\n\n with open('../Results/DBP.txt', 'a') as f:\n f.write(\"\\n\\n\")\n f.write(f\"{current_date}\\n\")\n f.write(f\"equal:{self.equal}; align: {self.align} rule3: {self.rule3} \\n\")\n f.write(f\"Reranker_test_batch result of {self.dataset}, side:{self.side}\\n\")\n f.write(f\"hard_count:{hard_count},\\n\")\n f.write(f\"dev_hit@1:{hit1}, \\n\")\n f.write(f\"hit@10:{hit10},\\n\")\n\n f.write(\"Average Token Usage per Case:\\n\")\n f.write(f\"- Prompt tokens: {prompt_tokens/hard_count:.2f}\\n\")\n f.write(f\"- Completion tokens: {completion_tokens/hard_count:.2f}\\n\")\n f.write(f\"- Total tokens: {total_tokens/hard_count:.2f}\\n\")\n f.write(f\"- Avg time:{avg_time/hard_count:.2f}\\n\")\n f.write(f\"- LLM ratio:{llm_ratio:4f}\\n\\n\")\n\n f.write(f\"[LLM_num : {LLM_num}]\\n\")\n f.write(f\"[LLM_right_num : {LLM_right_num}]\\n\")\n\n\n def rule1_2_prompt_gen_en(self, ent1, top5_list, equal_rule=False, align_rule=False):\n task_desp = f\"\"\"\n Task: Using the provided entity neighbor/attribute information, analyze the alignment possibility between source entity {ent1} and the candidate entity set.\n Requirements:\n 1. Determine whether source entity {ent1} can potentially align with any candidate entity; \n 2. We are not looking for entities that are \"exactly identical,\" but rather for those that are \"most likely to refer to the same thing.\" \n Minor discrepancies in non-essential attributes are allowed, \n as long as the core neiborhood, key attributes exhibit sufficient similarity.\n \n\n Output format: [\"candidate_entity_id\", \"candidate_entity_id\", ...] \n Note: Output only the list without any additional explanations.\n \"\"\"\n source_info = f\"\"\" Source entity information: ID {ent1}, \\n\"\"\"\n\n if self.Id_Name[ent1] != 'Null' and self.Id_Name[ent1] != 'Nul':\n source_info += f\"Name: {self.Id_Name[ent1]}, \\n\"\n if ent1 in self.Id_Name.keys() and self.Id_Name[ent1] in self.EntN_Attr.keys() and self.EntN_Attr[self.Id_Name[ent1]] != [] and self.EntN_Attr[self.Id_Name[ent1]] != '':\n source_info += f\"Attributes: {' ; '.join(self.EntN_Attr[self.Id_Name[ent1]])}, \\n\"\n\n equal_set = set()\n if equal_rule and ent1 in self.euql_ent_dict.keys(): \n for equal_ent in self.euql_ent_dict[ent1]:\n if equal_ent in self.Id_Name.keys() and self.Id_Name[equal_ent] in self.EntN_Attr.keys() and self.EntN_Attr[self.Name_id[equal_ent]] != []:\n equal_set.add(equal_ent)\n might_attr_list = self.EntN_Attr[self.Name_id[equal_ent]]\n might_attr = ';'.join(might_attr_list)\n source_info += f\"It MIGHT be equal type with {self.Id_Name[equal_ent]} (ID:{equal_ent}), so it MIGHT have the Attributes like: {might_attr};\\n\"\n\n\n align_set = set()\n if align_rule and ent1 in self.align_ent_dict.keys():\n candi_tuple_list = self.align_ent_dict[ent1]\n for rel_name, nei, aligned_nei, rel_propm_list in candi_tuple_list:\n align_set.add(nei)\n source_info += f\"It has a {rel_name} relation with entity {self.Id_Name[nei]} (ID: {nei}),\\n \\\n which is aligned with entity {self.Id_Name[aligned_nei]} (ID: {aligned_nei}),\\n \\\n SO, the aligned entity of source entity MIGHT have the relation {rel_name} with ID: {aligned_nei};\\n\"\n\n intersection = equal_set & align_set\n if intersection and self.rule3 is not None:\n for equal_ent in intersection:\n aligned_ = self.seed_dict[equal_ent]\n if aligned_ in self.Id_Name.keys() and self.Id_Name[aligned_] in self.EntN_Attr.keys() and self.EntN_Attr[self.Id_Name[aligned_]] != []:\n might_attr_list = self.EntN_Attr[self.Id_Name[aligned_]]\n might_attr = ';'.join(might_attr_list)\n source_info += f\"It MIGHT be equal type with {self.Id_Name[aligned_]}(ID:{aligned_}), so it MIGHT have the Attributes like: {might_attr};\\n\"\n \n length1 = len(source_info)\n if ent1 in self.Id_rel_dic:\n source_info += \"It has neighbors with the following relations:\\n\"\n rel_dic = self.Id_rel_dic[ent1]\n for ent_name in rel_dic.keys():\n rel_name = rel_dic[ent_name]\n length3 = len(source_info)\n if length3 - length1 > 500: # 控制relation的长度\n break\n if ent_name not in ('Nul', 'Null'):\n entity_id = self.Name_id.get(ent_name, 'unknown')\n source_info += f\"- It has a {rel_name} relation with entity {ent_name} (ID: {entity_id});\\n\"\n\n candidate_info = \"Candidate entity information:\\n\"\n for i, ent2 in enumerate(top5_list):\n candidate_info += f\"Candidate entity {i+1} (ID: {ent2}): \"\n length1 = len(candidate_info)\n if ent2 in self.Id_Name and self.Id_Name[ent2] != 'Null' and self.Id_Name[ent2] != 'Nul':\n candidate_info += f\"Name: {self.Id_Name[ent2]}, \"\n\n if ent2 in self.Id_Name.keys() and self.Id_Name[ent2] in self.EntN_Attr.keys() and self.EntN_Attr[self.Id_Name[ent2]] != []:\n candidate_info += f\"Attributes: {' ; '.join(self.EntN_Attr[self.Id_Name[ent2]])}, \\n\"\n\n length2 = len(candidate_info)\n\n if ent2 in self.Id_rel_dic:\n candidate_info += \"It has neighbors with the following relations:\\n\"\n rel_dic = self.Id_rel_dic[ent2]\n for ent_name in rel_dic.keys():\n rel_name = rel_dic[ent_name]\n length3 = len(candidate_info)\n if length3 - length2>500:\n break\n if ent_name not in ('Nul', 'Null'):\n entity_id = self.Name_id.get(ent_name, 'unknown')\n candidate_info += f\"- It has a {rel_name} relation with entity {ent_name} (ID: {entity_id});\\n\"\n\n full_prompt = task_desp + \"\\n\" + source_info + \"\\n\\n\" + candidate_info + \"\\nPlease output the result:\"\n \n return full_prompt\n\n def test_code(self):\n hit1 = 0\n hit10 = 0\n for epoch in range(0,1):\n if not len(self.en1_list):\n break\n else:\n for i, ent1 in enumerate(tqdm(self.en1_list)):\n topk_indices = self.topk_ST_e1(10)\n top_ent2_list = [self.en2_list[int(it)] for it in topk_indices[i]]\n\n if self.seed_dict[ent1] in top_ent2_list:\n hit10 += 1\n if self.seed_dict[ent1] == top_ent2_list[0]:\n hit1 += 1\n\n print(len(self.en1_list))\n print(hit1)\n print(hit10)\n\n# ts4 = DBP_Part_3('zh_en', None, equal_rule= None, align_rule=None, rule3 = None)\n# ts4.inference()\n\n# ts1 = DBP_Part_3('zh_en', None, equal_rule= True, align_rule=None, rule3 = None)\n# ts1.inference()\n\n# ts2 = DBP_Part_3('zh_en', None, equal_rule= None, align_rule=True, rule3 = None)\n# ts2.inference()\n\n# ts3 = DBP_Part_3('zh_en', None, equal_rule= True, align_rule=True, rule3 = None)\n# ts3.inference()\n\n# ts5 = DBP_Part_3('zh_en', None, equal_rule= True, align_rule=True, rule3 = True)\n# ts5.inference()\n\n\n# # ja-en\n# ts4 = DBP_Part_3('ja_en', None, equal_rule= None, align_rule=None, rule3 = None)\n# ts4.inference()\n\n# ts1 = DBP_Part_3('ja_en', None, equal_rule= True, align_rule=None, rule3 = None)\n# ts1.inference()\n\n# ts2 = DBP_Part_3('ja_en', None, equal_rule= None, align_rule=True, rule3 = None)\n# ts2.inference()\n\n# ts3 = DBP_Part_3('ja_en', None, equal_rule= True, align_rule=True, rule3 = None)\n# ts3.inference()\n\n# ts5 = DBP_Part_3('ja_en', None, equal_rule= True, align_rule=True, rule3 = True)\n# ts5.inference()\n\n\n# fr-en\n# ts4 = DBP_Part_3('fr_en', None, equal_rule= None, align_rule=None, rule3 = None)\n# ts4.inference()\n\n# ts1 = DBP_Part_3('fr_en', None, equal_rule= True, align_rule=None, rule3 = None)\n# ts1.inference()\n\n# ts2 = DBP_Part_3('fr_en', None, equal_rule= None, align_rule=True, rule3 = None)\n# ts2.inference()\n\n# ts3 = DBP_Part_3('fr_en', None, equal_rule= True, align_rule=True, rule3 = None)\n# ts3.inference()\n\n# ts5 = DBP_Part_3('fr_en', None, equal_rule= True, align_rule=True, rule3 = True)\n# ts5.inference()"}}}} |