BOW embedding layers for Keras

bow-embedding-keras.py

class BOW(Embedding):
    '''
    Bag-of-words embedding layer.
    
    The default embedding layer outputs the concatenation of [W_t * one_hot(input)];
    this layer simply sums over the resulting output.
    '''    
    def get_output_shape_for(self, input_shape):
        return (input_shape[0], self.output_dim)

    def call(self, x, mask=None):
        return K.sum(Embedding.call(self, x, mask), axis=1)
    
# The embedding layer uses a mask which isn't supported by the rest of the network.
# Our bag-of-words approach merges the masked and unmasked values however, so we
# just need a layer that removes the mask from the rest of the network.
class StripMask(Layer):
    def __init__(self, **kwargs):
        self.supports_mask = True
        Layer.__init__(self, **kwargs)
    
    def compute_mask(self, input_shape, input_mask):
        return None

      
class CosineDistance(Layer):
    def get_output_shape_for(self, input_shape):
        return (input_shape[0], 1)
    
    def call(self, input, mask=None):
        x = K.expand_dims(input[:, :HIDDEN_SIZE])
        y = K.expand_dims(input[:, HIDDEN_SIZE:])
        denominator = K.sqrt(K.batch_dot(x, x, axes=HIDDEN_SIZE) * 
                             K.batch_dot(y, y, axes=HIDDEN_SIZE))
        dot = K.batch_dot(x, y, axes=HIDDEN_SIZE)
        return dot / denominator