283 lines
11 KiB
Python
283 lines
11 KiB
Python
|
# _*_ coding:utf-8 _*_
|
|||
|
|
import logging
|
|||
|
|
import random
|
|||
|
|
import sys
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from torch import nn
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
import models.spo_net.etl_stl as etl
|
|||
|
|
from utils.data_util import Tokenizer
|
|||
|
|
from utils.optimizer_util import set_optimizer
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
tokenizer = Tokenizer('cpt/bert-base-chinese/vocab.txt', do_lower_case=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Trainer(object):
|
|||
|
|
|
|||
|
|
def __init__(self, args, data_loaders, examples, word_emb, ent_conf, spo_conf):
|
|||
|
|
|
|||
|
|
self.args = args
|
|||
|
|
self.device = torch.device("cuda:{}".format(args.device_id) if torch.cuda.is_available() else "cpu")
|
|||
|
|
self.n_gpu = torch.cuda.device_count()
|
|||
|
|
|
|||
|
|
self.id2rel = {item: key for key, item in spo_conf.items()}
|
|||
|
|
self.rel2id = spo_conf
|
|||
|
|
|
|||
|
|
random.seed(args.seed)
|
|||
|
|
np.random.seed(args.seed)
|
|||
|
|
torch.manual_seed(args.seed)
|
|||
|
|
if self.n_gpu > 0:
|
|||
|
|
torch.cuda.manual_seed_all(args.seed)
|
|||
|
|
|
|||
|
|
self.model = etl.ERENet(args, word_emb, ent_conf, spo_conf)
|
|||
|
|
self.model.to(self.device)
|
|||
|
|
# self.resume(args)
|
|||
|
|
logging.info('total gpu num is {}'.format(self.n_gpu))
|
|||
|
|
|
|||
|
|
if self.n_gpu > 1:
|
|||
|
|
self.model = nn.DataParallel(self.model.cuda(), device_ids=[0, 1])
|
|||
|
|
|
|||
|
|
train_dataloader, dev_dataloader = data_loaders
|
|||
|
|
train_eval, dev_eval = examples
|
|||
|
|
self.eval_file_choice = {
|
|||
|
|
"train": train_eval,
|
|||
|
|
"dev": dev_eval,
|
|||
|
|
}
|
|||
|
|
self.data_loader_choice = {
|
|||
|
|
"train": train_dataloader,
|
|||
|
|
"dev": dev_dataloader,
|
|||
|
|
}
|
|||
|
|
# args.use_bert = False
|
|||
|
|
self.optimizer = set_optimizer(args, self.model,
|
|||
|
|
train_steps=(int(len(train_eval) / args.train_batch_size) + 1) * args.epoch_num)
|
|||
|
|
|
|||
|
|
def train(self, args):
|
|||
|
|
|
|||
|
|
best_f1 = 0.0
|
|||
|
|
patience_stop = 0
|
|||
|
|
self.model.train()
|
|||
|
|
step_gap = 20
|
|||
|
|
for epoch in range(int(args.epoch_num)):
|
|||
|
|
|
|||
|
|
global_loss = 0.0
|
|||
|
|
|
|||
|
|
for step, batch in tqdm(enumerate(self.data_loader_choice[u"train"]), mininterval=5,
|
|||
|
|
desc=u'training at epoch : %d ' % epoch, leave=False, file=sys.stdout):
|
|||
|
|
|
|||
|
|
loss = self.forward(batch)
|
|||
|
|
|
|||
|
|
if step % step_gap == 0:
|
|||
|
|
global_loss += loss
|
|||
|
|
current_loss = global_loss / step_gap
|
|||
|
|
print(
|
|||
|
|
u"step {} / {} of epoch {}, train/loss: {}".format(step, len(self.data_loader_choice["train"]),
|
|||
|
|
epoch, current_loss))
|
|||
|
|
global_loss = 0.0
|
|||
|
|
|
|||
|
|
res_dev = self.eval_data_set("dev")
|
|||
|
|
if res_dev['f1'] >= best_f1:
|
|||
|
|
best_f1 = res_dev['f1']
|
|||
|
|
logging.info("** ** * Saving fine-tuned model ** ** * ")
|
|||
|
|
model_to_save = self.model.module if hasattr(self.model,
|
|||
|
|
'module') else self.model # Only save the model it-self
|
|||
|
|
output_model_file = args.output + "/pytorch_model.bin"
|
|||
|
|
torch.save(model_to_save.state_dict(), str(output_model_file))
|
|||
|
|
patience_stop = 0
|
|||
|
|
else:
|
|||
|
|
patience_stop += 1
|
|||
|
|
if patience_stop >= args.patience_stop:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
def resume(self, args):
|
|||
|
|
resume_model_file = args.output + "/pytorch_model.bin"
|
|||
|
|
logging.info("=> loading checkpoint '{}'".format(resume_model_file))
|
|||
|
|
checkpoint = torch.load(resume_model_file, map_location='cpu')
|
|||
|
|
self.model.load_state_dict(checkpoint)
|
|||
|
|
|
|||
|
|
def forward(self, batch, chosen=u'train', eval=False, answer_dict=None):
|
|||
|
|
|
|||
|
|
batch = tuple(t.to(self.device) for t in batch)
|
|||
|
|
if not eval:
|
|||
|
|
char_ids, word_ids, token_type_ids, subject_ids, subject_labels, object_labels = batch
|
|||
|
|
loss = self.model(char_ids=char_ids, word_ids=word_ids, token_type_ids=token_type_ids,
|
|||
|
|
subject_ids=subject_ids, subject_labels=subject_labels, object_labels=object_labels)
|
|||
|
|
if self.n_gpu > 1:
|
|||
|
|
loss = loss.mean() # mean() to average on multi-gpu.
|
|||
|
|
|
|||
|
|
loss.backward()
|
|||
|
|
loss = loss.item()
|
|||
|
|
self.optimizer.step()
|
|||
|
|
self.optimizer.zero_grad()
|
|||
|
|
return loss
|
|||
|
|
else:
|
|||
|
|
p_ids, char_ids, word_ids = batch
|
|||
|
|
eval_file = self.eval_file_choice[chosen]
|
|||
|
|
qids, subject_pred, po_pred = self.model(q_ids=p_ids,
|
|||
|
|
char_ids=char_ids,
|
|||
|
|
word_ids=word_ids,
|
|||
|
|
eval_file=eval_file, is_eval=eval)
|
|||
|
|
ans_dict = self.convert_spo_contour(qids, subject_pred, po_pred, eval_file,
|
|||
|
|
answer_dict, use_bert=self.args.use_bert)
|
|||
|
|
return ans_dict
|
|||
|
|
|
|||
|
|
def eval_data_set(self, chosen="dev"):
|
|||
|
|
|
|||
|
|
self.model.eval()
|
|||
|
|
|
|||
|
|
data_loader = self.data_loader_choice[chosen]
|
|||
|
|
eval_file = self.eval_file_choice[chosen]
|
|||
|
|
answer_dict = {i: [[], []] for i in range(len(eval_file))}
|
|||
|
|
|
|||
|
|
last_time = time.time()
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for _, batch in tqdm(enumerate(data_loader), mininterval=5, leave=False, file=sys.stdout):
|
|||
|
|
self.forward(batch, chosen, eval=True, answer_dict=answer_dict)
|
|||
|
|
used_time = time.time() - last_time
|
|||
|
|
logging.info('chosen {} took : {} sec'.format(chosen, used_time))
|
|||
|
|
res = self.evaluate(eval_file, answer_dict, chosen)
|
|||
|
|
self.model.train()
|
|||
|
|
return res
|
|||
|
|
|
|||
|
|
def show(self, chosen="dev"):
|
|||
|
|
|
|||
|
|
self.model.eval()
|
|||
|
|
answer_dict = {}
|
|||
|
|
|
|||
|
|
data_loader = self.data_loader_choice[chosen]
|
|||
|
|
eval_file = self.eval_file_choice[chosen]
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for _, batch in tqdm(enumerate(data_loader), mininterval=5, leave=False, file=sys.stdout):
|
|||
|
|
loss, answer_dict_ = self.forward(batch, chosen, eval=True)
|
|||
|
|
answer_dict.update(answer_dict_)
|
|||
|
|
# self.badcase_analysis(eval_file, answer_dict, chosen)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def evaluate(eval_file, answer_dict, chosen):
|
|||
|
|
|
|||
|
|
entity_em = 0
|
|||
|
|
entity_pred_num = 0
|
|||
|
|
entity_gold_num = 0
|
|||
|
|
|
|||
|
|
triple_em = 0
|
|||
|
|
triple_pred_num = 0
|
|||
|
|
triple_gold_num = 0
|
|||
|
|
X, Y, Z = 1e-10, 1e-10, 1e-10
|
|||
|
|
for key, value in answer_dict.items():
|
|||
|
|
triple_gold = eval_file[key].gold_answer
|
|||
|
|
entity_gold = eval_file[key].sub_entity_list
|
|||
|
|
|
|||
|
|
entity_pred, triple_pred = value
|
|||
|
|
|
|||
|
|
entity_em += len(set(entity_pred) & set(entity_gold))
|
|||
|
|
entity_pred_num += len(set(entity_pred))
|
|||
|
|
entity_gold_num += len(set(entity_gold))
|
|||
|
|
|
|||
|
|
# triple_em += len(set(triple_pred) & set(triple_gold))
|
|||
|
|
# triple_pred_num += len(set(triple_pred))
|
|||
|
|
# triple_gold_num += len(set(triple_gold))
|
|||
|
|
|
|||
|
|
R = set([SPO(spo) for spo in triple_pred])
|
|||
|
|
T = set([SPO(spo) for spo in triple_gold])
|
|||
|
|
|
|||
|
|
# R = set([spo for spo in triple_pred])
|
|||
|
|
# T = set([spo for spo in triple_gold])
|
|||
|
|
# if R != T:
|
|||
|
|
# print(R)
|
|||
|
|
# print(T)
|
|||
|
|
X += len(R & T)
|
|||
|
|
Y += len(R)
|
|||
|
|
Z += len(T)
|
|||
|
|
|
|||
|
|
f1, precision, recall = 2 * X / (Y + Z), X / Y, X / Z
|
|||
|
|
|
|||
|
|
entity_precision = 100.0 * entity_em / entity_pred_num if entity_pred_num > 0 else 0.
|
|||
|
|
entity_recall = 100.0 * entity_em / entity_gold_num if entity_gold_num > 0 else 0.
|
|||
|
|
entity_f1 = 2 * entity_recall * entity_precision / (entity_recall + entity_precision) if (
|
|||
|
|
entity_recall + entity_precision) != 0 else 0.0
|
|||
|
|
|
|||
|
|
print('============================================')
|
|||
|
|
print("{}/entity_em: {},\tentity_pred_num&entity_gold_num: {}\t{} ".format(chosen, entity_em, entity_pred_num,
|
|||
|
|
entity_gold_num))
|
|||
|
|
print(
|
|||
|
|
"{}/entity_f1: {}, \tentity_precision: {},\tentity_recall: {} ".format(chosen, entity_f1, entity_precision,
|
|||
|
|
entity_recall))
|
|||
|
|
print('============================================')
|
|||
|
|
print("{}/em: {},\tpre&gold: {}\t{} ".format(chosen, X, Y, Z))
|
|||
|
|
print("{}/f1: {}, \tPrecision: {},\tRecall: {} ".format(chosen, f1 * 100, precision * 100,
|
|||
|
|
recall * 100))
|
|||
|
|
return {'f1': f1, "recall": recall, "precision": precision}
|
|||
|
|
|
|||
|
|
def convert_spo_contour(self, qids, subject_preds, po_preds, eval_file, answer_dict, use_bert=False):
|
|||
|
|
|
|||
|
|
for qid, subject, po_pred in zip(qids.data.cpu().numpy(), subject_preds.data.cpu().numpy(),
|
|||
|
|
po_preds.data.cpu().numpy()):
|
|||
|
|
if qid == -1:
|
|||
|
|
continue
|
|||
|
|
tokens = eval_file[qid.item()].context
|
|||
|
|
|
|||
|
|
context = eval_file[qid.item()].context
|
|||
|
|
seq_len = min(len(context), self.args.max_len)
|
|||
|
|
|
|||
|
|
pred_entities = find_tag_position(po_pred, seq_len, self.id2rel)
|
|||
|
|
po_predict = []
|
|||
|
|
for o_1, o_2, tag in pred_entities:
|
|||
|
|
po_predict.append((tokens[subject[0]:subject[1] + 1],
|
|||
|
|
tag,
|
|||
|
|
tokens[o_1:o_2 + 1])
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if qid not in answer_dict:
|
|||
|
|
print('erro in answer_dict ')
|
|||
|
|
else:
|
|||
|
|
answer_dict[qid][0].append(
|
|||
|
|
tokens[subject[0]:subject[1] + 1])
|
|||
|
|
answer_dict[qid][1].extend(po_predict)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def find_tag_position(find_list, seq_len, id2rel):
|
|||
|
|
tag_list = list()
|
|||
|
|
j = 0
|
|||
|
|
while j < seq_len:
|
|||
|
|
end = j
|
|||
|
|
flag = True
|
|||
|
|
|
|||
|
|
if find_list[j] % 2 == 0 and find_list[j] != 0:
|
|||
|
|
start = j
|
|||
|
|
tag = id2rel[find_list[start]].split('-')[1]
|
|||
|
|
for k in range(start + 1, seq_len):
|
|||
|
|
if find_list[k] != find_list[start] + 1:
|
|||
|
|
end = k - 1
|
|||
|
|
flag = False
|
|||
|
|
break
|
|||
|
|
if flag:
|
|||
|
|
end = seq_len - 1
|
|||
|
|
tag_list.append((start, end, tag))
|
|||
|
|
j = end + 1
|
|||
|
|
return tag_list
|