mirror of
https://github.com/explosion/spaCy.git
synced 2024-09-22 20:09:18 +03:00
385 lines
16 KiB
Python
385 lines
16 KiB
Python
# coding: utf-8
|
|
from __future__ import unicode_literals
|
|
|
|
import os
|
|
import datetime
|
|
from os import listdir
|
|
from random import shuffle
|
|
import numpy as np
|
|
|
|
from examples.pipeline.wiki_entity_linking import run_el, training_set_creator, kb_creator
|
|
|
|
from spacy._ml import SpacyVectors, create_default_optimizer, zero_init
|
|
|
|
from thinc.api import chain, flatten_add_lengths, with_getitem, clone
|
|
from thinc.neural.util import get_array_module
|
|
from thinc.v2v import Model, Softmax, Maxout, Affine, ReLu
|
|
from thinc.t2v import Pooling, sum_pool, mean_pool
|
|
from thinc.t2t import ParametricAttention
|
|
from thinc.misc import Residual
|
|
|
|
from spacy.tokens import Doc
|
|
|
|
""" TODO: this code needs to be implemented in pipes.pyx"""
|
|
|
|
|
|
class EL_Model():
|
|
|
|
INPUT_DIM = 300
|
|
OUTPUT_DIM = 96
|
|
PRINT_LOSS = False
|
|
PRINT_F = True
|
|
EPS = 0.0000000005
|
|
|
|
labels = ["MATCH", "NOMATCH"]
|
|
name = "entity_linker"
|
|
|
|
def __init__(self, kb, nlp):
|
|
run_el._prepare_pipeline(nlp, kb)
|
|
self.nlp = nlp
|
|
self.kb = kb
|
|
|
|
self.entity_encoder = self._simple_encoder(in_width=self.INPUT_DIM, out_width=self.OUTPUT_DIM)
|
|
self.article_encoder = self._simple_encoder(in_width=self.INPUT_DIM, out_width=self.OUTPUT_DIM)
|
|
|
|
def train_model(self, training_dir, entity_descr_output, trainlimit=None, devlimit=None, to_print=True):
|
|
Doc.set_extension("entity_id", default=None)
|
|
|
|
train_instances, train_pos, train_neg, train_doc = self._get_training_data(training_dir,
|
|
entity_descr_output,
|
|
False,
|
|
trainlimit,
|
|
to_print)
|
|
|
|
dev_instances, dev_pos, dev_neg, dev_doc = self._get_training_data(training_dir,
|
|
entity_descr_output,
|
|
True,
|
|
devlimit,
|
|
to_print)
|
|
|
|
if to_print:
|
|
print("Training on", len(train_instances.values()), "articles")
|
|
print("Dev test on", len(dev_instances.values()), "articles")
|
|
print()
|
|
|
|
self.sgd_entity = self.begin_training(self.entity_encoder)
|
|
self.sgd_article = self.begin_training(self.article_encoder)
|
|
|
|
self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc)
|
|
|
|
losses = {}
|
|
|
|
instance_count = 0
|
|
|
|
for article_id, inst_cluster_set in train_instances.items():
|
|
# print("article", article_id)
|
|
article_doc = train_doc[article_id]
|
|
pos_ex_list = list()
|
|
neg_exs_list = list()
|
|
for inst_cluster in inst_cluster_set:
|
|
# print("inst_cluster", inst_cluster)
|
|
instance_count += 1
|
|
pos_ex_list.append(train_pos.get(inst_cluster))
|
|
neg_exs_list.append(train_neg.get(inst_cluster, []))
|
|
|
|
self.update(article_doc, pos_ex_list, neg_exs_list, losses=losses)
|
|
p, r, fscore = self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc)
|
|
if self.PRINT_F:
|
|
print(round(fscore, 1))
|
|
|
|
if to_print:
|
|
print("Trained on", instance_count, "instance clusters")
|
|
|
|
def _test_dev(self, dev_instances, dev_pos, dev_neg, dev_doc):
|
|
predictions = list()
|
|
golds = list()
|
|
|
|
for article_id, inst_cluster_set in dev_instances.items():
|
|
for inst_cluster in inst_cluster_set:
|
|
pos_ex = dev_pos.get(inst_cluster)
|
|
neg_exs = dev_neg.get(inst_cluster, [])
|
|
ex_to_id = dict()
|
|
|
|
if pos_ex and neg_exs:
|
|
ex_to_id[pos_ex] = pos_ex._.entity_id
|
|
for neg_ex in neg_exs:
|
|
ex_to_id[neg_ex] = neg_ex._.entity_id
|
|
|
|
article = inst_cluster.split(sep="_")[0]
|
|
entity_id = inst_cluster.split(sep="_")[1]
|
|
article_doc = dev_doc[article]
|
|
|
|
examples = list(neg_exs)
|
|
examples.append(pos_ex)
|
|
shuffle(examples)
|
|
|
|
best_entity, highest_prob = self._predict(examples, article_doc)
|
|
predictions.append(ex_to_id[best_entity])
|
|
golds.append(ex_to_id[pos_ex])
|
|
|
|
# TODO: use lowest_mse and combine with prior probability
|
|
p, r, F = run_el.evaluate(predictions, golds, to_print=False)
|
|
return p, r, F
|
|
|
|
def _predict(self, entities, article_doc):
|
|
doc_encoding = self.article_encoder([article_doc])
|
|
|
|
highest_prob = None
|
|
best_entity = None
|
|
|
|
entity_to_vector = dict()
|
|
for entity in entities:
|
|
entity_to_vector[entity] = self.entity_encoder([entity])
|
|
|
|
for entity in entities:
|
|
entity_encoding = entity_to_vector[entity]
|
|
prob = self._calculate_probability(doc_encoding, entity_encoding, entity_to_vector.values())
|
|
if not best_entity or prob > highest_prob:
|
|
highest_prob = prob
|
|
best_entity = entity
|
|
|
|
return best_entity, highest_prob
|
|
|
|
def _simple_encoder(self, in_width, out_width):
|
|
conv_depth = 1
|
|
cnn_maxout_pieces = 3
|
|
with Model.define_operators({">>": chain, "**": clone}):
|
|
# encoder = SpacyVectors \
|
|
# >> flatten_add_lengths \
|
|
# >> ParametricAttention(in_width)\
|
|
# >> Pooling(mean_pool) \
|
|
# >> Residual(zero_init(Maxout(in_width, in_width))) \
|
|
# >> zero_init(Affine(out_width, in_width, drop_factor=0.0))
|
|
encoder = SpacyVectors \
|
|
>> flatten_add_lengths \
|
|
>> with_getitem(0, Affine(in_width, in_width)) \
|
|
>> ParametricAttention(in_width) \
|
|
>> Pooling(sum_pool) \
|
|
>> Residual(ReLu(in_width, in_width)) ** conv_depth \
|
|
>> zero_init(Affine(out_width, in_width, drop_factor=0.0))
|
|
|
|
# >> zero_init(Affine(nr_class, width, drop_factor=0.0))
|
|
# >> logistic
|
|
|
|
# convolution = Residual(
|
|
# ExtractWindow(nW=1)
|
|
# >> LN(Maxout(width, width * 3, pieces=cnn_maxout_pieces))
|
|
# )
|
|
|
|
# embed = SpacyVectors >> LN(Maxout(width, width, pieces=3))
|
|
|
|
# encoder = SpacyVectors >> flatten_add_lengths >> convolution ** conv_depth
|
|
# encoder = with_flatten(embed >> convolution ** conv_depth, pad=conv_depth)
|
|
|
|
return encoder
|
|
|
|
def begin_training(self, model):
|
|
# TODO ? link_vectors_to_models(self.vocab)
|
|
sgd = create_default_optimizer(model.ops)
|
|
return sgd
|
|
|
|
def update(self, article_doc, true_entity_list, false_entities_list, drop=0., losses=None):
|
|
doc_encoding, article_bp = self.article_encoder.begin_update([article_doc], drop=drop)
|
|
doc_encoding = doc_encoding[0]
|
|
# print("doc", doc_encoding)
|
|
|
|
for i, true_entity in enumerate(true_entity_list):
|
|
try:
|
|
false_vectors = list()
|
|
false_entities = false_entities_list[i]
|
|
if len(false_entities) > 0:
|
|
# TODO: batch per doc
|
|
|
|
for false_entity in false_entities:
|
|
# TODO: one call only to begin_update ?
|
|
false_entity_encoding, false_entity_bp = self.entity_encoder.begin_update([false_entity], drop=drop)
|
|
false_entity_encoding = false_entity_encoding[0]
|
|
false_vectors.append(false_entity_encoding)
|
|
|
|
true_entity_encoding, true_entity_bp = self.entity_encoder.begin_update([true_entity], drop=drop)
|
|
true_entity_encoding = true_entity_encoding[0]
|
|
# true_gradient = self._calculate_true_gradient(doc_encoding, true_entity_encoding)
|
|
|
|
all_vectors = [true_entity_encoding]
|
|
all_vectors.extend(false_vectors)
|
|
|
|
# consensus_encoding = self._calculate_consensus(doc_encoding, true_entity_encoding)
|
|
|
|
true_prob = self._calculate_probability(doc_encoding, true_entity_encoding, all_vectors)
|
|
# print("true", true_prob, true_entity_encoding)
|
|
# print("true gradient", true_gradient)
|
|
# print()
|
|
|
|
all_probs = [true_prob]
|
|
for false_vector in false_vectors:
|
|
false_prob = self._calculate_probability(doc_encoding, false_vector, all_vectors)
|
|
# print("false", false_prob, false_vector)
|
|
# print("false gradient", false_gradient)
|
|
# print()
|
|
all_probs.append(false_prob)
|
|
|
|
loss = self._calculate_loss(true_prob, all_probs).astype(np.float32)
|
|
if self.PRINT_LOSS:
|
|
print(round(loss, 5))
|
|
|
|
#doc_gradient = self._calculate_doc_gradient(loss, doc_encoding, true_entity_encoding, false_vectors)
|
|
entity_gradient = self._calculate_entity_gradient(doc_encoding, true_entity_encoding, false_vectors)
|
|
# print("entity_gradient", entity_gradient)
|
|
# print("doc_gradient", doc_gradient)
|
|
# article_bp([doc_gradient.astype(np.float32)], sgd=self.sgd_article)
|
|
true_entity_bp([entity_gradient.astype(np.float32)], sgd=self.sgd_entity)
|
|
#true_entity_bp([true_gradient.astype(np.float32)], sgd=self.sgd_entity)
|
|
except Exception as e:
|
|
pass
|
|
|
|
|
|
# TODO: FIX
|
|
def _calculate_consensus(self, vector1, vector2):
|
|
if len(vector1) != len(vector2):
|
|
raise ValueError("To calculate consensus, both vectors should be of equal length")
|
|
|
|
avg = (vector2 + vector1) / 2
|
|
return avg
|
|
|
|
def _calculate_probability(self, vector1, vector2, allvectors):
|
|
""" Make sure that vector2 is included in allvectors """
|
|
if len(vector1) != len(vector2):
|
|
raise ValueError("To calculate similarity, both vectors should be of equal length")
|
|
|
|
vector1_t = vector1.transpose()
|
|
e = self._calculate_dot_exp(vector2, vector1_t)
|
|
e_sum = 0
|
|
for v in allvectors:
|
|
e_sum += self._calculate_dot_exp(v, vector1_t)
|
|
|
|
return float(e / (self.EPS + e_sum))
|
|
|
|
def _calculate_loss(self, true_prob, all_probs):
|
|
""" all_probs should include true_prob ! """
|
|
return -1 * np.log((self.EPS + true_prob) / (self.EPS + sum(all_probs)))
|
|
|
|
@staticmethod
|
|
def _calculate_doc_gradient(loss, doc_vector, true_vector, false_vectors):
|
|
gradient = np.zeros(len(doc_vector))
|
|
for i in range(len(doc_vector)):
|
|
min_false = min(x[i] for x in false_vectors)
|
|
max_false = max(x[i] for x in false_vectors)
|
|
|
|
if true_vector[i] > max_false:
|
|
if doc_vector[i] > 0:
|
|
gradient[i] = 0
|
|
else:
|
|
gradient[i] = -loss
|
|
elif true_vector[i] < min_false:
|
|
if doc_vector[i] > 0:
|
|
gradient[i] = loss
|
|
if doc_vector[i] < 0:
|
|
gradient[i] = 0
|
|
else:
|
|
target = 0 # non-distinctive vector positions should convert to 0
|
|
gradient[i] = doc_vector[i] - target
|
|
|
|
return gradient
|
|
|
|
def _calculate_true_gradient(self, doc_vector, entity_vector):
|
|
# sum_entity_vector = sum(entity_vector)
|
|
# gradient = [-sum_entity_vector/(self.EPS + np.exp(doc_vector[i] * entity_vector[i])) for i in range(len(doc_vector))]
|
|
gradient = [1 / (self.EPS + np.exp(doc_vector[i] * entity_vector[i])) for i in range(len(doc_vector))]
|
|
return np.asarray(gradient)
|
|
|
|
def _calculate_entity_gradient(self, doc_vector, true_vector, false_vectors):
|
|
entity_gradient = list()
|
|
prob_true = list()
|
|
false_prob_list = list()
|
|
for i in range(len(true_vector)):
|
|
doc_i = np.asarray([doc_vector[i]])
|
|
true_i = np.asarray([true_vector[i]])
|
|
falses_i = np.asarray([[fv[i]] for fv in false_vectors])
|
|
all_i = [true_i]
|
|
all_i.extend(falses_i)
|
|
|
|
prob_true_i = self._calculate_probability(doc_i, true_i, all_i)
|
|
prob_true.append(prob_true_i)
|
|
|
|
false_list = list()
|
|
all_probs_i = [prob_true_i]
|
|
for false_vector in falses_i:
|
|
false_prob_i = self._calculate_probability(doc_i, false_vector, all_i)
|
|
all_probs_i.append(false_prob_i)
|
|
false_list.append(false_prob_i)
|
|
false_prob_list.append(false_list)
|
|
|
|
sign_loss_i = 1
|
|
if doc_vector[i] * true_vector[i] < 0:
|
|
sign_loss_i = -1
|
|
|
|
loss_i = sign_loss_i * self._calculate_loss(prob_true_i, all_probs_i).astype(np.float32)
|
|
entity_gradient.append(loss_i)
|
|
# print("prob_true", prob_true)
|
|
# print("false_prob_list", false_prob_list)
|
|
return np.asarray(entity_gradient)
|
|
|
|
|
|
@staticmethod
|
|
def _calculate_dot_exp(vector1, vector2_transposed):
|
|
dot_product = vector1.dot(vector2_transposed)
|
|
dot_product = min(50, dot_product)
|
|
# dot_product = max(-10000, dot_product)
|
|
# print("DOT", dot_product)
|
|
e = np.exp(dot_product)
|
|
# print("E", e)
|
|
return e
|
|
|
|
def _get_training_data(self, training_dir, entity_descr_output, dev, limit, to_print):
|
|
id_to_descr = kb_creator._get_id_to_description(entity_descr_output)
|
|
|
|
correct_entries, incorrect_entries = training_set_creator.read_training_entities(training_output=training_dir,
|
|
collect_correct=True,
|
|
collect_incorrect=True)
|
|
|
|
|
|
instance_by_doc = dict()
|
|
local_vectors = list() # TODO: local vectors
|
|
doc_by_article = dict()
|
|
pos_entities = dict()
|
|
neg_entities = dict()
|
|
|
|
cnt = 0
|
|
for f in listdir(training_dir):
|
|
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")
|
|
cnt += 1
|
|
if article_id not in doc_by_article:
|
|
with open(os.path.join(training_dir, f), mode="r", encoding='utf8') as file:
|
|
text = file.read()
|
|
doc = self.nlp(text)
|
|
doc_by_article[article_id] = doc
|
|
instance_by_doc[article_id] = set()
|
|
|
|
for mention, entity_pos in correct_entries[article_id].items():
|
|
descr = id_to_descr.get(entity_pos)
|
|
if descr:
|
|
instance_by_doc[article_id].add(article_id + "_" + mention)
|
|
doc_descr = self.nlp(descr)
|
|
doc_descr._.entity_id = entity_pos
|
|
pos_entities[article_id + "_" + mention] = doc_descr
|
|
|
|
for mention, entity_negs in incorrect_entries[article_id].items():
|
|
for entity_neg in entity_negs:
|
|
descr = id_to_descr.get(entity_neg)
|
|
if descr:
|
|
doc_descr = self.nlp(descr)
|
|
doc_descr._.entity_id = entity_neg
|
|
descr_list = neg_entities.get(article_id + "_" + mention, [])
|
|
descr_list.append(doc_descr)
|
|
neg_entities[article_id + "_" + mention] = descr_list
|
|
|
|
if to_print:
|
|
print()
|
|
print("Processed", cnt, "training articles, dev=" + str(dev))
|
|
print()
|
|
return instance_by_doc, pos_entities, neg_entities, doc_by_article
|