Merge branch 'master' into gh-action-dep-lint

.github/workflows/actions.yml

@@ -65,3 +65,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 .

.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 .

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,

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

keras_nlp/layers/transformer_layer_utils.py

@@ -18,6 +18,26 @@
 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

keras_nlp/layers/transformer_layer_utils_test.py

@@ -24,7 +24,9 @@ def test_compute_causal_mask(self):
 
     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 @@ def test_merge_padding_and_attention_mask(self):
         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,
+            )

keras_nlp/models/backbone.py

@@ -18,12 +18,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.Backbone")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class Backbone(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

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."""
 

keras_nlp/models/task.py

@@ -17,13 +17,12 @@
 
 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."""