def evaluate(self, datasets):
"""
Parent function r esponsible for evaluating the ED model during the ED step. Note that
this is different from predict as this requires ground truth entities to be present.
:return: -
"""
dev_datasets = []
for dname, data in list(datasets.items()):
start = time.time()
dev_datasets.append(
(dname, self.get_data_items(data, dname, predict=True)))
for dname, data in dev_datasets:
predictions = self.__predict(data)
f1, recall, precision, total_nil = self.__eval(
datasets[dname], predictions)
print(
dname,
utils.tokgreen(
"Micro F1: {}, Recall: {}, Precision: {}".format(
f1, recall, precision)),
)
print("Total NIL: {}".format(total_nil))
print("----------------------------------")
def train_LR(self, datasets, model_path_lr, store_offline=True, threshold=0.3):
"""
Function that applies LR in an attempt to get confidence scores. Recall should be high,
because if it is low than we would have ignored a corrrect entity.
:return: -
"""
train_dataset = self.get_data_items(datasets['aida_train'], "train", predict=False)
dev_datasets = []
for dname, data in list(datasets.items()):
if dname == 'aida_train':
continue
dev_datasets.append((dname, self.get_data_items(data, dname, predict=True)))
model = LogisticRegression()
predictions = self.__predict(train_dataset, eval_raw=True)
X, y, meta = self.__create_dataset_LR(datasets, predictions, 'aida_train')
model.fit(X, y)
for dname, data in dev_datasets:
predictions = self.__predict(data, eval_raw=True)
X, y, meta = self.__create_dataset_LR(datasets, predictions, dname)
preds = model.predict_proba(X)
preds = np.array([x[1] for x in preds])
decisions = (preds >= threshold).astype(int)
print(utils.tokgreen('{}, F1-score: {}'.format(dname, f1_score(y, decisions))))
if store_offline:
with open('{}/lr_model.pkl'.format(model_path_lr), 'wb') as handle:
pkl.dump(model, handle, protocol=pkl.HIGHEST_PROTOCOL)
def train(self, org_train_dataset, org_dev_datasets):
"""
Responsible for training the ED model.
:return: -
"""
train_dataset = self.get_data_items(org_train_dataset,
"train",
predict=False)
dev_datasets = []
for dname, data in org_dev_datasets.items():
dev_datasets.append(
(dname, self.get_data_items(data, dname, predict=True)))
print("Creating optimizer")
optimizer = optim.Adam(
[p for p in self.model.parameters() if p.requires_grad],
lr=self.config["learning_rate"],
)
best_f1 = -1
not_better_count = 0
eval_after_n_epochs = self.config["eval_after_n_epochs"]
for e in range(self.config["n_epochs"]):
shuffle(train_dataset)
total_loss = 0
for dc, batch in enumerate(
train_dataset): # each document is a minibatch
self.model.train()
optimizer.zero_grad()
# convert data items to pytorch inputs
token_ids = [
m["context"][0] + m["context"][1]
if len(m["context"][0]) + len(m["context"][1]) > 0 else
[self.embeddings["word_voca"].unk_id] for m in batch
]
s_ltoken_ids = [m["snd_ctx"][0] for m in batch]
s_rtoken_ids = [m["snd_ctx"][1] for m in batch]
s_mtoken_ids = [m["snd_ment"] for m in batch]
entity_ids = Variable(
torch.LongTensor([
m["selected_cands"]["cands"] for m in batch
]).to(self.device))
true_pos = Variable(
torch.LongTensor([
m["selected_cands"]["true_pos"] for m in batch
]).to(self.device))
p_e_m = Variable(
torch.FloatTensor([
m["selected_cands"]["p_e_m"] for m in batch
]).to(self.device))
entity_mask = Variable(
torch.FloatTensor([
m["selected_cands"]["mask"] for m in batch
]).to(self.device))
token_ids, token_mask = utils.make_equal_len(
token_ids, self.embeddings["word_voca"].unk_id)
s_ltoken_ids, s_ltoken_mask = utils.make_equal_len(
s_ltoken_ids,
self.embeddings["snd_voca"].unk_id,
to_right=False)
s_rtoken_ids, s_rtoken_mask = utils.make_equal_len(
s_rtoken_ids, self.embeddings["snd_voca"].unk_id)
s_rtoken_ids = [l[::-1] for l in s_rtoken_ids]
s_rtoken_mask = [l[::-1] for l in s_rtoken_mask]
s_mtoken_ids, s_mtoken_mask = utils.make_equal_len(
s_mtoken_ids, self.embeddings["snd_voca"].unk_id)
token_ids = Variable(
torch.LongTensor(token_ids).to(self.device))
token_mask = Variable(
torch.FloatTensor(token_mask).to(self.device))
# too ugly but too lazy to fix it
self.model.s_ltoken_ids = Variable(
torch.LongTensor(s_ltoken_ids).to(self.device))
self.model.s_ltoken_mask = Variable(
torch.FloatTensor(s_ltoken_mask).to(self.device))
self.model.s_rtoken_ids = Variable(
torch.LongTensor(s_rtoken_ids).to(self.device))
self.model.s_rtoken_mask = Variable(
torch.FloatTensor(s_rtoken_mask).to(self.device))
self.model.s_mtoken_ids = Variable(
torch.LongTensor(s_mtoken_ids).to(self.device))
self.model.s_mtoken_mask = Variable(
torch.FloatTensor(s_mtoken_mask).to(self.device))
scores, ent_scores = self.model.forward(
token_ids,
token_mask,
entity_ids,
entity_mask,
p_e_m,
self.embeddings,
gold=true_pos.view(-1, 1),
)
loss = self.model.loss(scores, true_pos)
# loss = self.model.prob_loss(scores, true_pos)
loss.backward()
optimizer.step()
self.model.regularize(max_norm=100)
loss = loss.cpu().data.numpy()
total_loss += loss
print(
"epoch",
e,
"%0.2f%%" % (dc / len(train_dataset) * 100),
loss,
end="\r",
)
print("epoch", e, "total loss", total_loss,
total_loss / len(train_dataset))
if (e + 1) % eval_after_n_epochs == 0:
dev_f1 = 0
for dname, data in dev_datasets:
predictions = self.__predict(data)
f1, recall, precision, _ = self.__eval(
org_dev_datasets[dname], predictions)
print(
dname,
utils.tokgreen(
"Micro F1: {}, Recall: {}, Precision: {}".format(
f1, recall, precision)),
)
if dname == "aida_testA":
dev_f1 = f1
if (self.config["learning_rate"] == 1e-4
and dev_f1 >= self.config["dev_f1_change_lr"]):
eval_after_n_epochs = 2
best_f1 = dev_f1
not_better_count = 0
self.config["learning_rate"] = 1e-5
print("change learning rate to",
self.config["learning_rate"])
for param_group in optimizer.param_groups:
param_group["lr"] = self.config["learning_rate"]
if dev_f1 < best_f1:
not_better_count += 1
print("Not improving", not_better_count)
else:
not_better_count = 0
best_f1 = dev_f1
print("save model to", self.config["model_path"])
self.__save(self.config["model_path"])
if not_better_count == self.config["n_not_inc"]:
break
