#! -*- coding:utf-8 -*- # 三元组抽取任务，tplinker, cat方式实体部分收敛较快，关系部分收敛较慢 # 官方链接：https://github.com/131250208/TPlinker-joint-extraction # 数据集：http://ai.baidu.com/broad/download?dataset=sked import json from bert4torch.tokenizers import Tokenizer from bert4torch.models import build_transformer_model, BaseModel from bert4torch.snippets import sequence_padding, Callback, ListDataset from bert4torch.layers import TplinkerHandshakingKernel from tqdm import tqdm import torch import torch.nn as nn from torch.utils.data import DataLoader import torch.optim as optim maxlen = 50 batch_size = 64 config_path = 'F:/Projects/pretrain_ckpt/robert/[hit_torch_base]--chinese-roberta-wwm-ext-base/config.json' checkpoint_path = 'F:/Projects/pretrain_ckpt/robert/[hit_torch_base]--chinese-roberta-wwm-ext-base/pytorch_model.bin' dict_path = 'F:/Projects/pretrain_ckpt/robert/[hit_torch_base]--chinese-roberta-wwm-ext-base/vocab.txt' device = 'cuda' if torch.cuda.is_available() else 'cpu' # 加载标签字典 predicate2id, id2predicate = {}, {} with open('F:/Projects/data/corpus/relation_extraction/BD_Knowledge_Extraction/all_50_schemas', encoding='utf-8') as f: for l in f: l = json.loads(l) if l['predicate'] not in predicate2id: id2predicate[len(predicate2id)] = l['predicate'] predicate2id[l['predicate']] = len(predicate2id) # 建立分词器 tokenizer = Tokenizer(dict_path, do_lower_case=True) # 加载数据集 class MyDataset(ListDataset): @staticmethod def load_data(filename): """加载数据 单条格式：{'text': text, 'spo_list': [(s, p, o)]} """ D = [] with open(filename, encoding='utf-8') as f: for l in f: l = json.loads(l) D.append({'text': l['text'], 'spo_list': [(spo['subject'], spo['predicate'], spo['object']) for spo in l['spo_list']]}) return D def trans_ij2k(seq_len, i, j): '''把第i行，第j列转化成上三角flat后的序号 ''' if (i > seq_len - 1) or (j > seq_len - 1) or (i > j): return 0 return int(0.5*(2*seq_len-i+1)*i+(j-i)) map_ij2k = {(i, j): trans_ij2k(maxlen, i, j) for i in range(maxlen) for j in range(maxlen) if j >= i} map_k2ij = {v: k for k, v in map_ij2k.items()} def search(pattern, sequence): """从sequence中寻找子串pattern 如果找到，返回第一个下标；否则返回-1。 """ n = len(pattern) for i in range(len(sequence)): if sequence[i:i + n] == pattern: return i return -1 def collate_fn(batch): pair_len = maxlen * (maxlen+1)//2 # batch_entity_labels: [btz, pair_len] # batch_head_labels: [btz, rel_size, pair_len] # batch_tail_labels: [btz, rel_size, pair_len] batch_entity_labels = torch.zeros((len(batch), pair_len), dtype=torch.long, device=device) batch_head_labels = torch.zeros((len(batch), len(predicate2id), pair_len), dtype=torch.long, device=device) batch_tail_labels = torch.zeros((len(batch), len(predicate2id), pair_len), dtype=torch.long, device=device) batch_token_ids = [] for i, d in enumerate(batch): token_ids = tokenizer.encode(d['text'])[0][1:-1][:maxlen] # 这里要限制取前max_len个 batch_token_ids.append(token_ids) # 整理三元组 {s: [(o, p)]} for s, p, o in d['spo_list']: s = tokenizer.encode(s)[0][1:-1] p = predicate2id[p] o = tokenizer.encode(o)[0][1:-1] sh = search(s, token_ids) # 这里超过长度就会找不到 oh = search(o, token_ids) if sh != -1 and oh != -1: st, ot = sh+len(s)-1, oh+len(o)-1 batch_entity_labels[i, map_ij2k[sh, st]] = 1 batch_entity_labels[i, map_ij2k[oh, ot]] = 1 if sh <= oh: batch_head_labels[i, p, map_ij2k[sh, oh]] = 1 else: batch_head_labels[i, p, map_ij2k[oh, sh]] = 2 if st <= ot: batch_tail_labels[i, p, map_ij2k[st, ot]] = 1 else: batch_tail_labels[i, p, map_ij2k[ot, st]] = 2 batch_token_ids = torch.tensor(sequence_padding(batch_token_ids, length=maxlen), dtype=torch.long, device=device) return [batch_token_ids], [batch_entity_labels, batch_head_labels, batch_tail_labels] train_dataloader = DataLoader(MyDataset('F:/Projects/data/corpus/relation_extraction/BD_Knowledge_Extraction/train_data.json'), batch_size=batch_size, shuffle=True, collate_fn=collate_fn) valid_dataset = MyDataset('F:/Projects/data/corpus/relation_extraction/BD_Knowledge_Extraction/dev_data.json') valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, collate_fn=collate_fn) # 定义bert上的模型结构 class Model(BaseModel): def __init__(self): super().__init__() self.bert = build_transformer_model(config_path, checkpoint_path, segment_vocab_size=0) self.combine_fc = nn.Linear(768*2, 768) self.ent_fc = nn.Linear(768, 2) self.head_rel_fc = nn.Linear(768, len(predicate2id)*3) self.tail_rel_fc = nn.Linear(768, len(predicate2id)*3) self.handshaking_kernel = TplinkerHandshakingKernel(768, shaking_type='cat') def forward(self, inputs): last_hidden_state = self.bert(inputs) # [btz, seq_len, hdsz] shaking_hiddens = self.handshaking_kernel(last_hidden_state) # [btz, pair_len, hdsz] ent_shaking_outputs = self.ent_fc(shaking_hiddens) # [btz, pair_len, 2] btz, pair_len = shaking_hiddens.shape[:2] head_rel_shaking_outputs = self.head_rel_fc(shaking_hiddens).reshape(btz, -1, pair_len, 3) #[btz, predicate_num, pair_len, 3] tail_rel_shaking_outputs = self.tail_rel_fc(shaking_hiddens).reshape(btz, -1, pair_len, 3) return ent_shaking_outputs, head_rel_shaking_outputs, tail_rel_shaking_outputs model = Model().to(device) class MyLoss(nn.CrossEntropyLoss): def __init__(self, **kwargs): super().__init__(**kwargs) def forward(self, y_preds, y_trues): loss_list = [] for y_pred, y_true in zip(y_preds, y_trues): loss = super().forward(y_pred.view(-1, y_pred.size()[-1]), y_true.view(-1)) loss_list.append(loss) z = (2 * len(predicate2id) + 1) total_steps = 6000 # 前期实体识别的权重高一些，建议也可以设置为model.total_steps w_ent = max(1 / z + 1 - model.global_step / total_steps, 1 / z) w_rel = min((len(predicate2id) / z) * model.global_step / total_steps, (len(predicate2id) / z)) loss = w_ent*loss_list[0] + w_rel*loss_list[1] + w_rel*loss_list[2] return {'loss': loss, 'entity_loss': loss_list[0], 'head_loss': loss_list[1], 'tail_loss': loss_list[2]} model.compile(loss=MyLoss(), optimizer=optim.Adam(model.parameters(), 5e-5), metrics=['entity_loss', 'head_loss', 'tail_loss']) def extract_spoes(text): """抽取输入text所包含的三元组 """ def get_spots_fr_shaking_tag(shaking_tag): '''解析关系 ''' spots = [] for shaking_inds in shaking_tag.nonzero(): rel_id = shaking_inds[0].item() tag_id = shaking_tag[rel_id][shaking_inds[1]].item() matrix_inds = map_k2ij[shaking_inds[1].item()] # 保证前面是subject，后面是object if tag_id == 1: spot = (rel_id, matrix_inds[0], matrix_inds[1]) elif tag_id == 2: spot = (rel_id, matrix_inds[1], matrix_inds[0]) spots.append(spot) return spots tokens = tokenizer.tokenize(text)[1:-1] mapping = tokenizer.rematch(text, tokens) token_ids = tokenizer.encode(text)[0][1:-1] token_ids_ts = torch.tensor(sequence_padding([token_ids], length=maxlen), dtype=torch.long, device=device) outputs = model.predict([token_ids_ts]) outputs = [o[0].argmax(dim=-1) for o in outputs] # 抽取entity ent_matrix_spots = set() ent_text = set() for shaking_ind in outputs[0].nonzero(): shaking_ind_ = shaking_ind[0].item() # tag_id = outputs[0][shaking_ind_] matrix_inds = map_k2ij[shaking_ind_] spot = (matrix_inds[0], matrix_inds[1]) if (spot[0] < len(mapping)) and (spot[1] < len(mapping)): # 实体起始在mapping范围内 ent_matrix_spots.add(spot) ent_text.add(text[mapping[spot[0]][0]:mapping[spot[1]][-1] + 1]) # 识别对应的predicate head_rel_matrix_spots = get_spots_fr_shaking_tag(outputs[1]) tail_rel_matrix_spots = get_spots_fr_shaking_tag(outputs[2]) spoes = [] for rel_h, sh, oh in head_rel_matrix_spots: for rel_t, st, ot in tail_rel_matrix_spots: # 如果关系相同，且(sh, st)和(oh, ot)都在entity_maxtrix_spots中 if (rel_h == rel_t) and ((sh, st) in ent_matrix_spots) and ((oh, ot) in ent_matrix_spots): spoes.append((text[mapping[sh][0]:mapping[st][-1] + 1], id2predicate[rel_h], text[mapping[oh][0]:mapping[ot][-1] + 1])) return spoes, token_ids, ent_text class SPO(tuple): """用来存三元组的类 表现跟tuple基本一致，只是重写了 __hash__ 和 __eq__ 方法， 使得在判断两个三元组是否等价时容错性更好。 """ def __init__(self, spo): self.spox = (tuple(tokenizer.tokenize(spo[0])), spo[1], tuple(tokenizer.tokenize(spo[2]))) def __hash__(self): return self.spox.__hash__() def __eq__(self, spo): return self.spox == spo.spox def evaluate(data): """评估函数，计算f1、precision、recall """ X, Y, Z = 0, 1e-10, 1e-10 E1, E2 = 0, 1e-10 f = open('dev_pred.json', 'w', encoding='utf-8') pbar = tqdm() for d in data: spoes, token_ids, ent_text_pred = extract_spoes(d['text']) # spo_list是用来根据maxlen删减的 spo_list = [] for s, p, o in d['spo_list']: s_ = tokenizer.encode(s)[0][1:-1] o_ = tokenizer.encode(o)[0][1:-1] sh = search(s_, token_ids) # 这里超过长度就会找不到 oh = search(o_, token_ids) if sh != -1 and oh != -1: spo_list.append((s, p, o)) # 计算三元组的f1值 R = set([SPO(spo) for spo in spoes]) T = set([SPO(spo) for spo in spo_list]) X += len(R & T) Y += len(R) Z += len(T) f1, precision, recall = 2 * X / (Y + Z), X / Y, X / Z # 计算实体的指标 ent_text_truth = set([spo[0] for spo in spo_list] + [spo[-1] for spo in spo_list]) E1 += len(ent_text_pred & ent_text_truth) E2 += len(ent_text_truth) E_acc = E1 / E2 # 计算entity_matrix, head_matrix，tail_matrix的accuracy pbar.update() pbar.set_description('f1: %.5f, precision: %.5f, recall: %.5f, ent_acc: %.5f' % (f1, precision, recall, E_acc)) s = json.dumps({'text': d['text'], 'spo_list': list(T), 'spo_list_pred': list(R), 'new': list(R - T), 'lack': list(T - R)}, ensure_ascii=False, indent=4) f.write(s + '\n') pbar.close() f.close() return f1, precision, recall class Evaluator(Callback): """评估与保存 """ def __init__(self): self.best_val_f1 = 0. def on_epoch_end(self, steps, epoch, logs=None): f1, precision, recall = evaluate(valid_dataset.data) if f1 >= self.best_val_f1: self.best_val_f1 = f1 # model.save_weights('best_model.pt') print('f1: %.5f, precision: %.5f, recall: %.5f, best f1: %.5f\n' % (f1, precision, recall, self.best_val_f1)) if __name__ == '__main__': evaluator = Evaluator() model.fit(train_dataloader, steps_per_epoch=None, epochs=20, callbacks=[evaluator]) else: model.load_weights('best_model.pt')