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