-
Notifications
You must be signed in to change notification settings - Fork 1
/
sentence_label_predict_ace.py
101 lines (77 loc) · 2.86 KB
/
sentence_label_predict_ace.py
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
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "7"
import pickle
import torch
from transformers import AdamW
from transformers import get_linear_schedule_with_warmup
import config
from dataset import Dataset
from model import SentenceClassification, get_type_rep
from utils import save_model, load_model
def load_dataset():
filename = 'data/data_ace_all_mask.pk'
data = pickle.load(open(filename, 'rb'))
return data['train'], data['val'], data['test']
def evaluate(model, device, testset):
with torch.no_grad():
golds = []
predicts = []
for batch in testset.get_tqdm(device, False):
data_x, data_x_mask, data_labels, _, _, appendix = batch
logits = model.compute_logits(data_x, data_x_mask)
logits = torch.nn.Sigmoid()(logits)
golds.extend(data_labels.cpu().numpy())
predicts.extend(logits.cpu().numpy())
for i in range(10):
th = i / 10
n_gold, n_predict, n_correct = 0, 0, 0
for g, p in zip(golds, predicts):
n_gold += sum(g)
n_predict += sum(p >= th)
n_correct += sum(g * (p > th))
p = n_correct / n_predict
r = n_correct / n_gold
f1 = 2 * p * r / (p + r)
print(n_gold, n_predict, n_correct, p, r, f1)
print('___')
for i in range(1, len(config.idx2tag)):
th = 0.3
n_gold, n_predict, n_correct = 0, 0, 0
for g, p in zip(golds, predicts):
g = g[i:i+1]
p = p[i:i+1]
n_gold += sum(g)
n_predict += sum(p >= th)
n_correct += sum(g * (p > th))
p = n_correct / n_predict
r = n_correct / n_gold
f1 = 2 * p * r / (p + r)
print(config.idx2tag[i], n_gold, n_predict, n_correct, p, r, f1)
n_gold, n_predict, n_correct = 0, 0, 0
for i in range(1, len(config.idx2tag)):
if i in config.salient_type_set:
continue
th = 0.3
for g, p in zip(golds, predicts):
g = g[i:i+1]
p = p[i:i+1]
n_gold += sum(g)
n_predict += sum(p >= th)
n_correct += sum(g * (p > th))
p = n_correct / n_predict
r = n_correct / n_gold
f1 = 2 * p * r / (p + r)
print(config.idx2tag[i], n_gold, n_predict, n_correct, p, r, f1)
if __name__ == '__main__':
device = 'cuda'
batch_size = 10
train_data, val_data, test_data = load_dataset()
train_dataset = Dataset(batch_size, 150, train_data)
val_dataset = Dataset(batch_size, 100, val_data)
test_dataset = Dataset(batch_size, 80, test_data)
rep = get_type_rep().detach_().to(device)
model = SentenceClassification(config.bert_dir, len(config.idx2tag), rep)
load_model(model, 'model/ace_sentence9.ckp')
model.to(device)
model.eval()
evaluate(model, device, test_dataset)