spaCy/examples/pipeline/wiki_entity_linking/train_el.py

# 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 examples.pipeline.wiki_entity_linking import run_el, training_set_creator, kb_creator

from spacy._ml import SpacyVectors, create_default_optimizer, zero_init, logistic, Tok2Vec

from thinc.api import chain, concatenate, flatten_add_lengths, clone, with_flatten
from thinc.v2v import Model, Maxout, Affine, ReLu
from thinc.t2v import Pooling, mean_pool, sum_pool
from thinc.t2t import ParametricAttention
from thinc.misc import Residual
from thinc.misc import LayerNorm as LN

from spacy.matcher import PhraseMatcher
from spacy.tokens import Doc

""" TODO: this code needs to be implemented in pipes.pyx"""


class EL_Model:

    PRINT_INSPECT = False
    PRINT_TRAIN = True
    EPS = 0.0000000005
    CUTOFF = 0.5

    BATCH_SIZE = 5
    UPSAMPLE = True

    DOC_CUTOFF = 300    # number of characters from the doc context
    INPUT_DIM = 300     # dimension of pre-trained vectors

    # HIDDEN_1_WIDTH = 32   # 10
    HIDDEN_2_WIDTH = 32  # 6
    DESC_WIDTH = 64     # 4
    ARTICLE_WIDTH = 64   # 8
    SENT_WIDTH = 64

    DROP = 0.1
    LEARN_RATE = 0.0001
    EPOCHS = 20
    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(desc_width=self.DESC_WIDTH,
                        article_width=self.ARTICLE_WIDTH,
                        sent_width=self.SENT_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
        # (not sure if we need this. set L2 to 0 because it throws an error otherwsise)
        # np.seterr(all='raise')
        # alternative:
        np.seterr(divide="raise", over="warn", under="ignore", invalid="raise")

        train_ent, train_gold, train_desc, train_art, train_art_texts, train_sent, train_sent_texts = \
            self._get_training_data(training_dir, entity_descr_output, False, trainlimit, to_print=False)

        dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts = \
            self._get_training_data(training_dir, entity_descr_output, True, devlimit, to_print=False)

        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 entity in train_ent:
                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)

        if self.UPSAMPLE:
            if to_print:
                print()
                print("Upsampling, original training instances pos/neg:", train_pos_count, train_neg_count)

            # upsample positives to 50-50 distribution
            while train_pos_count < train_neg_count:
                train_ent.append(random.choice(train_pos_entities))
                train_pos_count += 1

            # upsample negatives to 50-50 distribution
            while train_neg_count < train_pos_count:
                train_ent.append(random.choice(train_neg_entities))
                train_neg_count += 1

        self._begin_training()

        if to_print:
            print()
            print("Training on", len(train_ent), "entities in", len(train_art_texts), "articles")
            print("Training instances pos/neg:", train_pos_count, train_neg_count)
            print()
            print("Dev test on", len(dev_ent), "entities in", len(dev_art_texts), "articles")
            print("Dev instances pos/neg:", dev_pos_count, dev_neg_count)
            print()
            print(" CUTOFF", self.CUTOFF)
            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(" HIDDEN_2_WIDTH", self.HIDDEN_2_WIDTH)
            print(" DROP", self.DROP)
            print(" LEARNING RATE", self.LEARN_RATE)
            print(" UPSAMPLE", self.UPSAMPLE)
            print()

        self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts,
                       print_string="dev_random", calc_random=True)

        self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts,
                       print_string="dev_pre", avg=True)
        print()

        processed = 0
        for i in range(self.EPOCHS):
            shuffle(train_ent)

            start = 0
            stop = min(self.BATCH_SIZE, len(train_ent))

            while start < len(train_ent):
                next_batch = train_ent[start:stop]

                golds = [train_gold[e] for e in next_batch]
                descs = [train_desc[e] for e in next_batch]
                article_texts = [train_art_texts[train_art[e]] for e in next_batch]
                sent_texts = [train_sent_texts[train_sent[e]] for e in next_batch]

                self.update(entities=next_batch, golds=golds, descs=descs, art_texts=article_texts, sent_texts=sent_texts)

                processed += len(next_batch)

                start = start + self.BATCH_SIZE
                stop = min(stop + self.BATCH_SIZE, len(train_ent))

            if self.PRINT_TRAIN:
                print()
                self._test_dev(train_ent, train_gold, train_desc, train_art, train_art_texts, train_sent, train_sent_texts,
                               print_string="train_inter_epoch " + str(i), avg=True)

            self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts,
                           print_string="dev_inter_epoch " + str(i), avg=True)

        if to_print:
            print()
            print("Trained on", processed, "entities across", self.EPOCHS, "epochs")

    def _test_dev(self, entities, gold_by_entity, desc_by_entity, art_by_entity, art_texts, sent_by_entity, sent_texts,
                  print_string, avg=True, calc_random=False):
        golds = [gold_by_entity[e] for e in entities]

        if calc_random:
            predictions = self._predict_random(entities=entities)

        else:
            desc_docs = self.nlp.pipe([desc_by_entity[e] for e in entities])
            article_docs = self.nlp.pipe([art_texts[art_by_entity[e]] for e in entities])
            sent_docs = self.nlp.pipe([sent_texts[sent_by_entity[e]] for e in entities])
            predictions = self._predict(entities=entities, article_docs=article_docs, sent_docs=sent_docs,
                                        desc_docs=desc_docs, avg=avg)

        # TODO: combine with prior probability
        p, r, f, acc = run_el.evaluate(predictions, golds, to_print=False, times_hundred=False)
        loss, gradient = self.get_loss(self.model.ops.asarray(predictions), self.model.ops.asarray(golds))

        print("p/r/F/acc/loss", print_string, round(p, 2), round(r, 2), round(f, 2), round(acc, 2), round(loss, 2))

        return loss, p, r, f

    def _predict(self, entities, article_docs, sent_docs, 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):
                doc_encodings = self.article_encoder(article_docs)
                desc_encodings = self.desc_encoder(desc_docs)
                sent_encodings = self.sent_encoder(sent_docs)

        else:
            doc_encodings = self.article_encoder(article_docs)
            desc_encodings = self.desc_encoder(desc_docs)
            sent_encodings = self.sent_encoder(sent_docs)

        concat_encodings = [list(doc_encodings[i]) + list(sent_encodings[i]) + list(desc_encodings[i]) for i in
                            range(len(entities))]

        np_array_list = np.asarray(concat_encodings)

        if avg:
            with self.model.use_params(self.sgd.averages):
                predictions = self.model(np_array_list)
        else:
            predictions = self.model(np_array_list)

        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]

        return predictions

    def _predict_random(self, entities, apply_threshold=True):
        if not apply_threshold:
            return [float(random.uniform(0, 1)) for _ in entities]
        else:
            return [float(1.0) if random.uniform(0, 1) > self.CUTOFF else float(0.0) for _ in entities]

    def _build_cnn_depr(self, embed_width, desc_width, article_width, sent_width, hidden_1_width, hidden_2_width):
        with Model.define_operators({">>": chain, "|": concatenate, "**": clone}):
            self.desc_encoder = self._encoder_depr(in_width=embed_width, hidden_with=hidden_1_width, end_width=desc_width)
            self.article_encoder = self._encoder_depr(in_width=embed_width, hidden_with=hidden_1_width, end_width=article_width)
            self.sent_encoder = self._encoder_depr(in_width=embed_width, hidden_with=hidden_1_width, end_width=sent_width)

            in_width = article_width + sent_width + desc_width
            out_width = hidden_2_width

            self.model = Affine(out_width, in_width) \
                >> LN(Maxout(out_width, out_width)) \
                >> Affine(1, out_width) \
                >> logistic

    def _build_cnn(self, desc_width, article_width, sent_width):
        with Model.define_operators({">>": chain, "|": concatenate, "**": clone}):
            self.desc_encoder = self._encoder(width=desc_width)
            self.article_encoder = self._encoder(width=article_width)
            self.sent_encoder = self._encoder(width=sent_width)

            in_width = desc_width + article_width + sent_width

            self.model = Affine(self.HIDDEN_2_WIDTH, in_width) \
                         >> LN(Maxout(self.HIDDEN_2_WIDTH, self.HIDDEN_2_WIDTH)) \
                         >> Affine(1, self.HIDDEN_2_WIDTH) \
                         >> logistic

            # output_layer = (
            #         zero_init(Affine(1, in_width, drop_factor=0.0)) >> logistic
            # )
            # self.model = output_layer
            self.model.nO = 1

    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)

    @staticmethod
    def _encoder_depr(in_width, hidden_with, end_width):
        conv_depth = 2
        cnn_maxout_pieces = 3

        with Model.define_operators({">>": chain}):
            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_desc = create_default_optimizer(self.desc_encoder.ops)
        self.sgd_desc.learn_rate = self.LEARN_RATE
        self.sgd_desc.L2 = self.L2

        self.sgd = create_default_optimizer(self.model.ops)
        self.sgd.learn_rate = self.LEARN_RATE
        self.sgd.L2 = self.L2

    @staticmethod
    def get_loss(predictions, golds):
        d_scores = (predictions - golds)
        gradient = d_scores.mean()
        loss = (d_scores ** 2).mean()
        return loss, gradient

    def update(self, entities, golds, descs, art_texts, sent_texts):
        golds = self.model.ops.asarray(golds)

        art_docs = self.nlp.pipe(art_texts)
        sent_docs = self.nlp.pipe(sent_texts)
        desc_docs = self.nlp.pipe(descs)

        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)
        desc_encodings, bp_desc = self.desc_encoder.begin_update(desc_docs, drop=self.DROP)

        concat_encodings = [list(doc_encodings[i]) + list(sent_encodings[i]) + list(desc_encodings[i])
                            for i in range(len(entities))]

        predictions, bp_model = self.model.begin_update(np.asarray(concat_encodings), drop=self.DROP)
        predictions = self.model.ops.flatten(predictions)

        # print("entities", entities)
        # print("predictions", predictions)
        # print("golds", golds)

        loss, gradient = self.get_loss(predictions, golds)

        gradient = float(gradient)
        # print("gradient", gradient)
        # print("loss", loss)

        model_gradient = bp_model(gradient, sgd=self.sgd)
        # print("model_gradient", model_gradient)

        # concat = doc + sent + desc, but doc is the same within this function
        sent_start = self.ARTICLE_WIDTH
        desc_start = self.ARTICLE_WIDTH + self.SENT_WIDTH
        doc_gradient = model_gradient[0][0:sent_start]
        sent_gradients = list()
        desc_gradients = list()
        for x in model_gradient:
            sent_gradients.append(list(x[sent_start:desc_start]))
            desc_gradients.append(list(x[desc_start:]))

        # print("doc_gradient", doc_gradient)
        # print("sent_gradients", sent_gradients)
        # print("desc_gradients", desc_gradients)

        bp_doc([doc_gradient], sgd=self.sgd_article)
        bp_sent(sent_gradients, sgd=self.sgd_sent)
        bp_desc(desc_gradients, sgd=self.sgd_desc)

    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)

        entities = set()
        gold_by_entity = dict()
        desc_by_entity = dict()
        article_by_entity = dict()
        text_by_article = dict()
        sentence_by_entity = dict()
        text_by_sentence = 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")
                    cnt += 1

                    # 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
                    article_terms = set()
                    entities_by_mention = dict()

                    for mention, entity_pos in correct_entries[article_id].items():
                        descr = id_to_descr.get(entity_pos)
                        if descr:
                            entity = "E_" + str(next_entity_nr) + "_" + article_id + "_" + mention
                            next_entity_nr += 1
                            gold_by_entity[entity] = 1
                            desc_by_entity[entity] = descr
                            article_terms.add(mention)
                            mention_entities = entities_by_mention.get(mention, set())
                            mention_entities.add(entity)
                            entities_by_mention[mention] = mention_entities

                    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:
                                entity = "E_" + str(next_entity_nr) + "_" + article_id + "_" + mention
                                next_entity_nr += 1
                                gold_by_entity[entity] = 0
                                desc_by_entity[entity] = descr
                                article_terms.add(mention)
                                mention_entities = entities_by_mention.get(mention, set())
                                mention_entities.add(entity)
                                entities_by_mention[mention] = mention_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(article_terms))

                    matcher.add("TerminologyList", None, *patterns)
                    matches = matcher(article_doc)

                    # store sentences
                    sentence_to_id = dict()
                    for match_id, start, end in matches:
                        span = article_doc[start:end]
                        sent_text = span.sent.text
                        sent_nr = sentence_to_id.get(sent_text,  None)
                        mention = span.text
                        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_to_id[sent_text] = sent_nr
                        mention_entities = entities_by_mention[mention]
                        for entity in mention_entities:
                            entities.add(entity)
                            sentence_by_entity[entity] = sent_nr
                            article_by_entity[entity] = article_id

        # remove entities that didn't have all data
        gold_by_entity = {k: v for k, v in gold_by_entity.items() if k in entities}
        desc_by_entity = {k: v for k, v in desc_by_entity.items() if k in entities}

        article_by_entity = {k: v for k, v in article_by_entity.items() if k in entities}
        text_by_article = {k: v for k, v in text_by_article.items() if k in article_by_entity.values()}

        sentence_by_entity = {k: v for k, v in sentence_by_entity.items() if k in entities}
        text_by_sentence = {k: v for k, v in text_by_sentence.items() if k in sentence_by_entity.values()}

        if to_print:
            print()
            print("Processed", cnt, "training articles, dev=" + str(dev))
            print()
        return list(entities), gold_by_entity, desc_by_entity, article_by_entity, text_by_article, \
               sentence_by_entity, text_by_sentence