# 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 = 5 # 96 PRINT_LOSS = True PRINT_F = True 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): for i, true_entity in enumerate(true_entity_list): for cnt in range(10): #try: false_vectors = list() false_entities = false_entities_list[i] if len(false_entities) > 0: # TODO: batch per doc doc_encoding, article_bp = self.article_encoder.begin_update([article_doc], drop=drop) doc_encoding = doc_encoding[0] print() print(cnt) print("doc", doc_encoding) 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] 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) 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) all_probs.append(false_prob) loss = self._calculate_loss(true_prob, all_probs).astype(np.float32) if self.PRINT_LOSS: print("loss", round(loss, 5)) doc_gradient = self._calculate_doc_gradient(loss, doc_encoding, true_entity_encoding, false_vectors) print("doc_gradient", doc_gradient) article_bp([doc_gradient.astype(np.float32)], sgd=self.sgd_article) #except Exception as e: #pass # TODO: FIX def _calculate_consensus(self, vector1, vector2): if len(vector1) != len(vector2): raise ValueError("To calculate consenus, 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 / e_sum) @staticmethod def _calculate_loss(true_prob, all_probs): """ all_probs should include true_prob ! """ return -1 * np.log(true_prob / 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 @staticmethod def _calculate_dot_exp(vector1, vector2_transposed): e = np.exp(vector1.dot(vector2_transposed)) 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