from keras.models import Sequential
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.embeddings import Embedding
from keras.layers.normalization import BatchNormalization
from keras.layers import Input, Convolution1D, GlobalMaxPooling1D, GlobalAveragePooling1D, \
Lambda, TimeDistributed, SpatialDropout1D, Reshape, RepeatVector
from keras.layers.merge import Dot, Concatenate, Multiply, Add
from keras.models import Model
from keras import backend as K
from keras.initializers import RandomUniform
DSSM_NUM_NEGS = 5
ATTENTION_DEEP_LEVEL = 5
def keras_diagonal(x):
seq_len = 10 if 'seq_len' not in K.params else K.params['seq_len']
if seq_len != 10: print("dim of diagonal matrix in keras_diagonal: %d" % seq_len)
return K.sum(K.eye(seq_len) * x, axis=1)
def elementwise_prod(x, y):
(_, vec_size,) = K.int_shape(x)
attention_level = K.params['attention_level']
return x*y[:, attention_level, 0:vec_size]
def attention_weighting_prod(attention, weights):
attention_level = K.params['attention_level']
return attention * weights[:, attention_level, :, :]
def repeat_vector(x, rep, axis):
return K.repeat(x, rep, axis)
def max_pooling(x):
return K.max(x, axis=1)
def mean_pooling(x):
return K.mean(x, axis=1)
def max_pooling_with_mask(x, query_mask):
# x is batch_size * |doc| * |query|
# query_mask is batch_size * |query| (with masks as 0)
return K.max(x, axis=1) * query_mask
def mean_pooling_with_mask(x, doc_mask, query_mask):
# x is batch_size * |doc| * |query|
# doc_mask is batch_size * |doc| (with masks as 0)
# query_mask is batch_size * |query| (with masks as 0)
ZERO_SHIFT = 0.1
doc_mask_sum = (K.sum(doc_mask, axis=-1, keepdims=True) + ZERO_SHIFT)
return query_mask * K.batch_dot(x, doc_mask, axes=[1, 1]) / doc_mask_sum
def add_embed_layer(vocab_emb, vocab_size, embed_size, train_embed, dropout_rate):
emb_layer = Sequential()
if vocab_emb is not None:
print("Embedding with initialized weights")
print(vocab_size, embed_size)
emb_layer.add(Embedding(input_dim=vocab_size, output_dim=embed_size, weights=[vocab_emb],
trainable=train_embed, mask_zero=False))
else:
print("Embedding with random weights")
emb_layer.add(Embedding(input_dim=vocab_size, output_dim=embed_size, trainable=True, mask_zero=False,
embeddings_initializer=RandomUniform(-0.05, 0.05)))
emb_layer.add(SpatialDropout1D(dropout_rate))
return emb_layer
def add_conv_layer(input_list, layer_name, nb_filters, kernel_size, padding, dropout_rate=0.1,
activation='relu', strides=1, attention_level=0, conv_option="normal", prev_conv_tensors=None):
conv_layer = Convolution1D(filters=nb_filters, kernel_size=kernel_size, padding=padding,
activation=activation, strides=strides, name=layer_name)
max_pooling_layer = GlobalMaxPooling1D()
dropout_layer = Dropout(dropout_rate)
output_list, conv_output_list = [], []
for i in range(len(input_list)):
input = input_list[i]
conv_tensor = conv_layer(input)
if conv_option == "ResNet":
conv_tensor = Add()([conv_tensor, prev_conv_tensors[i][-1]])
dropout_tensor = dropout_layer(conv_tensor)
#conv_pooling_tensor = max_pooling_layer(conv_tensor)
output_list.append(dropout_tensor)
#conv_output_list.append(conv_pooling_tensor)
conv_output_list.append(conv_tensor)
return output_list, conv_output_list
def add_attention_layer(query_embedding, doc_embedding, layer_name, query_mask=None, doc_mask=None, mask=False):
dot_prod = Dot(axes=-1, name=layer_name)([doc_embedding, query_embedding])
norm_sim = Activation('softmax')(dot_prod)
if mask:
max_sim = Lambda(lambda x: max_pooling_with_mask(x[0], x[1]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][2],))([norm_sim, query_mask])
mean_sim = Lambda(lambda x: mean_pooling_with_mask(x[0], x[1], x[2]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][2],))([norm_sim, doc_mask, query_mask])
else:
max_sim = Lambda(max_pooling, output_shape=lambda inp_shp: (inp_shp[0], inp_shp[2], ))(norm_sim)
mean_sim = Lambda(mean_pooling, output_shape=lambda inp_shp: (inp_shp[0], inp_shp[2],))(norm_sim)
return norm_sim, max_sim, mean_sim
def add_attention_layer_with_query_weighting(query_embedding, doc_embedding, layer_name, attention_level,
query_weight, query_mask=None, doc_mask=None, mask=False):
"""
Dot -> softmax -> pooling -> (mask) -> weighting
"""
dot_prod = Dot(axes=-1, name=layer_name)([doc_embedding, query_embedding])
norm_sim = Activation('softmax')(dot_prod)
if mask:
max_sim = Lambda(lambda x: max_pooling_with_mask(x[0], x[1]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][2],))([norm_sim, query_mask])
mean_sim = Lambda(lambda x: mean_pooling_with_mask(x[0], x[1], x[2]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][2],))([norm_sim, doc_mask, query_mask])
else:
max_sim = Lambda(max_pooling, output_shape=lambda inp_shp: (inp_shp[0], inp_shp[2], ))(norm_sim)
mean_sim = Lambda(mean_pooling, output_shape=lambda inp_shp: (inp_shp[0], inp_shp[2],))(norm_sim)
if attention_level <= 1:
setattr(K, 'params', {'attention_level': attention_level})
max_sim = Lambda(lambda x: elementwise_prod(x[0], x[1]),
output_shape=lambda inp_shp: (inp_shp[0][0], inp_shp[0][1],))([max_sim, query_weight])
mean_sim = Lambda(lambda x: elementwise_prod(x[0], x[1]),
output_shape=lambda inp_shp: (inp_shp[0][0], inp_shp[0][1]))([mean_sim, query_weight])
return norm_sim, max_sim, mean_sim
# doc weighting is not tested yet, please use query weighting instead.
def add_attention_layer_with_doc_weighting(query_embedding, doc_embedding, layer_name, attention_level,
query_weight, doc_weight, max_query_len, max_doc_len,
query_mask=None, doc_mask=None, mask=False):
dot_prod = Dot(axes=-1, name=layer_name)([doc_embedding, query_embedding])
norm_sim = Activation('softmax')(dot_prod)
reshaped_query_weight = Reshape((ATTENTION_DEEP_LEVEL, max_query_len, 1),
input_shape=(ATTENTION_DEEP_LEVEL, max_query_len,))(query_weight)
repeated_query_weight = Lambda(lambda x: repeat_vector(x[0], x[1], x[2]), output_shape=lambda inp_shp:(
inp_shp[0][0], inp_shp[0][1], max_query_len, max_doc_len,))([reshaped_query_weight, max_doc_len, -1])
reshaped_doc_weight = Reshape((ATTENTION_DEEP_LEVEL, 1, max_doc_len),
input_shape=(ATTENTION_DEEP_LEVEL, max_doc_len,))(doc_weight)
repeated_doc_weight = Lambda(lambda x: repeat_vector(x[0], x[1], x[2]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][1], max_query_len, max_doc_len,))([reshaped_doc_weight, max_query_len, -2])
weight_product = Multiply()([repeated_query_weight, repeated_doc_weight])
transformed_weight_product = Dense(max_doc_len, activation='relu')(weight_product)
setattr(K, 'params', {'attention_level': attention_level})
norm_sim = Lambda(lambda x: attention_weighting_prod(x[0], x[1]), output_shape=lambda inp_shp:(
inp_shp[0][0], inp_shp[0][1], inp_shp[0][2]))([norm_sim, transformed_weight_product])
if mask:
max_sim = Lambda(lambda x: max_pooling_with_mask(x[0], x[1]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][2],))([norm_sim, query_mask])
mean_sim = Lambda(lambda x: mean_pooling_with_mask(x[0], x[1], x[2]), output_shape=lambda inp_shp: (
inp_shp[0][0], inp_shp[0][2],))([norm_sim, doc_mask, query_mask])
else:
max_sim = Lambda(max_pooling, output_shape=lambda inp_shp: (inp_shp[0], inp_shp[2], ))(norm_sim)
mean_sim = Lambda(mean_pooling, output_shape=lambda inp_shp: (inp_shp[0], inp_shp[2],))(norm_sim)
return norm_sim, max_sim, mean_sim
########################## Our model implementation #########################
def create_attention_model(max_query_len, max_doc_len, max_url_len, vocab_size, embedding_matrix, nb_filters,
embed_size=300, dropout_rate=0.1, trainable=True, weighting=False, mask=False,
conv_option="normal", model_option="complete"):
print('create attention model...')
model = Sequential()
query_word_input = Input(shape=(max_query_len['word'], ), name="query_word_input")
doc_word_input = Input(shape=(max_doc_len['word'],), name="doc_word_input")
query_char_input = Input(shape=(max_query_len['3gram'],), name="query_3gram_input")
doc_char_input = Input(shape=(max_doc_len['3gram'],), name="doc_3gram_input")
url_char_input = Input(shape=(max_url_len['url'],), name="url_3gram_input")
input_list = [query_word_input, doc_word_input, query_char_input, doc_char_input, url_char_input]
# Define Mask
query_word_mask, query_char_mask, doc_word_mask, doc_char_mask, url_char_mask = None, None, None, None, None
if mask:
query_word_mask = Input(shape=(max_query_len['word'], ), name="query_word_mask")
doc_word_mask = Input(shape=(max_doc_len['word'], ), name="doc_word_mask")
query_char_mask = Input(shape=(max_query_len['3gram'],), name="query_3gram_mask")
doc_char_mask = Input(shape=(max_doc_len['3gram'],), name="doc_3gram_mask")
url_char_mask = Input(shape=(max_url_len['url'],), name="url_3gram_mask")
input_list.extend([query_word_mask, doc_word_mask, query_char_mask, url_char_mask])
query_word_weight, query_char_weight, doc_word_weight, doc_char_weight, url_char_weight = None, None, None, None, None
if weighting == 'query':
query_word_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_query_len['word'], ), name="query_word_weight")
query_char_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_query_len['3gram'],), name="query_3gram_weight")
#external_feat = Input(shape=(2, ), name='overlap_feat')
input_list.extend([query_word_weight, query_char_weight])
#input_list.append(external_feat)
elif weighting == 'doc':
query_word_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_query_len['word'],), name="query_word_weight")
query_char_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_query_len['3gram'],), name="query_3gram_weight")
doc_word_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_doc_len['word'],), name="doc_word_weight")
doc_char_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_doc_len['3gram'],), name="doc_3gram_weight")
url_char_weight = Input(shape=(ATTENTION_DEEP_LEVEL, max_url_len['url'],), name="url_3gram_weight")
input_list.extend([query_word_weight, query_char_weight, doc_word_weight, doc_char_weight, url_char_weight])
char_weight_candidates = [doc_char_weight, url_char_weight]
# Create query-doc word-to-word attention layer
query_word_embedding_layer = add_embed_layer(embedding_matrix, vocab_size['word'], embed_size, trainable,
dropout_rate)
query_embedding = query_word_embedding_layer(query_word_input)
doc_embedding = query_word_embedding_layer(doc_word_input)
norm_sim_list, max_sim_list, mean_sim_list = [], [], []
conv_embedding_list = [[query_embedding], [doc_embedding]]
for i in range(ATTENTION_DEEP_LEVEL):
if i > 0:
output_list, conv_output_list = add_conv_layer([query_embedding, doc_embedding], "word-conv%d" % i,
nb_filters, 2, "same", dropout_rate, strides=1,
attention_level=i, conv_option=conv_option,
prev_conv_tensors=conv_embedding_list)
query_embedding, doc_embedding = output_list[0], output_list[1]
conv_embedding_list[0].append(conv_output_list[0])
conv_embedding_list[1].append(conv_output_list[1])
if weighting == 'query':
norm_sim, max_sim, mean_sim = add_attention_layer_with_query_weighting(
query_embedding, doc_embedding, "word-attention%d" % i, i, query_word_weight,
query_word_mask, doc_word_mask, mask)
elif weighting == 'doc':
norm_sim, max_sim, mean_sim = add_attention_layer_with_doc_weighting(
query_embedding, doc_embedding, "word-attention%d" % i, i, query_word_weight, doc_word_weight,
max_query_len['word'], max_doc_len['word'], query_word_mask, doc_word_mask, mask)
norm_sim_list.append(norm_sim)
max_sim_list.append(max_sim)
mean_sim_list.append(mean_sim)
# Create query-doc char-to-char attention layer
char_embedding_layer = add_embed_layer(None, vocab_size['char'], embed_size, True, dropout_rate)
query_char_embedding_layer = url_char_embedding_layer = doc_char_embedding_layer = char_embedding_layer
doc_char_embedding = doc_char_embedding_layer(doc_char_input)
url_char_embedding = url_char_embedding_layer(url_char_input)
char_embedding_candidates = [doc_char_embedding, url_char_embedding]
char_mask_candidates = [doc_char_mask, url_char_mask]
max_doc_len_candidates = [max_doc_len["3gram"], max_url_len["url"]]
j = 0
for char_embedding, char_mask, char_weight, max_doc_len_tmp in zip(char_embedding_candidates, char_mask_candidates,
char_weight_candidates, max_doc_len_candidates):
query_embedding = query_char_embedding_layer(query_char_input)
conv_embedding_list = [[query_embedding], [char_embedding]]
setattr(K, 'params', {'max_doc_len': max_doc_len_tmp, 'max_query_len': max_query_len['3gram']})
for i in range(ATTENTION_DEEP_LEVEL):
if i > 0:
output_list, conv_output_list = add_conv_layer([query_embedding, char_embedding], "3gram-conv%d" % j,
nb_filters, 4, "same", dropout_rate, strides=1,
attention_level=i, conv_option=conv_option,
prev_conv_tensors=conv_embedding_list)
query_embedding, char_embedding = output_list[0], output_list[1]
conv_embedding_list[0].append(conv_output_list[0])
conv_embedding_list[1].append(conv_output_list[1])
if weighting == 'query':
norm_sim2, max_sim2, mean_sim2 = add_attention_layer_with_query_weighting(
query_embedding, char_embedding, "url-attention%d" % j, i, query_char_weight,
query_char_mask, char_mask, mask)
elif weighting == 'doc':
norm_sim2, max_sim2, mean_sim2 = add_attention_layer_with_doc_weighting(
query_embedding, char_embedding, "url-attention%d" % j, i, query_char_weight, char_weight,
max_query_len['3gram'], max_url_len['url'], query_char_mask, char_mask, mask)
else:
norm_sim2, max_sim2, mean_sim2 = add_attention_layer(query_embedding, char_embedding, "url-attention%d" % j,
query_char_mask, char_mask, mask)
norm_sim_list.append(norm_sim2)
max_sim_list.append(max_sim2)
mean_sim_list.append(mean_sim2)
j += 1
max_sim_list.extend(mean_sim_list)
feature_vector = Concatenate(axis=-1, name="feature_vector")(max_sim_list)
#if weighting:
# feature_vector = Concatenate(axis=-1, name="external_feature_vector")([feature_vector, external_feat])
feature_vector1 = Dense(150, activation='relu', name="feature_vector1")(feature_vector)
feature_vector2 = Dense(50, activation='relu', name="feature_vector2")(feature_vector1)
prediction = Dense(1, activation='sigmoid', name="prediction")(feature_vector2)
return Model(input_list, [prediction])
