def get_cell(cell_type, size, layers=1, direction='unidirectional'):
if cell_type == "layer_norm_basic":
cell = LayerNormBasicLSTMCell(size)
elif cell_type == "lstm_block_fused":
cell = tf.contrib.rnn.LSTMBlockFusedCell(size)
elif cell_type == "cudnn_lstm":
cell = CudnnLSTM(layers, size, direction=direction)
elif cell_type == "cudnn_gru":
cell = CudnnGRU(layers, size, direction=direction)
elif cell_type == "lstm_block":
cell = LSTMBlockCell(size)
elif cell_type == "gru_block":
cell = GRUBlockCell(size)
elif cell_type == "rnn":
cell = BasicRNNCell(size)
elif cell_type == "cudnn_rnn":
cell = CudnnRNNTanh(layers, size)
else:
cell = BasicLSTMCell(size)
return cell
def __call__(self, is_train, scope=None):
return GRUBlockCell(self.num_units)
def _init(self):
ExtractionQAModel._init(self)
if self._composition == "GRU":
if self._layer_norm:
rnn_constructor = lambda size: FusedRNNCellAdaptor(LayerNormGRUCell(size), use_dynamic_rnn=True)
else:
rnn_constructor = lambda size: FusedRNNCellAdaptor(GRUBlockCell(size), use_dynamic_rnn=True)
elif self._composition == "RNN":
rnn_constructor = lambda size: FusedRNNCellAdaptor(BasicRNNCell(size), use_dynamic_rnn=True)
else:
if self._layer_norm:
rnn_constructor = lambda size: FusedRNNCellAdaptor(LayerNormLSTMCell(size), use_dynamic_rnn=True)
else:
rnn_constructor = lambda size: LSTMBlockFusedCell(size)
with tf.device(self._device0):
self._eval = tf.get_variable("is_eval", initializer=False, trainable=False)
self._set_train = self._eval.initializer
self._set_eval = self._eval.assign(True)
self.context_mask = tfutil.mask_for_lengths(self.context_length, self._batch_size, self.embedder.max_length)
question_binary_mask = tfutil.mask_for_lengths(self.question_length,
self.question_embedder.batch_size,
self.question_embedder.max_length,
value=1.0,
mask_right=False)
with tf.variable_scope("preprocessing_layer"):
question_binary_mask = tf.gather(question_binary_mask, self.context_partition)
self._embedded_question_not_dropped = tf.gather(self._embedded_question_not_dropped, self.context_partition)
# context
if self._with_features:
mask = tf.get_variable("attention_mask", [1, 1, self._embedded_question_not_dropped.get_shape()[-1].value],
initializer=tf.constant_initializer(1.0))
# compute word wise features
#masked_question = self.question_embedder.output * mask
# [B, Q, L]
q2c_scores = tf.matmul(self._embedded_question_not_dropped * mask,
self._embedded_context_not_dropped, adjoint_b=True)
q2c_scores = q2c_scores + tf.expand_dims(self.context_mask, 1)
#c2q_weights = tf.reduce_max(q2c_scores / (tf.reduce_max(q2c_scores, [2], keep_dims=True) + 1e-5), [1])
q2c_weights = tf.reduce_sum(tf.nn.softmax(q2c_scores) * \
tf.expand_dims(question_binary_mask, 2), [1])
# [B, L , 1]
self.context_features = tf.concat(axis=2, values=[tf.expand_dims(self._word_in_question, 2),
#tf.expand_dims(c2q_weights, 2),
tf.expand_dims(q2c_weights, 2)])
embedded_ctxt = tf.concat(axis=2, values=[self.embedded_context, self.context_features])
in_question_feature = tf.ones(tf.stack([self.question_embedder.batch_size,
self.question_embedder.max_length, 2]))
embedded_question = tf.concat(axis=2, values=[self.embedded_question, in_question_feature])
else:
embedded_ctxt = self.embedded_context
embedded_question = self.embedded_question
if self._with_question_type_features:
# Need to add another zero vector so that the total number
# of features is even, for LSTM performance reasons.
question_type_features = tf.stack([self._is_factoid,
self._is_list,
self._is_yesno,
tf.zeros(tf.shape(self._is_list),
dtype=tf.bool)],
axis=1)
question_type_features = tf.cast(question_type_features, tf.float32)
question_type_features = tf.expand_dims(question_type_features, 1)
embedded_question = tf.concat(axis=2, values=[embedded_question,
tf.tile(question_type_features,
tf.stack([1, tf.shape(embedded_question)[1], 1]))])
question_type_features = tf.gather(question_type_features, self.context_partition)
embedded_ctxt = tf.concat(axis=2, values=[embedded_ctxt,
tf.tile(question_type_features,
tf.stack([1, tf.shape(embedded_ctxt)[1], 1]))])
if self._with_entity_tag_features:
embedded_question = tf.concat(axis=2, values=[embedded_question,
tf.cast(self._question_tags, tf.float32)])
embedded_ctxt = tf.concat(axis=2, values=[embedded_ctxt,
tf.cast(self._context_tags, tf.float32)])
self.encoded_question = self._preprocessing_layer(rnn_constructor, embedded_question,
self.question_length, projection_scope="question_proj")
self.encoded_ctxt = self._preprocessing_layer(rnn_constructor, embedded_ctxt, self.context_length,
share_rnn=True, projection_scope="context_proj",
num_fusion_layers=self._num_intrafusion_layers)
# single time attention over question
attention_scores = tf.contrib.layers.fully_connected(self.encoded_question, 1,
activation_fn=None,
weights_initializer=None,
biases_initializer=None,
scope="attention")
attention_scores = attention_scores + tf.expand_dims(
tfutil.mask_for_lengths(self.question_length, self.question_embedder.batch_size,
self.question_embedder.max_length), 2)
attention_weights = tf.nn.softmax(attention_scores, 1)
self.question_attention_weights = attention_weights
self.question_representation = tf.reduce_sum(attention_weights * self.encoded_question, [1])
# Multiply question features for each paragraph
self.encoded_question = tf.gather(self.encoded_question, self.context_partition)
self.question_representation_per_context = tf.gather(self.question_representation, self.context_partition)
self.question_length = tf.gather(self.question_length, self.context_partition)
if self._with_inter_fusion:
with tf.variable_scope("inter_fusion"):
with tf.variable_scope("associative") as vs:
mask = tf.get_variable("attention_mask", [1, 1, self.size], initializer=tf.constant_initializer(1.0))
mask = tf.nn.relu(mask)
for i in range(1):
# [B, Q, L]
inter_scores = tf.matmul(self.encoded_question * mask, self.encoded_ctxt, adjoint_b=True)
inter_scores = inter_scores + tf.expand_dims(self.context_mask, 1)
inter_weights = tf.nn.softmax(inter_scores)
inter_weights = inter_weights * tf.expand_dims(question_binary_mask, 2)
# [B, L, Q] x [B, Q, S] -> [B, L, S]
co_states = tf.matmul(inter_weights, self.encoded_question, adj_x=True)
u = tf.contrib.layers.fully_connected(tf.concat(axis=2, values=[self.encoded_ctxt, co_states]), self.size,
activation_fn=tf.sigmoid,
biases_initializer=tf.constant_initializer(1.0),
scope="update_gate")
self.encoded_ctxt = u * self.encoded_ctxt + (1.0 - u) * co_states
vs.reuse_variables()
with tf.variable_scope("recurrent") as vs:
self.encoded_ctxt.set_shape([None, None, self.size])
self.encoded_ctxt = dynamic_rnn(GatedAggregationRNNCell(self.size),
tf.reverse_sequence(self.encoded_ctxt, self.context_length, 1),
self.context_length,
dtype=tf.float32, time_major=False, scope="backward")[0]
self.encoded_ctxt = dynamic_rnn(GatedAggregationRNNCell(self.size),
tf.reverse_sequence(self.encoded_ctxt, self.context_length, 1),
self.context_length,
dtype=tf.float32, time_major=False, scope="forward")[0]
# No matching layer, so set matched_output to encoded_ctxt (for compatibility)
self.matched_output = self.encoded_ctxt
with tf.variable_scope("pointer_layer"):
self.predicted_context_indices, \
self._start_scores, self._start_pointer, self.start_probs, \
self._end_scores, self._end_pointer, self.end_probs = \
self._spn_answer_layer(self.question_representation_per_context, self.encoded_ctxt)
self.yesno_added = False
if self._with_yesno:
self.add_yesno(add_model_scope=False)
self._train_variables = [p for p in tf.trainable_variables() if self.name in p.name]