embedding.py

from keras import backend as K
from keras.layers import Embedding

class OntoAwareEmbedding(Embedding):
    '''
    This class modifies two aspects of the Embedding class in Keras:
    1. Higher order inputs: Embedding already works with inputs of any shape, except that the output shape
    returned by it assumes that the input it 2D. Changing it.
    2. Sense priors: The expected input shape is (num_samples, num_words, num_senses, num_hyps+1). The +1 at the
    end is the word index appended at the end of each sense indices vector. This is to define an additional real
    value for each word, which will act as the sense prior in OntoLSTM. OntoLSTM is reponsible for handling this
    correctly. The output shape is (num_samples, num_words, num_senses, num_hyps, embedding_dim + 1).
    '''
    input_ndim = 4

    def __init__(self, word_index_size, synset_index_size, embedding_dim, set_sense_priors=True,
                 tune_embedding=True, **kwargs):
        self.embedding_dim = embedding_dim
        self.word_index_size = word_index_size
        self.synset_index_size = synset_index_size
        self.set_sense_priors = set_sense_priors
        # We have a separate "tune_embedding" field instead of using trainable because we have two sets of
        # parameters here: the embedding weights, and sense prior weights. We may want to fix only one of
        # them at a time.
        self.tune_embedding = tune_embedding
        # Convincing Embedding to return an embedding of the right shape. The output_dim of this layer is embedding_dim+1
        kwargs['output_dim'] = self.embedding_dim
        kwargs['input_dim'] = self.synset_index_size
        self.onto_aware_embedding_weights = None
        super(OntoAwareEmbedding, self).__init__(**kwargs)

    @staticmethod
    def _get_initial_sense_priors(shape, rate_range=None, name=None):
        # This returns a Keras variable with the initial values all being 0.5.
        if rate_range is None:
            low, high = 0.01, 0.99
        else:
            low, high = rate_range
        return K.random_uniform_variable(shape, low, high, name=name)

    def build(self, input_shape):
        # input shape is (batch_size, num_words, num_senses, num_hyps)
        self.num_senses = input_shape[-2]
        self.num_hyps = input_shape[-1] - 1  # -1 because the last value is a word index
        # embedding of size 1.
        if self.set_sense_priors:
            self.sense_priors = self._get_initial_sense_priors((self.word_index_size, 1), name='{}_sense_priors'.format(self.name))
        else:
            # OntoLSTM makes sense proabilities uniform if the passed sense parameters are zero.
            self.sense_priors = K.zeros((self.word_index_size, 1))  # uniform sense probs
        # Keeping aside the initial weights to not let Embedding set them. It wouldn't know what sense priors are.
        if self.initial_weights is not None:
            self.onto_aware_embedding_weights = self.initial_weights
            self.initial_weights = None
        # The following method will set self.trainable_weights
        super(OntoAwareEmbedding, self).build(input_shape)  # input_shape will not be used by Embedding's build.
        if not self.tune_embedding:
            # Move embedding to non_trainable_weights
            self._non_trainable_weights.append(self._trainable_weights.pop())

        if self.set_sense_priors:
            self._trainable_weights.append(self.sense_priors)

        if self.onto_aware_embedding_weights is not None:
            self.set_weights(self.onto_aware_embedding_weights)
 
    def call(self, x, mask=None):
        # Remove the word indices at the end before making a call to Embedding.
        x_synsets = x[:, :, :, :-1]  # (num_samples, num_words, num_senses, num_hyps)
        # Taking the last index from the first sense. The last index in all the senses will be the same.
        x_word_index = x[:, :, 0, -1]  # (num_samples, num_words)
        # (num_samples, num_words, num_senses, num_hyps, embedding_dim)
        synset_embedding = super(OntoAwareEmbedding, self).call(x_synsets, mask=None)
        # (num_samples, num_words, 1, 1)
        sense_prior_embedding = K.expand_dims(K.gather(self.sense_priors, x_word_index))
        # Now tile sense_prior_embedding and concatenate it with synset_embedding.
        # (num_samples, num_words, num_senses, num_hyps, 1)
        tiled_sense_prior_embedding = K.expand_dims(K.tile(sense_prior_embedding, (1, 1, self.num_senses, self.num_hyps)))
        synset_embedding_with_priors = K.concatenate([synset_embedding, tiled_sense_prior_embedding])
        return synset_embedding_with_priors

    def compute_mask(self, x, mask=None):
        # Since the output dim is different, we need to change the mask size
        embedding_mask = super(OntoAwareEmbedding, self).compute_mask(x, mask)
        return embedding_mask[:, :, :, :-1] if embedding_mask is not None else None

    def get_output_shape_for(self, input_shape):
        return input_shape[:3] + (self.num_hyps, self.embedding_dim+1)

    def get_config(self):
        config = {"word_index_size": self.word_index_size,
                  "synset_index_size": self.synset_index_size,
                  "embedding_dim": self.embedding_dim,
                  "set_sense_priors": self.set_sense_priors
                 }
        base_config = super(OntoAwareEmbedding, self).get_config()
        config.update(base_config)
        return config