File size: 5,395 Bytes
3d3bb36 | 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 | from utils import *
from model import *
from torch.utils import data
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
def get_triple_list(sub_head_ids, sub_tail_ids, model, encoded_text, text, mask, offset_mapping):
id2rel, _ = get_rel()
triple_list = []
for sub_head_id in sub_head_ids:
sub_tail_ids = sub_tail_ids[sub_tail_ids >= sub_head_id]
if len(sub_tail_ids) == 0:
continue
sub_tail_id = sub_tail_ids[0]
if mask[sub_head_id] == 0 or mask[sub_tail_id] == 0:
continue
# 根据位置信息反推出 subject 文本内容
sub_head_pos_id = offset_mapping[sub_head_id][0]
sub_tail_pos_id = offset_mapping[sub_tail_id][1]
subject_text = text[sub_head_pos_id:sub_tail_pos_id]
# 根据 subject 计算出对应 object 和 relation
sub_head_seq = torch.tensor(multihot(len(mask), sub_head_id)).to(DEVICE)
sub_tail_seq = torch.tensor(multihot(len(mask), sub_tail_id)).to(DEVICE)
pred_obj_head, pred_obj_tail = model.get_objs_for_specific_sub(\
encoded_text.unsqueeze(0), sub_head_seq.unsqueeze(0), sub_tail_seq.unsqueeze(0))
# 按分类找对应关系
pred_obj_head = pred_obj_head[0].T
pred_obj_tail = pred_obj_tail[0].T
for j in range(len(pred_obj_head)):
obj_head_ids = torch.where(pred_obj_head[j] > OBJ_HEAD_BAR)[0]
obj_tail_ids = torch.where(pred_obj_tail[j] > OBJ_TAIL_BAR)[0]
for obj_head_id in obj_head_ids:
obj_tail_ids = obj_tail_ids[obj_tail_ids >= obj_head_id]
if len(obj_tail_ids) == 0:
continue
obj_tail_id = obj_tail_ids[0]
if mask[obj_head_id] == 0 or mask[obj_tail_id] == 0:
continue
# 根据位置信息反推出 object 文本内容，mapping中已经有移位，不需要再加1
obj_head_pos_id = offset_mapping[obj_head_id][0]
obj_tail_pos_id = offset_mapping[obj_tail_id][1]
object_text = text[obj_head_pos_id:obj_tail_pos_id]
triple_list.append((subject_text, id2rel[j], object_text))
return list(set(triple_list))
def report(model, encoded_text, pred_y, batch_text, batch_mask):
# 计算三元结构，和统计指标
pred_sub_head, pred_sub_tail, _, _ = pred_y
true_triple_list = batch_text['triple_list']
pred_triple_list = []
correct_num, predict_num, gold_num = 0, 0, 0
# 遍历batch
for i in range(len(pred_sub_head)):
text = batch_text['text'][i]
true_triple_item = true_triple_list[i]
mask = batch_mask[i]
offset_mapping = batch_text['offset_mapping'][i]
sub_head_ids = torch.where(pred_sub_head[i] > SUB_HEAD_BAR)[0]
sub_tail_ids = torch.where(pred_sub_tail[i] > SUB_TAIL_BAR)[0]
pred_triple_item = get_triple_list(sub_head_ids, sub_tail_ids, model, \
encoded_text[i], text, mask, offset_mapping)
# 统计个数
correct_num += len(set(true_triple_item) & set(pred_triple_item))
predict_num += len(set(pred_triple_item))
gold_num += len(set(true_triple_item))
pred_triple_list.append(pred_triple_item)
precision = correct_num / (predict_num + EPS)
recall = correct_num / (gold_num + EPS)
f1_score = 2 * precision * recall / (precision + recall + EPS)
print('\tcorrect_num:', correct_num, 'predict_num:', predict_num, 'gold_num:', gold_num)
print('\tprecision:%.3f' % precision, 'recall:%.3f' % recall, 'f1_score:%.3f' % f1_score)
if __name__ == '__main__':
model = CasRel().to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
dataset = Dataset()
for e in range(EPOCH):
loader = data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=dataset.collate_fn)
for b, (batch_mask, batch_x, batch_y) in enumerate(loader):
# print(batch_x)
# exit()
batch_text, batch_sub_rnd = batch_x
batch_sub, batch_obj_rel = batch_y
# 整理input数据并预测
input_mask = torch.tensor(batch_mask).to(DEVICE)
input = (
torch.tensor(batch_text['input_ids']).to(DEVICE),
torch.tensor(batch_sub_rnd['head_seq']).to(DEVICE),
torch.tensor(batch_sub_rnd['tail_seq']).to(DEVICE),
)
encoded_text, pred_y = model(input, input_mask)
# 整理target数据并计算损失
true_y = (
torch.tensor(batch_sub['heads_seq']).to(DEVICE),
torch.tensor(batch_sub['tails_seq']).to(DEVICE),
torch.tensor(batch_obj_rel['heads_mx']).to(DEVICE),
torch.tensor(batch_obj_rel['tails_mx']).to(DEVICE),
)
loss = model.loss_fn(true_y, pred_y, input_mask)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if b % 5 == 0:
print('>> epoch:', e, 'batch:', b, 'loss:', loss.item())
# print('>> epoch:', e, 'batch:', b, 'loss:', loss.item())
if b % 500 == 0:
report(model, encoded_text, pred_y, batch_text, batch_mask)
if e % 3 == 0:
torch.save(model, MODEL_DIR + f'model_{e}.pth') |