spaCy/examples/pipeline/wiki_entity_linking/train_el.py

493 lines
21 KiB
Python

# coding: utf-8
from __future__ import unicode_literals
import os
import datetime
from os import listdir
import numpy as np
import random
from random import shuffle
from thinc.neural._classes.convolution import ExtractWindow
from thinc.neural.util import get_array_module
from examples.pipeline.wiki_entity_linking import run_el, training_set_creator, kb_creator
from spacy._ml import SpacyVectors, create_default_optimizer, zero_init, cosine
from thinc.api import chain, concatenate, flatten_add_lengths, clone, with_flatten
from thinc.v2v import Model, Maxout, Affine
from thinc.t2v import Pooling, mean_pool
from thinc.t2t import ParametricAttention
from thinc.misc import Residual
from thinc.misc import LayerNorm as LN
# from spacy.cli.pretrain import get_cossim_loss
from spacy.matcher import PhraseMatcher
class EL_Model:
PRINT_INSPECT = False
PRINT_BATCH_LOSS = False
EPS = 0.0000000005
BATCH_SIZE = 100
DOC_CUTOFF = 300 # number of characters from the doc context
INPUT_DIM = 300 # dimension of pre-trained vectors
HIDDEN_1_WIDTH = 32
DESC_WIDTH = 64
ARTICLE_WIDTH = 128
SENT_WIDTH = 64
DROP = 0.4
LEARN_RATE = 0.005
EPOCHS = 10
L2 = 1e-6
name = "entity_linker"
def __init__(self, kb, nlp):
run_el._prepare_pipeline(nlp, kb)
self.nlp = nlp
self.kb = kb
self._build_cnn(embed_width=self.INPUT_DIM,
desc_width=self.DESC_WIDTH,
article_width=self.ARTICLE_WIDTH,
sent_width=self.SENT_WIDTH,
hidden_1_width=self.HIDDEN_1_WIDTH)
def train_model(self, training_dir, entity_descr_output, trainlimit=None, devlimit=None, to_print=True):
np.seterr(divide="raise", over="warn", under="ignore", invalid="raise")
id_to_descr = kb_creator._get_id_to_description(entity_descr_output)
train_ent, train_gold, train_desc, train_art, train_art_texts, train_sent, train_sent_texts = \
self._get_training_data(training_dir, id_to_descr, False, trainlimit, to_print=False)
train_clusters = list(train_ent.keys())
dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts = \
self._get_training_data(training_dir, id_to_descr, True, devlimit, to_print=False)
dev_clusters = list(dev_ent.keys())
dev_pos_count = len([g for g in dev_gold.values() if g])
dev_neg_count = len([g for g in dev_gold.values() if not g])
# inspect data
if self.PRINT_INSPECT:
for cluster, entities in train_ent.items():
print()
for entity in entities:
print("entity", entity)
print("gold", train_gold[entity])
print("desc", train_desc[entity])
print("sentence ID", train_sent[entity])
print("sentence text", train_sent_texts[train_sent[entity]])
print("article ID", train_art[entity])
print("article text", train_art_texts[train_art[entity]])
print()
train_pos_entities = [k for k, v in train_gold.items() if v]
train_neg_entities = [k for k, v in train_gold.items() if not v]
train_pos_count = len(train_pos_entities)
train_neg_count = len(train_neg_entities)
self._begin_training()
if to_print:
print()
print("Training on", len(train_clusters), "entity clusters in", len(train_art_texts), "articles")
print("Training instances pos/neg:", train_pos_count, train_neg_count)
print()
print("Dev test on", len(dev_clusters), "entity clusters in", len(dev_art_texts), "articles")
print("Dev instances pos/neg:", dev_pos_count, dev_neg_count)
print()
print(" DOC_CUTOFF", self.DOC_CUTOFF)
print(" INPUT_DIM", self.INPUT_DIM)
print(" HIDDEN_1_WIDTH", self.HIDDEN_1_WIDTH)
print(" DESC_WIDTH", self.DESC_WIDTH)
print(" ARTICLE_WIDTH", self.ARTICLE_WIDTH)
print(" SENT_WIDTH", self.SENT_WIDTH)
print(" DROP", self.DROP)
print(" LEARNING RATE", self.LEARN_RATE)
print(" BATCH SIZE", self.BATCH_SIZE)
print()
dev_random = self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts,
calc_random=True)
print("acc", "dev_random", round(dev_random, 2))
dev_pre = self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts,
avg=True)
print("acc", "dev_pre", round(dev_pre, 2))
print()
processed = 0
for i in range(self.EPOCHS):
shuffle(train_clusters)
start = 0
stop = min(self.BATCH_SIZE, len(train_clusters))
while start < len(train_clusters):
next_batch = {c: train_ent[c] for c in train_clusters[start:stop]}
processed += len(next_batch.keys())
self.update(entity_clusters=next_batch, golds=train_gold, descs=train_desc,
art_texts=train_art_texts, arts=train_art,
sent_texts=train_sent_texts, sents=train_sent)
start = start + self.BATCH_SIZE
stop = min(stop + self.BATCH_SIZE, len(train_clusters))
train_acc = self._test_dev(train_ent, train_gold, train_desc, train_art, train_art_texts, train_sent, train_sent_texts, avg=True)
dev_acc = self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts, avg=True)
print(i, "acc train/dev", round(train_acc, 2), round(dev_acc, 2))
if to_print:
print()
print("Trained on", processed, "entity clusters across", self.EPOCHS, "epochs")
def _test_dev(self, entity_clusters, golds, descs, arts, art_texts, sents, sent_texts, avg=True, calc_random=False):
correct = 0
incorrect = 0
if calc_random:
for cluster, entities in entity_clusters.items():
correct_entities = [e for e in entities if golds[e]]
assert len(correct_entities) == 1
entities = list(entities)
shuffle(entities)
if calc_random:
predicted_entity = random.choice(entities)
if predicted_entity in correct_entities:
correct += 1
else:
incorrect += 1
else:
all_clusters = list()
arts_list = list()
sents_list = list()
for cluster in entity_clusters.keys():
all_clusters.append(cluster)
arts_list.append(art_texts[arts[cluster]])
sents_list.append(sent_texts[sents[cluster]])
art_docs = list(self.nlp.pipe(arts_list))
sent_docs = list(self.nlp.pipe(sents_list))
for i, cluster in enumerate(all_clusters):
entities = entity_clusters[cluster]
correct_entities = [e for e in entities if golds[e]]
assert len(correct_entities) == 1
entities = list(entities)
shuffle(entities)
desc_docs = self.nlp.pipe([descs[e] for e in entities])
sent_doc = sent_docs[i]
article_doc = art_docs[i]
predicted_index = self._predict(article_doc=article_doc, sent_doc=sent_doc,
desc_docs=desc_docs, avg=avg)
if entities[predicted_index] in correct_entities:
correct += 1
else:
incorrect += 1
if correct == incorrect == 0:
return 0
acc = correct / (correct + incorrect)
return acc
def _predict(self, article_doc, sent_doc, desc_docs, avg=True, apply_threshold=True):
if avg:
with self.article_encoder.use_params(self.sgd_article.averages) \
and self.desc_encoder.use_params(self.sgd_desc.averages)\
and self.sent_encoder.use_params(self.sgd_sent.averages):
desc_encodings = self.desc_encoder(desc_docs)
doc_encoding = self.article_encoder([article_doc])
sent_encoding = self.sent_encoder([sent_doc])
else:
desc_encodings = self.desc_encoder(desc_docs)
doc_encoding = self.article_encoder([article_doc])
sent_encoding = self.sent_encoder([sent_doc])
concat_encoding = [list(doc_encoding[0]) + list(sent_encoding[0])]
if avg:
with self.cont_encoder.use_params(self.sgd_cont.averages):
cont_encodings = self.cont_encoder(np.asarray([concat_encoding[0]]))
else:
cont_encodings = self.cont_encoder(np.asarray([concat_encoding[0]]))
context_enc = np.transpose(cont_encodings)
highest_sim = -5
best_i = -1
for i, desc_enc in enumerate(desc_encodings):
sim = cosine(desc_enc, context_enc)
if sim >= highest_sim:
best_i = i
highest_sim = sim
return best_i
def _build_cnn(self, embed_width, desc_width, article_width, sent_width, hidden_1_width):
self.desc_encoder = self._encoder(in_width=embed_width, hidden_with=hidden_1_width, end_width=desc_width)
self.cont_encoder = self._context_encoder(embed_width=embed_width, article_width=article_width,
sent_width=sent_width, hidden_width=hidden_1_width,
end_width=desc_width)
# def _encoder(self, width):
# tok2vec = Tok2Vec(width=width, embed_size=2000, pretrained_vectors=self.nlp.vocab.vectors.name, cnn_maxout_pieces=3,
# subword_features=False, conv_depth=4, bilstm_depth=0)
#
# return tok2vec >> flatten_add_lengths >> Pooling(mean_pool)
def _context_encoder(self, embed_width, article_width, sent_width, hidden_width, end_width):
self.article_encoder = self._encoder(in_width=embed_width, hidden_with=hidden_width, end_width=article_width)
self.sent_encoder = self._encoder(in_width=embed_width, hidden_with=hidden_width, end_width=sent_width)
model = Affine(end_width, article_width+sent_width, drop_factor=0.0)
return model
@staticmethod
def _encoder(in_width, hidden_with, end_width):
conv_depth = 2
cnn_maxout_pieces = 3
with Model.define_operators({">>": chain, "**": clone}):
convolution = Residual((ExtractWindow(nW=1) >>
LN(Maxout(hidden_with, hidden_with * 3, pieces=cnn_maxout_pieces))))
encoder = SpacyVectors \
>> with_flatten(LN(Maxout(hidden_with, in_width)) >> convolution ** conv_depth, pad=conv_depth) \
>> flatten_add_lengths \
>> ParametricAttention(hidden_with)\
>> Pooling(mean_pool) \
>> Residual(zero_init(Maxout(hidden_with, hidden_with))) \
>> zero_init(Affine(end_width, hidden_with, drop_factor=0.0))
# TODO: ReLu or LN(Maxout) ?
# sum_pool or mean_pool ?
return encoder
def _begin_training(self):
self.sgd_article = create_default_optimizer(self.article_encoder.ops)
self.sgd_article.learn_rate = self.LEARN_RATE
self.sgd_article.L2 = self.L2
self.sgd_sent = create_default_optimizer(self.sent_encoder.ops)
self.sgd_sent.learn_rate = self.LEARN_RATE
self.sgd_sent.L2 = self.L2
self.sgd_cont = create_default_optimizer(self.cont_encoder.ops)
self.sgd_cont.learn_rate = self.LEARN_RATE
self.sgd_cont.L2 = self.L2
self.sgd_desc = create_default_optimizer(self.desc_encoder.ops)
self.sgd_desc.learn_rate = self.LEARN_RATE
self.sgd_desc.L2 = self.L2
def get_loss(self, pred, gold, targets):
loss, gradients = self.get_cossim_loss(pred, gold, targets)
return loss, gradients
def get_cossim_loss(self, yh, y, t):
# Add a small constant to avoid 0 vectors
# print()
# print("yh", yh)
# print("y", y)
# print("t", t)
yh = yh + 1e-8
y = y + 1e-8
# https://math.stackexchange.com/questions/1923613/partial-derivative-of-cosine-similarity
xp = get_array_module(yh)
norm_yh = xp.linalg.norm(yh, axis=1, keepdims=True)
norm_y = xp.linalg.norm(y, axis=1, keepdims=True)
mul_norms = norm_yh * norm_y
cos = (yh * y).sum(axis=1, keepdims=True) / mul_norms
# print("cos", cos)
d_yh = (y / mul_norms) - (cos * (yh / norm_yh ** 2))
# print("abs", xp.abs(cos - t))
loss = xp.abs(cos - t).sum()
# print("loss", loss)
# print("d_yh", d_yh)
inverse = np.asarray([int(t[i][0]) * d_yh[i] for i in range(len(t))])
# print("inverse", inverse)
return loss, -inverse
def update(self, entity_clusters, golds, descs, art_texts, arts, sent_texts, sents):
arts_list = list()
sents_list = list()
descs_list = list()
targets = list()
for cluster, entities in entity_clusters.items():
art = art_texts[arts[cluster]]
sent = sent_texts[sents[cluster]]
for e in entities:
if golds[e]:
arts_list.append(art)
sents_list.append(sent)
descs_list.append(descs[e])
targets.append([1])
# else:
# arts_list.append(art)
# sents_list.append(sent)
# descs_list.append(descs[e])
# targets.append([-1])
desc_docs = self.nlp.pipe(descs_list)
desc_encodings, bp_desc = self.desc_encoder.begin_update(desc_docs, drop=self.DROP)
art_docs = self.nlp.pipe(arts_list)
sent_docs = self.nlp.pipe(sents_list)
doc_encodings, bp_doc = self.article_encoder.begin_update(art_docs, drop=self.DROP)
sent_encodings, bp_sent = self.sent_encoder.begin_update(sent_docs, drop=self.DROP)
concat_encodings = [list(doc_encodings[i]) + list(sent_encodings[i]) for i in
range(len(targets))]
cont_encodings, bp_cont = self.cont_encoder.begin_update(np.asarray(concat_encodings), drop=self.DROP)
loss, cont_gradient = self.get_loss(cont_encodings, desc_encodings, targets)
# loss, desc_gradient = self.get_loss(desc_encodings, cont_encodings, targets)
# cont_gradient = cont_gradient / 2
# desc_gradient = desc_gradient / 2
# bp_desc(desc_gradient, sgd=self.sgd_desc)
if self.PRINT_BATCH_LOSS:
print("batch loss", loss)
context_gradient = bp_cont(cont_gradient, sgd=self.sgd_cont)
# gradient : concat (doc+sent) vs. desc
sent_start = self.ARTICLE_WIDTH
sent_gradients = list()
doc_gradients = list()
for x in context_gradient:
doc_gradients.append(list(x[0:sent_start]))
sent_gradients.append(list(x[sent_start:]))
bp_doc(doc_gradients, sgd=self.sgd_article)
bp_sent(sent_gradients, sgd=self.sgd_sent)
def _get_training_data(self, training_dir, id_to_descr, dev, limit, to_print):
correct_entries, incorrect_entries = training_set_creator.read_training_entities(training_output=training_dir,
collect_correct=True,
collect_incorrect=True)
entities_by_cluster = dict()
gold_by_entity = dict()
desc_by_entity = dict()
article_by_cluster = dict()
text_by_article = dict()
sentence_by_cluster = dict()
text_by_sentence = dict()
sentence_by_text = dict()
cnt = 0
next_entity_nr = 1
next_sent_nr = 1
files = listdir(training_dir)
shuffle(files)
for f in files:
if not limit or cnt < limit:
if dev == run_el.is_dev(f):
article_id = f.replace(".txt", "")
if cnt % 500 == 0 and to_print:
print(datetime.datetime.now(), "processed", cnt, "files in the training dataset")
try:
# parse the article text
with open(os.path.join(training_dir, f), mode="r", encoding='utf8') as file:
text = file.read()
article_doc = self.nlp(text)
truncated_text = text[0:min(self.DOC_CUTOFF, len(text))]
text_by_article[article_id] = truncated_text
# process all positive and negative entities, collect all relevant mentions in this article
for mention, entity_pos in correct_entries[article_id].items():
cluster = article_id + "_" + mention
descr = id_to_descr.get(entity_pos)
entities = set()
if descr:
entity = "E_" + str(next_entity_nr) + "_" + cluster
next_entity_nr += 1
gold_by_entity[entity] = 1
desc_by_entity[entity] = descr
entities.add(entity)
entity_negs = incorrect_entries[article_id][mention]
for entity_neg in entity_negs:
descr = id_to_descr.get(entity_neg)
if descr:
entity = "E_" + str(next_entity_nr) + "_" + cluster
next_entity_nr += 1
gold_by_entity[entity] = 0
desc_by_entity[entity] = descr
entities.add(entity)
found_matches = 0
if len(entities) > 1:
entities_by_cluster[cluster] = entities
# find all matches in the doc for the mentions
# TODO: fix this - doesn't look like all entities are found
matcher = PhraseMatcher(self.nlp.vocab)
patterns = list(self.nlp.tokenizer.pipe([mention]))
matcher.add("TerminologyList", None, *patterns)
matches = matcher(article_doc)
# store sentences
for match_id, start, end in matches:
span = article_doc[start:end]
if mention == span.text:
found_matches += 1
sent_text = span.sent.text
sent_nr = sentence_by_text.get(sent_text, None)
if sent_nr is None:
sent_nr = "S_" + str(next_sent_nr) + article_id
next_sent_nr += 1
text_by_sentence[sent_nr] = sent_text
sentence_by_text[sent_text] = sent_nr
article_by_cluster[cluster] = article_id
sentence_by_cluster[cluster] = sent_nr
if found_matches == 0:
# print("Could not find neg instances or sentence matches for", mention, "in", article_id)
entities_by_cluster.pop(cluster, None)
article_by_cluster.pop(cluster, None)
sentence_by_cluster.pop(cluster, None)
for entity in entities:
gold_by_entity.pop(entity, None)
desc_by_entity.pop(entity, None)
cnt += 1
except:
print("Problem parsing article", article_id)
if to_print:
print()
print("Processed", cnt, "training articles, dev=" + str(dev))
print()
return entities_by_cluster, gold_by_entity, desc_by_entity, article_by_cluster, text_by_article, \
sentence_by_cluster, text_by_sentence