def transformer_text_encoder(inputs,
target_space,
hparams,
name=None):
"""Transformer text encoder over inputs with unmasked full attention.
Args:
inputs: Tensor of shape [batch, length, 1, hparams.hidden_size].
target_space: int. Used for encoding inputs under a target space id.
hparams: tf.contrib.training.HParams.
name: string, variable scope.
Returns:
encoder_output: Tensor of shape [batch, length, hparams.hidden_size].
ed: Tensor of shape [batch, 1, 1, length]. Encoder-decoder attention bias
for any padded tokens.
"""
with tf.variable_scope(name, default_name="transformer_text_encoder"):
inputs = common_layers.flatten4d3d(inputs)
[
encoder_input,
encoder_self_attention_bias,
ed,
] = transformer_layers.transformer_prepare_encoder(
inputs, target_space=target_space, hparams=hparams)
encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
encoder_output = transformer_layers.transformer_encoder(
encoder_input, encoder_self_attention_bias, hparams)
return encoder_output, ed
def transformer_text_encoder(inputs, target_space, hparams, name=None):
"""Transformer text encoder over inputs with unmasked full attention.
Args:
inputs: Tensor of shape [batch, length, 1, hparams.hidden_size].
target_space: int. Used for encoding inputs under a target space id.
hparams: HParams.
name: string, variable scope.
Returns:
encoder_output: Tensor of shape [batch, length, hparams.hidden_size].
ed: Tensor of shape [batch, 1, 1, length]. Encoder-decoder attention bias
for any padded tokens.
"""
with tf.variable_scope(name, default_name="transformer_text_encoder"):
inputs = common_layers.flatten4d3d(inputs)
[
encoder_input,
encoder_self_attention_bias,
ed,
] = transformer_layers.transformer_prepare_encoder(
inputs, target_space=target_space, hparams=hparams)
encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
encoder_output = transformer_layers.transformer_encoder(
encoder_input, encoder_self_attention_bias, hparams)
return encoder_output, ed
def model(self, x, sen_len):
x, enc_self_att_bias, _ = transformer_prepare_encoder(
x, self.target_space, self.hparams)
# x shape (batch_size, sequence_len, hidden_size)
# enc_self_att_bias shape (batch_size, 1, 1, sequence_len)
x = transformer_encoder(x, enc_self_att_bias, self.hparams)
# shape (batch_size, sequence_length, hidden_size)
indices = tf.stack(
[tf.range(0,
tf.shape(x)[0]), sen_len - 1
], # the last indices of each sentence
axis=1)
final_output_state = tf.gather_nd(x,
indices,
name="select_last_word_in_sentence")
# shape: (batch_size, hidden_size)
return final_output_state