diff --git a/examples/pipeline/wiki_entity_linking/train_el.py b/examples/pipeline/wiki_entity_linking/train_el.py index e213f0955..2d218ed60 100644 --- a/examples/pipeline/wiki_entity_linking/train_el.py +++ b/examples/pipeline/wiki_entity_linking/train_el.py @@ -52,27 +52,25 @@ class EL_Model: # raise errors instead of runtime warnings in case of int/float overflow np.seterr(all='raise') - Doc.set_extension("entity_id", default=None) + train_inst, train_pos, train_neg, train_texts = self._get_training_data(training_dir, + entity_descr_output, + False, + trainlimit, + to_print=False) - train_inst, train_pos, train_neg, train_doc = self._get_training_data(training_dir, - entity_descr_output, - False, - trainlimit, - to_print=False) - - dev_inst, dev_pos, dev_neg, dev_doc = self._get_training_data(training_dir, - entity_descr_output, - True, - devlimit, - to_print=False) + dev_inst, dev_pos, dev_neg, dev_texts = self._get_training_data(training_dir, + entity_descr_output, + True, + devlimit, + to_print=False) self._begin_training() print() - self._test_dev(train_inst, train_pos, train_neg, train_doc, print_string="train_random", calc_random=True) - self._test_dev(dev_inst, dev_pos, dev_neg, dev_doc, print_string="dev_random", calc_random=True) + self._test_dev(train_inst, train_pos, train_neg, train_texts, print_string="train_random", calc_random=True) + self._test_dev(dev_inst, dev_pos, dev_neg, dev_texts, print_string="dev_random", calc_random=True) print() - self._test_dev(train_inst, train_pos, train_neg, train_doc, print_string="train_pre", avg=False) - self._test_dev(dev_inst, dev_pos, dev_neg, dev_doc, print_string="dev_pre", avg=False) + self._test_dev(train_inst, train_pos, train_neg, train_texts, print_string="train_pre", avg=False) + self._test_dev(dev_inst, dev_pos, dev_neg, dev_texts, print_string="dev_pre", avg=False) instance_pos_count = 0 instance_neg_count = 0 @@ -99,26 +97,22 @@ class EL_Model: # print() # print(article_count, "Training on article", article_id) article_count += 1 - article_docs = list() + article_text = train_texts[article_id] entities = list() golds = list() 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 - self.update(article_docs=article_docs, entities=entities, golds=golds) + self.update(article_text=article_text, entities=entities, golds=golds) # dev eval - # self._test_dev(dev_inst, dev_pos, dev_neg, dev_doc, print_string="dev_inter", avg=False) - self._test_dev(dev_inst, dev_pos, dev_neg, dev_doc, print_string="dev_inter_avg", avg=True) - print() + self._test_dev(dev_inst, dev_pos, dev_neg, dev_texts, print_string="dev_inter_avg", avg=True) except ValueError as e: print("Error in article id", article_id) @@ -127,13 +121,9 @@ class EL_Model: print("Trained on", instance_pos_count, "/", instance_neg_count, "instances pos/neg") if self.PRINT_TRAIN: - # print() - # self._test_dev(train_inst, train_pos, train_neg, train_doc, print_string="train_post", avg=False) - self._test_dev(train_inst, train_pos, train_neg, train_doc, print_string="train_post_avg", avg=True) - # self._test_dev(dev_inst, dev_pos, dev_neg, dev_doc, print_string="dev_post", avg=False) - # self._test_dev(dev_inst, dev_pos, dev_neg, dev_doc, print_string="dev_post_avg", avg=True) + self._test_dev(train_inst, train_pos, train_neg, train_texts, print_string="train_post_avg", avg=True) - def _test_dev(self, instances, pos, neg, doc, print_string, avg=False, calc_random=False): + def _test_dev(self, instances, pos, neg, texts_by_id, print_string, avg=False, calc_random=False): predictions = list() golds = list() @@ -144,22 +134,18 @@ class EL_Model: article = inst_cluster.split(sep="_")[0] entity_id = inst_cluster.split(sep="_")[1] - article_doc = doc[article] + article_doc = self.nlp(texts_by_id[article]) + entities = [self.nlp(pos_ex)] + golds.append(float(1.0)) + for neg_ex in neg_exs: + entities.append(self.nlp(neg_ex)) + golds.append(float(0.0)) if calc_random: - prediction = self._predict_random(entity=pos_ex) + preds = self._predict_random(entities=entities) 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)) + preds = self._predict(article_doc=article_doc, entities=entities, avg=avg) + predictions.extend(preds) # TODO: combine with prior probability p, r, f = run_el.evaluate(predictions, golds, to_print=False) @@ -172,39 +158,38 @@ class EL_Model: return loss, p, r, f - def _predict(self, article_doc, entity, avg=False, apply_threshold=True): + def _predict(self, article_doc, entities, avg=False, apply_threshold=True): if avg: with self.article_encoder.use_params(self.sgd_article.averages) \ and self.entity_encoder.use_params(self.sgd_entity.averages): doc_encoding = self.article_encoder([article_doc])[0] - entity_encoding = self.entity_encoder([entity])[0] + entity_encodings = self.entity_encoder(entities) else: doc_encoding = self.article_encoder([article_doc])[0] - entity_encoding = self.entity_encoder([entity])[0] + entity_encodings = self.entity_encoder(entities) - concat_encoding = list(entity_encoding) + list(doc_encoding) - np_array = np.asarray([concat_encoding]) + concat_encodings = [list(entity_encodings[i]) + list(doc_encoding) for i in range(len(entities))] + np_array_list = np.asarray(concat_encodings) if avg: - with self.model.use_params(self.sgd.averages): - prediction = self.model(np_array) + with self.model.use_params(self.sgd.averages): + predictions = self.model(np_array_list) else: - prediction = self.model(np_array) + predictions = self.model(np_array_list) - if not apply_threshold: - return float(prediction) - if prediction > self.CUTOFF: - return float(1.0) - return float(0.0) + predictions = self.model.ops.flatten(predictions) + predictions = [float(p) for p in predictions] + if apply_threshold: + predictions = [float(1.0) if p > self.CUTOFF else float(0.0) for p in predictions] - def _predict_random(self, entity, apply_threshold=True): - r = random.uniform(0, 1) + return predictions + + def _predict_random(self, entities, apply_threshold=True): if not apply_threshold: - return r - if r > self.CUTOFF: - return float(1.0) - return float(0.0) + return [float(random.uniform(0,1)) for e in entities] + else: + return [float(1.0) if random.uniform(0,1) > self.CUTOFF else float(0.0) for e in entities] def _build_cnn(self, hidden_entity_width, hidden_article_width): with Model.define_operators({">>": chain, "|": concatenate, "**": clone}): @@ -252,20 +237,27 @@ class EL_Model: loss = (d_scores ** 2).sum() return loss, d_scores - def update(self, article_docs, entities, golds, apply_threshold=True): - doc_encodings, bp_doc = self.article_encoder.begin_update(article_docs, drop=self.DROP) - # print("doc_encodings", len(doc_encodings), doc_encodings) + # TODO: multiple docs/articles + def update(self, article_text, entities, golds, apply_threshold=True): + article_doc = self.nlp(article_text) + doc_encodings, bp_doc = self.article_encoder.begin_update([article_doc], drop=self.DROP) + doc_encoding = doc_encodings[0] - entity_encodings, bp_entity = self.entity_encoder.begin_update(entities, drop=self.DROP) + entity_docs = list(self.nlp.pipe(entities)) + # print("entity_docs", type(entity_docs)) + + entity_encodings, bp_entity = self.entity_encoder.begin_update(entity_docs, drop=self.DROP) # print("entity_encodings", len(entity_encodings), entity_encodings) - concat_encodings = [list(entity_encodings[i]) + list(doc_encodings[i]) for i in range(len(entities))] + concat_encodings = [list(entity_encodings[i]) + list(doc_encoding) for i in range(len(entities))] # print("concat_encodings", len(concat_encodings), concat_encodings) predictions, bp_model = self.model.begin_update(np.asarray(concat_encodings), drop=self.DROP) predictions = self.model.ops.flatten(predictions) + # print("predictions", predictions) golds = self.model.ops.asarray(golds) + # print("golds", golds) loss, d_scores = self.get_loss(predictions, golds) @@ -275,7 +267,7 @@ class EL_Model: if self.PRINT_F and self.PRINT_TRAIN: predictions_f = [x for x in predictions] if apply_threshold: - predictions_f = [1.0 if x > self.CUTOFF else 0.0 for x in predictions_f] + predictions_f = [float(1.0) if x > self.CUTOFF else float(0.0) for x in predictions_f] p, r, f = run_el.evaluate(predictions_f, golds, to_print=False) print("p/r/F train", round(p, 1), round(r, 1), round(f, 1)) @@ -286,17 +278,17 @@ class EL_Model: model_gradient = bp_model(d_scores, sgd=self.sgd) # print("model_gradient", model_gradient) - doc_gradient = list() - entity_gradient = list() + # concat = entity + doc, but doc is the same within this function (TODO: multiple docs/articles) + doc_gradient = model_gradient[0][self.ENTITY_WIDTH:] + entity_gradients = list() for x in model_gradient: - doc_gradient.append(list(x[0:self.ARTICLE_WIDTH])) - entity_gradient.append(list(x[self.ARTICLE_WIDTH:])) + entity_gradients.append(list(x[0:self.ENTITY_WIDTH])) # print("doc_gradient", doc_gradient) - # print("entity_gradient", entity_gradient) + # print("entity_gradients", entity_gradients) - bp_doc(doc_gradient, sgd=self.sgd_article) - bp_entity(entity_gradient, sgd=self.sgd_entity) + bp_doc([doc_gradient], sgd=self.sgd_article) + bp_entity(entity_gradients, sgd=self.sgd_entity) 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) @@ -305,9 +297,9 @@ class EL_Model: collect_correct=True, collect_incorrect=True) - instance_by_doc = dict() + instance_by_article = dict() local_vectors = list() # TODO: local vectors - doc_by_article = dict() + text_by_article = dict() pos_entities = dict() neg_entities = dict() @@ -319,33 +311,28 @@ class EL_Model: 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: + if article_id not in text_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() + text_by_article[article_id] = text + instance_by_article[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 + instance_by_article[article_id].add(article_id + "_" + mention) + pos_entities[article_id + "_" + mention] = 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) + descr_list.append(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 + return instance_by_article, pos_entities, neg_entities, text_by_article diff --git a/examples/pipeline/wiki_entity_linking/wiki_nel_pipeline.py b/examples/pipeline/wiki_entity_linking/wiki_nel_pipeline.py index 6f021597f..23c12bfe6 100644 --- a/examples/pipeline/wiki_entity_linking/wiki_nel_pipeline.py +++ b/examples/pipeline/wiki_entity_linking/wiki_nel_pipeline.py @@ -111,7 +111,7 @@ if __name__ == "__main__": print("STEP 6: training", datetime.datetime.now()) my_nlp = spacy.load('en_core_web_md') trainer = EL_Model(kb=my_kb, nlp=my_nlp) - trainer.train_model(training_dir=TRAINING_DIR, entity_descr_output=ENTITY_DESCR, trainlimit=1000, devlimit=200) + trainer.train_model(training_dir=TRAINING_DIR, entity_descr_output=ENTITY_DESCR, trainlimit=10, devlimit=10) print() # STEP 7: apply the EL algorithm on the dev dataset