mirror of
https://github.com/explosion/spaCy.git
synced 2024-11-11 20:28:20 +03:00
493 lines
21 KiB
Python
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
|
|
|