Feat/make transformer decoder callable without causal mask #1083
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
| use_causal_mask: bool, defaults to True. If true, causal mask | ||
| (masking out future input) is applied on the decoder sequence. |
There was a problem hiding this comment.
naming and documentation from previous version where this was in: https://github.com/keras-team/keras-nlp/blob/cb0fa028971475879911ddf042a1473037775ee6/keras_nlp/layers/transformer_decoder.py#L191-L192
There was a problem hiding this comment.
this lgtm, and matches the multi-head attention layer in core Keras
| use_causal_mask: bool, defaults to True. If true, causal mask | ||
| (masking out future input) is applied on the decoder sequence. |
There was a problem hiding this comment.
this lgtm, and matches the multi-head attention layer in core Keras
| if decoder_mask is not None: | ||
| self_attention_mask = tf.minimum(decoder_mask, self_attention_mask) | ||
| if use_causal_mask: | ||
| batch_size = tf.shape(decoder_sequence)[0] |
There was a problem hiding this comment.
I think this branching is getting complex enough that we should split a private method for this
self_attention_mask = self._compute_self_attention_mask(
decoder_sequence,
self_attention_cache,
self_attention_cache_update_index,
use_causal_mask,
)
|
/gcbrun |
Fixes #1062