Repository: keras-team/keras-nlp
Pull request #790 (2023-02-27 23:46:48)

Function merge_padding_and_attention_mask does not return an output with the desired shape when both padding and attention masks are given
Function merge_padding_and_attention_mask does not return an output with the desired shape when both padding and attention masks are given.

See the linked issue here: https://github.com/keras-team/keras-nlp/issues/783#issue-1600168236

Currently, the TransformerEncoder layers breaks if we call it with both padding and attention masks.

Changes: 
- Remove duplicated `tf.newaxis`
- Create a check function to control the shapes of the padding and attention masks
- Update unit tests to make sure that the error is raised if we give `merge_padding_and_attention_mask` masks with bad shapes.

google-cla[bot]: 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](https://github.com/keras-team/keras-nlp/pull/790/checks?check_run_id=11640772538) 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.
atharvapurdue: Hi @abodinier , great job!
I think there is small change required. `merged_mask` is assigned but never used. The check should pass after that

Patch:
diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
index 343d8203..14a5ab7b 100644
--- a/.github/workflows/actions.yml
+++ b/.github/workflows/actions.yml
@@ -34,6 +34,9 @@ jobs:
       - name: Test with pytest
         run: |
           pytest --cov=keras_nlp --cov-report xml:coverage.xml
+      - name: Run integration tests
+        run: |
+          python pip_build.py --install && cd integration_tests && pytest .
   format:
     name: Check the code format
     runs-on: ubuntu-latest
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index fd62f43f..1c998b05 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -39,3 +39,6 @@ jobs:
       - name: Test with pytest
         run: |
           pytest --cov=keras_nlp --cov-report xml:coverage.xml
+      - name: Run integration tests
+        run: |
+          python pip_build.py --install && cd integration_tests && pytest .
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 6b51e960..d5af724c 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -97,6 +97,10 @@ recommend using `conda` to install tensorflow dependencies (such as CUDA), and
 `pip` to install python packages from PyPI. The exact method will depend on your
 OS.
 
+**Note**: Please be careful not to use the `tensorflow` pre-packaged with conda,
+which is incompatible with `tensorflow-text` on PyPi, and follow the
+instructions below.
+
 ### Linux (recommended)
 
 To setup a complete environment with TensorFlow, a local install of keras-nlp,
diff --git a/ROADMAP.md b/ROADMAP.md
index 4d92b2fe..ca7c112f 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -148,7 +148,7 @@ are trying to explain to the reader.
 This will continue to be a key investment area for the library. If you have an
 idea for a guide or example, please open an issue to discuss.
 
-By the end of 2022, most new NLP examples on keras.io should be use
+By the end of 2022, most new NLP examples on keras.io should use
 KerasNLP library.
 
 ## Citation bar
diff --git a/keras_nlp/tests/integration_tests/__init__.py b/integration_tests/__init__.py
similarity index 100%
rename from keras_nlp/tests/integration_tests/__init__.py
rename to integration_tests/__init__.py
diff --git a/keras_nlp/tests/integration_tests/basic_usage_test.py b/integration_tests/basic_usage_test.py
similarity index 100%
rename from keras_nlp/tests/integration_tests/basic_usage_test.py
rename to integration_tests/basic_usage_test.py
diff --git a/keras_nlp/tests/integration_tests/import_test.py b/integration_tests/import_test.py
similarity index 100%
rename from keras_nlp/tests/integration_tests/import_test.py
rename to integration_tests/import_test.py
diff --git a/keras_nlp/layers/transformer_layer_utils.py b/keras_nlp/layers/transformer_layer_utils.py
index 2ce24d8a..9bf88aec 100644
--- a/keras_nlp/layers/transformer_layer_utils.py
+++ b/keras_nlp/layers/transformer_layer_utils.py
@@ -18,6 +18,26 @@ import tensorflow as tf
 from absl import logging
 
 
+def _check_masks_shapes(inputs, padding_mask, attention_mask):
+    mask = padding_mask
+    if hasattr(inputs, "_keras_mask") and mask is None:
+        mask = inputs._keras_mask
+    if mask is not None:
+        if mask._rank() != 2:
+            raise ValueError(
+                "`padding_mask` should have shape "
+                "(batch_size, target_length). "
+                f"Received shape `{mask.shape}`."
+            )
+    if attention_mask is not None:
+        if attention_mask._rank() != 3:
+            raise ValueError(
+                "`attention_mask` should have shape "
+                "(batch_size, target_length, source_length). "
+                f"Received shape `{mask.shape}`."
+            )
+
+
 def compute_causal_mask(batch_size, input_length, output_length, cache_index=0):
     """Compute a causal attention mask for a transformer decoder.
 
@@ -60,6 +80,7 @@ def merge_padding_and_attention_mask(
         A merged 2D mask or None. If only `padding_mask` is provided, the
         returned mask is padding_mask with one additional axis.
     """
+    _check_masks_shapes(inputs, padding_mask, attention_mask)
     mask = padding_mask
     if hasattr(inputs, "_keras_mask"):
         if mask is None:
@@ -80,8 +101,5 @@ def merge_padding_and_attention_mask(
         if mask is None:
             return attention_mask
         else:
-            return tf.minimum(
-                mask[:, tf.newaxis, :],
-                attention_mask,
-            )
+            return tf.minimum(mask, attention_mask)
     return mask
diff --git a/keras_nlp/models/backbone.py b/keras_nlp/models/backbone.py
index 8639c5ca..78e7d46c 100644
--- a/keras_nlp/models/backbone.py
+++ b/keras_nlp/models/backbone.py
@@ -18,12 +18,11 @@ import os
 
 from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
 
-@keras_nlp_export("keras_nlp.models.Backbone")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class Backbone(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
diff --git a/keras_nlp/models/preprocessor.py b/keras_nlp/models/preprocessor.py
index ad8702de..546cc763 100644
--- a/keras_nlp/models/preprocessor.py
+++ b/keras_nlp/models/preprocessor.py
@@ -14,12 +14,11 @@
 
 from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
 
-@keras_nlp_export("keras_nlp.models.Preprocessor")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class Preprocessor(keras.layers.Layer):
     """Base class for model preprocessors."""
 
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
index 0d3f1d54..8ea1e243 100644
--- a/keras_nlp/models/task.py
+++ b/keras_nlp/models/task.py
@@ -17,13 +17,12 @@ import os
 
 from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.utils.pipeline_model import PipelineModel
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
 
-@keras_nlp_export("keras_nlp.models.Task")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class Task(PipelineModel):
     """Base class for Task models."""
 


Test patch:
diff --git a/keras_nlp/layers/transformer_layer_utils_test.py b/keras_nlp/layers/transformer_layer_utils_test.py
index 82ee1d93..462aaab1 100644
--- a/keras_nlp/layers/transformer_layer_utils_test.py
+++ b/keras_nlp/layers/transformer_layer_utils_test.py
@@ -24,7 +24,9 @@ class TransformerEncoderTest(tf.test.TestCase):
 
     def test_merge_padding_and_attention_mask(self):
         padding_mask = tf.convert_to_tensor([[1, 1, 0]])
-        attention_mask = tf.convert_to_tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]])
+        attention_mask = tf.convert_to_tensor(
+            [[[0, 0, 1], [0, 1, 0], [1, 0, 0]]]
+        )
         inputs = tf.random.uniform(shape=[1, 3, 2])
         merged_mask = utils.merge_padding_and_attention_mask(
             inputs,
@@ -34,3 +36,26 @@ class TransformerEncoderTest(tf.test.TestCase):
         self.assertTrue(
             (merged_mask.numpy() == [[0, 0, 0], [0, 1, 0], [1, 0, 0]]).all()
         )
+
+    def test_bad_mask_shapes(self):
+        with self.assertRaises(ValueError):
+            padding_mask = tf.convert_to_tensor([[[1, 1, 0], [1, 0, 0]]])
+            attention_mask = tf.convert_to_tensor(
+                [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
+            )
+            inputs = tf.random.uniform(shape=[1, 3, 2])
+            utils.merge_padding_and_attention_mask(
+                inputs,
+                padding_mask,
+                attention_mask,
+            )
+
+        with self.assertRaises(ValueError):
+            padding_mask = tf.convert_to_tensor([[1, 1, 0]])
+            attention_mask = tf.convert_to_tensor([[0, 0, 1], [1, 0, 0]])
+            inputs = tf.random.uniform(shape=[1, 3, 2])
+            utils.merge_padding_and_attention_mask(
+                inputs,
+                padding_mask,
+                attention_mask,
+            )