def test():
wdir = os.path.join(ROOT_DIR, 'weights')
best_weight = os.path.join(wdir, 'best.pt')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
model = myModel()
model.double()
model.to(device)
assert os.path.isdir(wdir), 'weight folder does not exist!'
assert os.path.isfile(best_weight), 'the best weight file does not exist!'
try:
chkpt = torch.load(best_weight, map_location=device)
print('best_fitness', chkpt['epoch'])
model.load_state_dict(chkpt['model'])
if chkpt.get('training_results') is not None:
print(chkpt['training_results'])
del chkpt
except Exception:
SystemExit('Error: Can not load pretrained model!')
test_dataset = CommentDataset(TEST_DATASET, WORD2VEC_MODEL_FILE,
MAX_LEN_FILE)
test_loader = DataLoader(test_dataset,
batch_size=64,
num_workers=min(os.cpu_count(), 64),
shuffle=False)
print('\n\nTesting...')
model.eval()
test_acc = 0
y_true = []
y_pred = []
pbar = tqdm(enumerate(test_loader), total=len(test_loader))
for i, (sentences, labels) in pbar:
sentences = sentences.to(device)
labels = labels.to(device)
with torch.no_grad():
preds = model(sentences)
test_acc += (preds.argmax(1) == labels).sum().item()
y_pred += preds.argmax(1).to('cpu').data.numpy().tolist()
y_true += labels.to('cpu').data.numpy().tolist()
y_pred = np.asarray(y_pred)
y_true = np.asarray(y_true)
s_test = "Acc: {:.3f}% || Precision: {:.3f} || Recall: {:.3f} || F1-score: {:.3f}".format(
test_acc / len(test_dataset) * 100, precision_score(y_true, y_pred),
recall_score(y_true, y_pred), f1_score(y_true, y_pred))
print(s_test)
with open('test_results.txt', 'w+') as f:
f.write(s_test)
torch.cuda.empty_cache()
def main():
with torch.no_grad():
model = myModel(pre_train_dir=args.pretrained_model_path,
dropout_rate=0.5).to(device)
model.load_state_dict(torch.load(args.checkpoints))
sentences, _ = load_data(args.test_path)
with open('./output/result.json', 'w', encoding='utf-8') as f:
for sent in tqdm(sentences):
encode_dict = tokenizer.encode_plus(sent,
max_length=args.max_length,
pad_to_max_length=True)
input_ids = encode_dict['input_ids']
input_seg = encode_dict['token_type_ids']
input_mask = encode_dict['attention_mask']
input_ids = torch.Tensor([input_ids]).long()
input_seg = torch.Tensor([input_seg]).long()
input_mask = torch.Tensor([input_mask]).float()
span_logits = model(input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
is_training=False)
span_logits = torch.argmax(span_logits, dim=-1)[0].to(
torch.device('cpu')).numpy().tolist()
args_index = sapn_decode(span_logits)
text_mapping = get_mapping(sent)
entity_list = []
for k in args_index:
try:
dv = 0
while text_mapping[k[0] - 1 + dv] == []:
dv += 1
start_split = text_mapping[k[0] - 1 + dv]
while text_mapping[k[1] - 1 + dv] == []:
dv += 1
end_split = text_mapping[k[1] - 1 + dv]
argument = sent[start_split[0]:end_split[-1] + 1]
entity_type = k[2]
entity_list.append({
'type': entity_type,
'argument': argument
})
except:
pass
result = {'text': sent, 'entity_list': entity_list}
json_data = json.dumps(result, ensure_ascii=False)
f.write(json_data + '\n')
def main():
with torch.no_grad():
model = myModel(pre_train_dir=args.pretrained_model_path,
dropout_rate=0.5).to(device)
model.load_state_dict(torch.load(args.checkpoints))
sentences = load_data(args.test_path)
with open('./output/result.json', 'w', encoding='utf-8') as f:
for sentence in tqdm(sentences):
type_sent_list = [
entity_type + '[SEP]' + sentence
for entity_type in label2id.keys()
]
input_ids = []
input_seg = []
input_mask = []
for type_sent in type_sent_list:
encode_dict = tokenizer.encode_plus(
type_sent,
max_length=args.max_length,
pad_to_max_length=True)
input_ids.append(encode_dict['input_ids'])
input_seg.append(encode_dict['token_type_ids'])
input_mask.append(encode_dict['attention_mask'])
input_ids = torch.Tensor(input_ids).long()
input_seg = torch.Tensor(input_seg).long()
input_mask = torch.Tensor(input_mask).float()
start_logit, end_logit = model(
input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
is_training=False)
start_logit = start_logit.to(
torch.device('cpu')).numpy().tolist()
end_logit = end_logit.to(torch.device('cpu')).numpy().tolist()
entity_list = []
for i, type_sent in enumerate(type_sent_list):
text_start_id, text_end_id, text_mapping = get_actual_id(
sentence, type_sent)
args_dict = dict()
entity_type = type_sent.split('[SEP]')[0]
# args_index = get_start_end_i(start_logits[i],end_logits[i],span_logits[i],text_start_id,text_end_id)
args_index = extract_entity_from_start_end_ids(
start_logit[i], end_logit[i], text_start_id,
text_end_id)
# args_index=sapn_decode(span_logits[i])
for k in args_index:
tmp = {}
dv = 0
while text_mapping[k[0] - text_start_id + dv] == []:
dv += 1
start_split = text_mapping[k[0] - text_start_id + dv]
dv = 0
while text_mapping[k[1] - text_start_id + dv] == []:
dv -= 1
end_split = text_mapping[k[1] - text_start_id + dv]
tmp['type'] = entity_type
tmp['argument'] = sentence[start_split[0]:
end_split[-1] + 1] #抽取实体
entity_list.append(tmp)
result = {}
result['text'] = sentence
result['entity_list'] = entity_list
json_data = json.dumps(result, ensure_ascii=False)
f.write(json_data + '\n')
def train(opts):
epochs = opts.epochs
batch_size = opts.batch_size
init_seeds(42)
wdir = os.path.join(ROOT_DIR, 'weights')
if not opts.resume and not os.path.isdir(wdir):
os.mkdir(wdir)
best_weight = os.path.join(wdir, 'best.pt')
last_weight = os.path.join(wdir, 'last.pt')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = myModel()
model.double()
model.to(device)
optimizer = optim.SGD(model.parameters(), lr=opts.lr, momentum=opts.momentum, weight_decay=opts.weight_decay)
criterion = nn.CrossEntropyLoss()
start_epoch = 0
best_fitness = 0
if opts.resume:
assert os.path.isdir(wdir), 'weight folder does not exist!'
assert os.path.isfile(last_weight), 'the laster weight file does not exist!'
try:
chkpt = torch.load(last_weight, map_location=device)
model.load_state_dict(chkpt['model'])
if chkpt['optimizer'] is not None:
optimizer.load_state_dict(chkpt['optimizer'])
if chkpt.get('training_results') is not None:
print(chkpt['training_results'])
start_epoch = chkpt['epoch'] + 1
del chkpt
except Exception:
SystemExit('Error: Can not load pretrained model!')
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[round(opts.epochs*x) for x in [0.8,0.9]], gamma=0.2)
scheduler.last_epoch = start_epoch-1
train_dataset = CommentDataset(TRAIN_DATASET,WORD2VEC_MODEL_FILE, MAX_LEN_FILE)
val_dataset = CommentDataset(VAL_DATASET, WORD2VEC_MODEL_FILE, MAX_LEN_FILE)
train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=min(os.cpu_count(), batch_size), shuffle=False, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=64, num_workers=min(os.cpu_count(), batch_size), shuffle=False, pin_memory=True)
for epoch in range(start_epoch, epochs):
save_best = False
model.train()
train_loss = 0
train_acc = 0
print('\n\nTraining...')
pbar = tqdm(enumerate(train_loader), total=len(train_loader))
for i, (sentences, labels) in pbar:
sentences = sentences.to(device)
labels = labels.to(device)
optimizer.zero_grad()
preds = model(sentences)
loss = criterion(preds, labels)
loss.backward()
optimizer.step()
train_loss += loss.item()
train_acc += (preds.argmax(1) == labels).sum().item()
s_train = "Epoch: {} || Train Loss: {:.6f} ".format(
epoch+1, train_loss/len(train_dataset))
pbar.set_description(s_train)
print(s_train)
print('\n\nValidating...')
model.eval()
val_loss = 0
val_acc = 0
for (sentences,labels) in val_loader:
sentences = sentences.to(device)
labels = labels.to(device)
with torch.no_grad():
preds = model(sentences)
loss = criterion(preds, labels)
val_loss += loss.item()
val_acc += (preds.argmax(1) == labels).sum().item()
s_val = "Epoch: {} || Val Loss: {:.6f} || Val Acc: {:.3f} ".format(
epoch+1, val_loss/len(val_dataset), val_acc/len(val_dataset))
print(s_val)
if best_fitness < val_acc:
save_best = True
print('Saving...')
best_fitness = val_acc
chkpt = {'epoch':epoch,
'best_fitness':best_fitness,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(chkpt,best_weight)
del chkpt
if epoch > 0:
scheduler.step()
if opts.save_result:
with open('train_result.txt', 'a+') as f:
f.write(s_train + '\n')
if save_best:
f.write('Saved at:\n')
f.write(s_val + '\n\n')
torch.cuda.empty_cache()
def train():
train_data = data_prepro.yield_data(args.train_path)
test_data = data_prepro.yield_data(args.test_path)
model = myModel(pre_train_dir=args.pretrained_model_path,
dropout_rate=0.5).to(device)
# model.load_state_dict(torch.load(args.checkpoints))
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'gamma', 'beta']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay_rate':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay_rate':
0.0
}]
optimizer = optim.AdamW(params=optimizer_grouped_parameters,
lr=args.learning_rate)
schedule = WarmUp_LinearDecay(optimizer=optimizer,
init_rate=args.learning_rate,
warm_up_epoch=args.warm_up_epoch,
decay_epoch=args.decay_epoch)
span_loss_func = span_loss.Span_loss().to(device)
span_acc = metrics.metrics_span().to(device)
step = 0
best = 0
for epoch in range(args.epoch):
for item in train_data:
step += 1
input_ids, input_mask, input_seg = item["input_ids"], item[
"input_mask"], item["input_seg"]
span_label, span_mask = item['span_label'], item["span_mask"]
optimizer.zero_grad()
span_logits = model(input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
is_training=True)
span_loss_v = span_loss_func(span_logits, span_label.to(device),
span_mask.to(device))
loss = span_loss_v
loss = loss.float().mean().type_as(loss)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
max_norm=args.clip_norm)
schedule.step()
# optimizer.step()
if step % 100 == 0:
span_logits = torch.nn.functional.softmax(span_logits, dim=-1)
recall, precise, span_f1 = span_acc(span_logits,
span_label.to(device))
logger.info(
'epoch %d, step %d, loss %.4f, recall %.4f, precise %.4f, span_f1 %.4f'
% (epoch, step, loss, recall, precise, span_f1))
with torch.no_grad():
count = 0
span_f1 = 0
recall = 0
precise = 0
for item in test_data:
count += 1
input_ids, input_mask, input_seg = item["input_ids"], item[
"input_mask"], item["input_seg"]
span_label, span_mask = item['span_label'], item["span_mask"]
optimizer.zero_grad()
span_logits = model(input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
is_training=False)
tmp_recall, tmp_precise, tmp_span_f1 = span_acc(
span_logits, span_label.to(device))
span_f1 += tmp_span_f1
recall += tmp_recall
precise += tmp_precise
span_f1 = span_f1 / count
recall = recall / count
precise = precise / count
logger.info('-----eval----')
logger.info(
'epoch %d, step %d, loss %.4f, recall %.4f, precise %.4f, span_f1 %.4f'
% (epoch, step, loss, recall, precise, span_f1))
logger.info('-----eval----')
if best < span_f1:
best = span_f1
torch.save(model.state_dict(), f=args.checkpoints)
logger.info('-----save the best model----')
def train():
train_data = data_prepro.yield_data(args.train_path)
test_data = data_prepro.yield_data(args.test_path)
model = myModel(pre_train_dir=args.pretrained_model_path, dropout_rate=0.5).to(device)
# model.load_state_dict(torch.load(args.checkpoints))
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'gamma', 'beta']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay_rate': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay_rate': 0.0}
]
optimizer = optim.AdamW(params=optimizer_grouped_parameters, lr=args.learning_rate)
schedule = WarmUp_LinearDecay(
optimizer = optimizer,
init_rate = args.learning_rate,
warm_up_epoch = args.warm_up_epoch,
decay_epoch = args.decay_epoch
)
loss_func = cross_entropy_loss.cross_entropy().to(device)
acc_func = metrics.metrics_func().to(device)
# start_acc = metrics.metrics_start().to(device)
# end_acc = metrics.metrics_end().to(device)
step=0
best=0
for epoch in range(args.epoch):
for item in train_data:
step+=1
input_ids, input_mask, input_seg = item["input_ids"], item["input_mask"], item["input_seg"]
start_label ,end_label, span_label, seq_mask = item["start_label"],item["end_label"],item['span_label'],item["seq_mask"]
seq_id = item["seq_id"]
optimizer.zero_grad()
start_logits,end_logits = model(
input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
is_training=True
)
start_end_loss = loss_func(start_logits,end_logits,start_label.to(device),end_label.to(device),seq_mask.to(device))
loss = start_end_loss
loss = loss.float().mean().type_as(loss)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=args.clip_norm)
schedule.step()
# optimizer.step()
if step%50 == 0:
start_logits = torch.nn.functional.softmax(start_logits, dim=-1)
end_logits = torch.nn.functional.softmax(end_logits, dim=-1)
_,_,start_f1=acc_func(start_logits,start_label.to(device),seq_mask.to(device))
_,_,end_f1=acc_func(end_logits,end_label.to(device),seq_mask.to(device))
logger.info('epoch %d, step %d, loss %.4f, start_f1 %.4f, end_f1 %.4f'% (
epoch,step,loss,start_f1,end_f1))
with torch.no_grad():
start_f1=0
end_f1=0
count=0
flag_f1=0
for item in test_data:
input_ids, input_mask, input_seg = item["input_ids"], item["input_mask"], item["input_seg"]
start_label,end_label,span_label,seq_mask = item["start_label"],item["end_label"],item['span_label'],item["seq_mask"]
seq_id = item["seq_id"]
optimizer.zero_grad()
start_logits,end_logits = model(
input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
is_training=False
)
_,_,tmp_f1_start=acc_func(start_logits,start_label.to(device),seq_mask.to(device))
start_f1+=tmp_f1_start
_,_,tmp_f1_end=acc_func(end_logits,end_label.to(device),seq_mask.to(device))
end_f1+=tmp_f1_end
count+=1
start_f1=start_f1/count
end_f1=end_f1/count
logger.info('-----eval----')
logger.info('epoch %d, step %d, loss %.4f, start_f1 %.4f, end_f1 %.4f'% (
epoch,step,loss,start_f1,end_f1))
logger.info('-----eval----')
if best < start_f1+end_f1:
best=start_f1+end_f1
torch.save(model.state_dict(), f=args.checkpoints)
logger.info('-----save the best model----')
def train():
train_data = data_prepro.yield_data(args.train_path)
test_data = data_prepro.yield_data(args.test_path)
model = myModel(pre_train_dir=args.pretrained_model_path,
dropout_rate=0.5).to(device)
# model.load_state_dict(torch.load(args.checkpoints),False)
no_decay = ["bias", "LayerNorm.weight"]
param_optimizer = list(model.named_parameters())
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay_rate':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay_rate':
0.0
}]
optimizer = optim.AdamW(params=optimizer_grouped_parameters,
lr=args.learning_rate)
schedule = WarmUp_LinearDecay(optimizer=optimizer,
init_rate=args.learning_rate,
warm_up_epoch=args.warm_up_epoch,
decay_epoch=args.decay_epoch)
acc_func = metrics.metrics_func().to(device)
step = 0
best_f1 = 0
for epoch in range(args.epoch):
for item in train_data:
step += 1
input_ids, input_mask, input_seg = item["input_ids"], item[
"input_mask"], item["input_seg"]
labels = item["labels"]
token_word = item["token_word"]
optimizer.zero_grad()
loss, logits = model(input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
token_word=token_word.to(device),
labels=labels)
loss = loss.float().mean().type_as(loss)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
max_norm=args.clip_norm)
schedule.step()
if step % 100 == 0:
recall, precise, f1 = acc_func(logits, labels.to(device))
logger.info(
'epoch %d, step %d, loss %.4f, recall %.4f, precise %.4f, f1 %.4f'
% (epoch, step, loss, recall, precise, f1))
with torch.no_grad():
recall = 0
precise = 0
f1 = 0
count = 0
for item in test_data:
count += 1
input_ids, input_mask, input_seg = item["input_ids"], item[
"input_mask"], item["input_seg"]
labels = item["labels"]
token_word = item["token_word"]
loss, logits = model(input_ids=input_ids.to(device),
input_mask=input_mask.to(device),
input_seg=input_seg.to(device),
token_word=token_word.to(device),
labels=labels)
tmp_recall, tmp_precise, tmp_f1 = acc_func(
logits, labels.to(device))
recall += tmp_recall
precise += tmp_precise
f1 += tmp_f1
recall = recall / count
precise = precise / count
f1 = f1 / count
logger.info('-----eval----')
logger.info(
'epoch %d, step %d, loss %.4f, recall %.4f, precise %.4f, f1 %.4f'
% (epoch, step, loss, recall, precise, f1))
logger.info('-----eval----')
if best_f1 < f1:
best_f1 = f1
torch.save(model.state_dict(), f=args.checkpoints)
logger.info('-----save the best model----')