mirror of
https://github.com/explosion/spaCy.git
synced 2024-09-23 12:29:18 +03:00
299 lines
13 KiB
Python
299 lines
13 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 thinc.neural._classes.convolution import ExtractWindow
|
|
|
|
from examples.pipeline.wiki_entity_linking import run_el, training_set_creator, kb_creator
|
|
|
|
from spacy._ml import SpacyVectors, create_default_optimizer, zero_init, logistic
|
|
|
|
from thinc.api import chain, concatenate, flatten_add_lengths, with_getitem, clone, with_flatten
|
|
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, max_pool
|
|
from thinc.t2t import ParametricAttention
|
|
from thinc.misc import Residual
|
|
from thinc.misc import LayerNorm as LN
|
|
|
|
from spacy.tokens import Doc
|
|
|
|
""" TODO: this code needs to be implemented in pipes.pyx"""
|
|
|
|
|
|
class EL_Model:
|
|
|
|
PRINT_LOSS = True
|
|
PRINT_F = True
|
|
EPS = 0.0000000005
|
|
CUTOFF = 0.5
|
|
|
|
INPUT_DIM = 300
|
|
ENTITY_WIDTH = 64
|
|
ARTICLE_WIDTH = 64
|
|
HIDDEN_1_WIDTH = 256
|
|
HIDDEN_2_WIDTH = 64
|
|
|
|
name = "entity_linker"
|
|
|
|
def __init__(self, kb, nlp):
|
|
run_el._prepare_pipeline(nlp, kb)
|
|
self.nlp = nlp
|
|
self.kb = kb
|
|
|
|
self._build_cnn(hidden_entity_width=self.ENTITY_WIDTH, hidden_article_width=self.ARTICLE_WIDTH)
|
|
|
|
def train_model(self, training_dir, entity_descr_output, trainlimit=None, devlimit=None, to_print=True):
|
|
# raise errors instead of runtime warnings in case of int/float overflow
|
|
np.seterr(all='raise')
|
|
|
|
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=False)
|
|
|
|
dev_instances, dev_pos, dev_neg, dev_doc = self._get_training_data(training_dir,
|
|
entity_descr_output,
|
|
True,
|
|
devlimit,
|
|
to_print=False)
|
|
self._begin_training()
|
|
|
|
if self.PRINT_F:
|
|
_, _, f_avg_train = -3.42, -3.42, -3.42 # self._test_dev(train_instances, train_pos, train_neg, train_doc, avg=True)
|
|
_, _, f_nonavg_train = self._test_dev(train_instances, train_pos, train_neg, train_doc, avg=False)
|
|
_, _, f_random_train = self._test_dev(train_instances, train_pos, train_neg, train_doc, calc_random=True)
|
|
_, _, f_avg_dev = -3.42, -3.42, -3.42 # self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc, avg=True)
|
|
_, _, f_nonavg_dev = self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc, avg=False)
|
|
_, _, f_random_dev = self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc, calc_random=True)
|
|
|
|
print("random F train", round(f_random_train, 1))
|
|
print("random F dev", round(f_random_dev, 1))
|
|
print()
|
|
print("avg/nonavg F train", round(f_avg_train, 1), round(f_nonavg_train, 1))
|
|
print("avg/nonavg F dev", round(f_avg_dev, 1), round(f_nonavg_dev, 1))
|
|
print()
|
|
|
|
instance_pos_count = 0
|
|
instance_neg_count = 0
|
|
|
|
if to_print:
|
|
print("Training on", len(train_instances.values()), "articles")
|
|
print("Dev test on", len(dev_instances.values()), "articles")
|
|
print()
|
|
|
|
article_docs = list()
|
|
entities = list()
|
|
golds = list()
|
|
for article_id, inst_cluster_set in train_instances.items():
|
|
for inst_cluster in inst_cluster_set:
|
|
article_docs.append(train_doc[article_id])
|
|
entities.append(train_pos.get(inst_cluster))
|
|
golds.append(float(1.0))
|
|
instance_pos_count += 1
|
|
for neg_entity in train_neg.get(inst_cluster, []):
|
|
article_docs.append(train_doc[article_id])
|
|
entities.append(neg_entity)
|
|
golds.append(float(0.0))
|
|
instance_neg_count += 1
|
|
|
|
for x in range(10):
|
|
print("Updating", x)
|
|
self.update(article_docs=article_docs, entities=entities, golds=golds)
|
|
|
|
# eval again
|
|
if self.PRINT_F:
|
|
_, _, f_avg_train = -3.42, -3.42, -3.42 # self._test_dev(train_instances, train_pos, train_neg, train_doc, avg=True)
|
|
_, _, f_nonavg_train = self._test_dev(train_instances, train_pos, train_neg, train_doc, avg=False)
|
|
_, _, f_avg_dev = -3.42, -3.42, -3.42 # self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc, avg=True)
|
|
_, _, f_nonavg_dev = self._test_dev(dev_instances, dev_pos, dev_neg, dev_doc, avg=False)
|
|
|
|
print("avg/nonavg F train", round(f_avg_train, 1), round(f_nonavg_train, 1))
|
|
print("avg/nonavg F dev", round(f_avg_dev, 1), round(f_nonavg_dev, 1))
|
|
print()
|
|
|
|
if to_print:
|
|
print("Trained on", instance_pos_count, "/", instance_neg_count, "instances pos/neg")
|
|
|
|
def _test_dev(self, dev_instances, dev_pos, dev_neg, dev_doc, avg=False, calc_random=False):
|
|
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, [])
|
|
|
|
article = inst_cluster.split(sep="_")[0]
|
|
entity_id = inst_cluster.split(sep="_")[1]
|
|
article_doc = dev_doc[article]
|
|
|
|
if calc_random:
|
|
prediction = self._predict_random(entity=pos_ex)
|
|
else:
|
|
prediction = self._predict(article_doc=article_doc, entity=pos_ex, avg=avg)
|
|
predictions.append(prediction)
|
|
golds.append(float(1.0))
|
|
|
|
for neg_ex in neg_exs:
|
|
if calc_random:
|
|
prediction = self._predict_random(entity=neg_ex)
|
|
else:
|
|
prediction = self._predict(article_doc=article_doc, entity=neg_ex, avg=avg)
|
|
predictions.append(prediction)
|
|
golds.append(float(0.0))
|
|
|
|
# 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, article_doc, entity, avg=False, apply_threshold=True):
|
|
if avg:
|
|
with self.sgd.use_params(self.model.averages):
|
|
doc_encoding = self.article_encoder([article_doc])
|
|
entity_encoding = self.entity_encoder([entity])
|
|
return self.model(np.append(entity_encoding, doc_encoding)) # TODO list
|
|
|
|
doc_encoding = self.article_encoder([article_doc])[0]
|
|
entity_encoding = self.entity_encoder([entity])[0]
|
|
concat_encoding = list(entity_encoding) + list(doc_encoding)
|
|
np_array = np.asarray([concat_encoding])
|
|
prediction = self.model(np_array)
|
|
if not apply_threshold:
|
|
return float(prediction)
|
|
if prediction > self.CUTOFF:
|
|
return float(1.0)
|
|
return float(0.0)
|
|
|
|
def _predict_random(self, entity, apply_threshold=True):
|
|
r = random.uniform(0, 1)
|
|
if not apply_threshold:
|
|
return r
|
|
if r > self.CUTOFF:
|
|
return float(1.0)
|
|
return float(0.0)
|
|
|
|
def _build_cnn(self, hidden_entity_width, hidden_article_width):
|
|
with Model.define_operators({">>": chain, "|": concatenate, "**": clone}):
|
|
self.entity_encoder = self._encoder(in_width=self.INPUT_DIM, hidden_width=hidden_entity_width) # entity encoding
|
|
self.article_encoder = self._encoder(in_width=self.INPUT_DIM, hidden_width=hidden_article_width) # doc encoding
|
|
|
|
hidden_input_with = hidden_entity_width + hidden_article_width
|
|
hidden_output_with = self.HIDDEN_1_WIDTH
|
|
|
|
convolution_2 = Residual((ExtractWindow(nW=1) >> LN(Maxout(hidden_output_with, hidden_output_with * 3))))
|
|
|
|
self.model = Affine(hidden_output_with, hidden_input_with) \
|
|
>> LN(Maxout(hidden_output_with, hidden_output_with)) \
|
|
>> convolution_2 \
|
|
>> Affine(self.HIDDEN_2_WIDTH, hidden_output_with) \
|
|
>> Affine(1, self.HIDDEN_2_WIDTH) \
|
|
>> logistic
|
|
|
|
@staticmethod
|
|
def _encoder(in_width, hidden_width):
|
|
with Model.define_operators({">>": chain}):
|
|
encoder = SpacyVectors \
|
|
>> flatten_add_lengths \
|
|
>> ParametricAttention(in_width)\
|
|
>> Pooling(mean_pool) \
|
|
>> Residual(zero_init(Maxout(in_width, in_width))) \
|
|
>> zero_init(Affine(hidden_width, in_width, drop_factor=0.0))
|
|
|
|
return encoder
|
|
|
|
def _begin_training(self):
|
|
self.sgd = create_default_optimizer(self.model.ops)
|
|
|
|
def update(self, article_docs, entities, golds, drop=0.):
|
|
doc_encodings, bp_doc = self.article_encoder.begin_update(article_docs, drop=drop)
|
|
entity_encodings, bp_encoding = self.entity_encoder.begin_update(entities, drop=drop)
|
|
concat_encodings = [list(entity_encodings[i]) + list(doc_encodings[i]) for i in range(len(entities))]
|
|
|
|
predictions, bp_model = self.model.begin_update(np.asarray(concat_encodings), drop=drop)
|
|
|
|
predictions = self.model.ops.flatten(predictions)
|
|
golds = self.model.ops.asarray(golds)
|
|
|
|
# print("predictions", predictions)
|
|
# print("golds", golds)
|
|
|
|
d_scores = (predictions - golds) # / predictions.shape[0]
|
|
# print("d_scores (1)", d_scores)
|
|
|
|
loss = (d_scores ** 2).sum()
|
|
|
|
if self.PRINT_LOSS:
|
|
print("loss train", round(loss, 5))
|
|
|
|
d_scores = d_scores.reshape((-1, 1))
|
|
d_scores = d_scores.astype(np.float32)
|
|
# print("d_scores (2)", d_scores)
|
|
|
|
model_gradient = bp_model(d_scores, sgd=self.sgd)
|
|
|
|
doc_gradient = [x[0:self.ARTICLE_WIDTH] for x in model_gradient]
|
|
entity_gradient = [x[self.ARTICLE_WIDTH:] for x in model_gradient]
|
|
|
|
bp_doc(doc_gradient)
|
|
bp_encoding(entity_gradient)
|
|
|
|
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
|