set context = encoded when memory is None
when memory is None, add the else clause and set: `context = encoded`
and add `if last_attention is not None:`
It can be used to build bi-directional self-attention decoder (which memory is None) for the Masked LM Task.
@@ -204,6 +204,8 @@ class SelfAttentionDecoder(decoder.Decoder):
context,
mode,
dropout=self.dropout)
+ else:
+ context = encoded
with tf.variable_scope("ffn"):
transformed = transformer.feed_forward(
@@ -219,10 +221,13 @@ class SelfAttentionDecoder(decoder.Decoder):
inputs = transformed
+ if last_attention is not None:
# The first head of the last layer is returned.
first_head_attention = last_attention[:, 0]
if cache is not None and "attn" in cache:
cache["attn"] = tf.concat([cache["attn"], first_head_attention], 1)
+ else:
+ first_head_attention = None
outputs = transformer.norm(inputs)
return outputs, first_head_attention
