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 thinc.neural.util import get_array_module

from examples.pipeline.wiki_entity_linking import run_el, training_set_creator, kb_creator

from spacy._ml import SpacyVectors, create_default_optimizer, zero_init, cosine

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

# from spacy.cli.pretrain import get_cossim_loss
from spacy.matcher import PhraseMatcher


class EL_Model:

    PRINT_INSPECT = False
    PRINT_BATCH_LOSS = False
    EPS = 0.0000000005

    BATCH_SIZE = 100

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

    HIDDEN_1_WIDTH = 32
    DESC_WIDTH = 64
    ARTICLE_WIDTH = 128
    SENT_WIDTH = 64

    DROP = 0.4
    LEARN_RATE = 0.005
    EPOCHS = 10
    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(embed_width=self.INPUT_DIM,
                        desc_width=self.DESC_WIDTH,
                        article_width=self.ARTICLE_WIDTH,
                        sent_width=self.SENT_WIDTH,
                        hidden_1_width=self.HIDDEN_1_WIDTH)

    def train_model(self, training_dir, entity_descr_output, trainlimit=None, devlimit=None, to_print=True):
        np.seterr(divide="raise", over="warn", under="ignore", invalid="raise")

        id_to_descr = kb_creator._get_id_to_description(entity_descr_output)

        train_ent, train_gold, train_desc, train_art, train_art_texts, train_sent, train_sent_texts = \
            self._get_training_data(training_dir, id_to_descr, False, trainlimit, to_print=False)
        train_clusters = list(train_ent.keys())

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

        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 cluster, entities in train_ent.items():
                print()
                for entity in entities:
                    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)

        self._begin_training()

        if to_print:
            print()
            print("Training on", len(train_clusters), "entity clusters in", len(train_art_texts), "articles")
            print("Training instances pos/neg:", train_pos_count, train_neg_count)
            print()
            print("Dev test on", len(dev_clusters), "entity clusters in", len(dev_art_texts), "articles")
            print("Dev instances pos/neg:", dev_pos_count, dev_neg_count)
            print()
            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(" DROP", self.DROP)
            print(" LEARNING RATE", self.LEARN_RATE)
            print(" BATCH SIZE", self.BATCH_SIZE)
            print()

        dev_random = self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts,
                                    calc_random=True)
        print("acc", "dev_random", round(dev_random, 2))

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

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

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

            while start < len(train_clusters):
                next_batch = {c: train_ent[c] for c in train_clusters[start:stop]}
                processed += len(next_batch.keys())

                self.update(entity_clusters=next_batch, golds=train_gold, descs=train_desc,
                            art_texts=train_art_texts, arts=train_art,
                            sent_texts=train_sent_texts, sents=train_sent)

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

            train_acc = self._test_dev(train_ent, train_gold, train_desc, train_art, train_art_texts, train_sent, train_sent_texts, avg=True)
            dev_acc = self._test_dev(dev_ent, dev_gold, dev_desc, dev_art, dev_art_texts, dev_sent, dev_sent_texts, avg=True)

            print(i, "acc train/dev", round(train_acc, 2), round(dev_acc, 2))

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

    def _test_dev(self, entity_clusters, golds, descs, arts, art_texts, sents, sent_texts, avg=True, calc_random=False):
        correct = 0
        incorrect = 0

        if calc_random:
            for cluster, entities in entity_clusters.items():
                correct_entities = [e for e in entities if golds[e]]
                assert len(correct_entities) == 1

                entities = list(entities)
                shuffle(entities)

                if calc_random:
                    predicted_entity = random.choice(entities)
                    if predicted_entity in correct_entities:
                        correct += 1
                    else:
                        incorrect += 1

        else:
            all_clusters = list()
            arts_list = list()
            sents_list = list()

            for cluster in entity_clusters.keys():
                all_clusters.append(cluster)
                arts_list.append(art_texts[arts[cluster]])
                sents_list.append(sent_texts[sents[cluster]])

            art_docs = list(self.nlp.pipe(arts_list))
            sent_docs = list(self.nlp.pipe(sents_list))

            for i, cluster in enumerate(all_clusters):
                entities = entity_clusters[cluster]
                correct_entities = [e for e in entities if golds[e]]
                assert len(correct_entities) == 1

                entities = list(entities)
                shuffle(entities)

                desc_docs = self.nlp.pipe([descs[e] for e in entities])
                sent_doc = sent_docs[i]
                article_doc = art_docs[i]

                predicted_index = self._predict(article_doc=article_doc, sent_doc=sent_doc,
                                                desc_docs=desc_docs, avg=avg)
                if entities[predicted_index] in correct_entities:
                    correct += 1
                else:
                    incorrect += 1

        if correct == incorrect == 0:
            return 0

        acc = correct / (correct + incorrect)
        return acc

    def _predict(self, article_doc, sent_doc, 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)\
                 and self.sent_encoder.use_params(self.sgd_sent.averages):
                desc_encodings = self.desc_encoder(desc_docs)
                doc_encoding = self.article_encoder([article_doc])
                sent_encoding = self.sent_encoder([sent_doc])

        else:
            desc_encodings = self.desc_encoder(desc_docs)
            doc_encoding = self.article_encoder([article_doc])
            sent_encoding = self.sent_encoder([sent_doc])

        concat_encoding = [list(doc_encoding[0]) + list(sent_encoding[0])]

        if avg:
            with self.cont_encoder.use_params(self.sgd_cont.averages):
                cont_encodings = self.cont_encoder(np.asarray([concat_encoding[0]]))

        else:
            cont_encodings = self.cont_encoder(np.asarray([concat_encoding[0]]))

        context_enc = np.transpose(cont_encodings)

        highest_sim = -5
        best_i = -1
        for i, desc_enc in enumerate(desc_encodings):
            sim = cosine(desc_enc, context_enc)
            if sim >= highest_sim:
                best_i = i
                highest_sim = sim

        return best_i

    def _build_cnn(self, embed_width, desc_width, article_width, sent_width, hidden_1_width):
        self.desc_encoder = self._encoder(in_width=embed_width, hidden_with=hidden_1_width, end_width=desc_width)
        self.cont_encoder = self._context_encoder(embed_width=embed_width, article_width=article_width,
                                                     sent_width=sent_width, hidden_width=hidden_1_width,
                                                     end_width=desc_width)


    # 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)

    def _context_encoder(self, embed_width, article_width, sent_width, hidden_width, end_width):
        self.article_encoder = self._encoder(in_width=embed_width, hidden_with=hidden_width, end_width=article_width)
        self.sent_encoder = self._encoder(in_width=embed_width, hidden_with=hidden_width, end_width=sent_width)

        model = Affine(end_width, article_width+sent_width, drop_factor=0.0)
        return model

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

        with Model.define_operators({">>": chain, "**": clone}):
            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_cont = create_default_optimizer(self.cont_encoder.ops)
        self.sgd_cont.learn_rate = self.LEARN_RATE
        self.sgd_cont.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

    def get_loss(self, pred, gold, targets):
        loss, gradients = self.get_cossim_loss(pred, gold, targets)
        return loss, gradients

    def get_cossim_loss(self, yh, y, t):
        # Add a small constant to avoid 0 vectors
        # print()
        # print("yh", yh)
        # print("y", y)
        # print("t", t)
        yh = yh + 1e-8
        y = y + 1e-8
        # https://math.stackexchange.com/questions/1923613/partial-derivative-of-cosine-similarity
        xp = get_array_module(yh)
        norm_yh = xp.linalg.norm(yh, axis=1, keepdims=True)
        norm_y = xp.linalg.norm(y, axis=1, keepdims=True)
        mul_norms = norm_yh * norm_y
        cos = (yh * y).sum(axis=1, keepdims=True) / mul_norms
        # print("cos", cos)
        d_yh = (y / mul_norms) - (cos * (yh / norm_yh ** 2))
        # print("abs", xp.abs(cos - t))
        loss = xp.abs(cos - t).sum()
        # print("loss", loss)
        # print("d_yh", d_yh)
        inverse = np.asarray([int(t[i][0]) * d_yh[i] for i in range(len(t))])
        # print("inverse", inverse)
        return loss, -inverse

    def update(self, entity_clusters, golds, descs, art_texts, arts, sent_texts, sents):
        arts_list = list()
        sents_list = list()
        descs_list = list()
        targets = list()

        for cluster, entities in entity_clusters.items():
            art = art_texts[arts[cluster]]
            sent = sent_texts[sents[cluster]]
            for e in entities:
                if golds[e]:
                    arts_list.append(art)
                    sents_list.append(sent)
                    descs_list.append(descs[e])
                    targets.append([1])
                # else:
                #    arts_list.append(art)
                #    sents_list.append(sent)
                #    descs_list.append(descs[e])
                #    targets.append([-1])

        desc_docs = self.nlp.pipe(descs_list)
        desc_encodings, bp_desc = self.desc_encoder.begin_update(desc_docs, drop=self.DROP)

        art_docs = self.nlp.pipe(arts_list)
        sent_docs = self.nlp.pipe(sents_list)

        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)

        concat_encodings = [list(doc_encodings[i]) + list(sent_encodings[i]) for i in
                            range(len(targets))]
        cont_encodings, bp_cont = self.cont_encoder.begin_update(np.asarray(concat_encodings), drop=self.DROP)

        loss, cont_gradient = self.get_loss(cont_encodings, desc_encodings, targets)

        # loss, desc_gradient = self.get_loss(desc_encodings, cont_encodings, targets)
        # cont_gradient = cont_gradient / 2
        # desc_gradient = desc_gradient / 2
        # bp_desc(desc_gradient, sgd=self.sgd_desc)

        if self.PRINT_BATCH_LOSS:
            print("batch loss", loss)

        context_gradient = bp_cont(cont_gradient, sgd=self.sgd_cont)

        # gradient : concat (doc+sent) vs. desc
        sent_start = self.ARTICLE_WIDTH
        sent_gradients = list()
        doc_gradients = list()
        for x in context_gradient:
            doc_gradients.append(list(x[0:sent_start]))
            sent_gradients.append(list(x[sent_start:]))

        bp_doc(doc_gradients, sgd=self.sgd_article)
        bp_sent(sent_gradients, sgd=self.sgd_sent)

    def _get_training_data(self, training_dir, id_to_descr, dev, limit, to_print):
        correct_entries, incorrect_entries = training_set_creator.read_training_entities(training_output=training_dir,
                                                                                         collect_correct=True,
                                                                                         collect_incorrect=True)

        entities_by_cluster = dict()
        gold_by_entity = dict()
        desc_by_entity = dict()
        article_by_cluster = dict()
        text_by_article = dict()
        sentence_by_cluster = dict()
        text_by_sentence = dict()
        sentence_by_text = 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")

                    try:
                        # 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
                        for mention, entity_pos in correct_entries[article_id].items():
                            cluster = article_id + "_" + mention
                            descr = id_to_descr.get(entity_pos)
                            entities = set()
                            if descr:
                                entity = "E_" + str(next_entity_nr) + "_" + cluster
                                next_entity_nr += 1
                                gold_by_entity[entity] = 1
                                desc_by_entity[entity] = descr
                                entities.add(entity)

                                entity_negs = incorrect_entries[article_id][mention]
                                for entity_neg in entity_negs:
                                    descr = id_to_descr.get(entity_neg)
                                    if descr:
                                        entity = "E_" + str(next_entity_nr) + "_" + cluster
                                        next_entity_nr += 1
                                        gold_by_entity[entity] = 0
                                        desc_by_entity[entity] = descr
                                        entities.add(entity)

                            found_matches = 0
                            if len(entities) > 1:
                                entities_by_cluster[cluster] = 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([mention]))

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

                                # store sentences
                                for match_id, start, end in matches:
                                    span = article_doc[start:end]
                                    if mention == span.text:
                                        found_matches += 1
                                        sent_text = span.sent.text
                                        sent_nr = sentence_by_text.get(sent_text,  None)
                                        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_by_text[sent_text] = sent_nr
                                        article_by_cluster[cluster] = article_id
                                        sentence_by_cluster[cluster] = sent_nr

                            if found_matches == 0:
                                # print("Could not find neg instances or sentence matches for", mention, "in", article_id)
                                entities_by_cluster.pop(cluster, None)
                                article_by_cluster.pop(cluster, None)
                                sentence_by_cluster.pop(cluster, None)
                                for entity in entities:
                                    gold_by_entity.pop(entity, None)
                                    desc_by_entity.pop(entity, None)
                        cnt += 1
                    except:
                        print("Problem parsing article", article_id)

        if to_print:
            print()
            print("Processed", cnt, "training articles, dev=" + str(dev))
            print()
        return entities_by_cluster, gold_by_entity, desc_by_entity, article_by_cluster, text_by_article, \
               sentence_by_cluster, text_by_sentence