Repository: keras-team/keras-nlp
Pull request #61 (2022-03-25 00:01:15)

Add a class to generate LM masks
I am adding a class to generate LM masks, which can be useful for models pretrained with MLM  (masked language modeling) tasks, such as BERT, RoBERTa and etc. 

This is only an initial commit, let's discuss the API interface from here, thanks!


Patch:
diff --git a/examples/machine_translation/train.py b/examples/machine_translation/train.py
index c65125c6..648585bc 100644
--- a/examples/machine_translation/train.py
+++ b/examples/machine_translation/train.py
@@ -77,9 +77,12 @@ def run_training(model, train_ds, val_ds):
 
 
 def main(_):
-    (train_ds, val_ds, test_ds), (
-        eng_tokenizer,
-        spa_tokenizer,
+    (
+        (train_ds, val_ds, test_ds),
+        (
+            eng_tokenizer,
+            spa_tokenizer,
+        ),
     ) = get_dataset_and_tokenizer(
         FLAGS.sequence_length, FLAGS.vocab_size, FLAGS.batch_size
     )
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index fe677bbe..9e246621 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -14,6 +14,7 @@
 
 from keras_nlp.layers.fnet_encoder import FNetEncoder
 from keras_nlp.layers.position_embedding import PositionEmbedding
+from keras_nlp.layers.preprocessing import MLMMaskGenerator
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
 from keras_nlp.layers.transformer_decoder import TransformerDecoder
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
diff --git a/keras_nlp/layers/preprocessing/__init__.py b/keras_nlp/layers/preprocessing/__init__.py
new file mode 100644
index 00000000..c8059499
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_nlp.layers.preprocessing.mlm_mask_generator import MLMMaskGenerator
diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator.py b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
new file mode 100644
index 00000000..59daa236
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
@@ -0,0 +1,202 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+import tensorflow_text as tf_text
+from tensorflow import keras
+
+
+class MLMMaskGenerator(keras.layers.Layer):
+    """Layer that applies language model masking.
+
+    This layer is useful for preparing inputs for masked languaged modeling
+    (MLM) tasks. It follows the masking strategy described in the [original BERT
+    paper](https://arxiv.org/abs/1810.04805). Given tokenized text,
+    it randomly selects certain number of tokens for masking. Then for each
+    selected token, it has a chance (configurable) to be replaced by
+    "mask token" or random token, or stay unchanged.
+
+    Users should use this layer with `tf.data` to generate masks.
+
+    Args:
+        vocabulary_size: int, the size of the vocabulary.
+        mask_selection_rate: float, the probability of a token is selected for
+            masking.
+        mask_token_id: int. The id of mask token.
+        mask_selection_length: int, defaults to None. Maximum number of tokens
+            selected for  masking in each sequence. If set, the output
+            `mask_positions`, `mask_ids` and `mask_weights` will be padded
+            to dense tensors of length `mask_selection_length`,
+            otherwise the output will be a RaggedTensor.
+        unselectable_token_ids: A list of tokens, defaults to [0] (the default
+            `padding_token_id`).
+        mask_token_rate: float, defaults to 0.8. `mask_token_rate` must be
+            between 0 and 1 which indicates how often the mask_token is
+            substituted for tokens selected for masking.
+        random_token_rate: float, defaults to 0.1. `random_token_rate` must be
+            between 0 and 1 which indicates how often a random token is
+            substituted for tokens selected for masking. Default is 0.1.
+            Note: mask_token_rate + random_token_rate <= 1,  and for
+            (1 - mask_token_rate - random_token_rate), the token will not be
+            changed.
+
+    Input:
+        A 1D integer tensor of shape [sequence_length] or a 2D integer tensor
+        of shape [batch_size, sequence_length], or a 2D integer RaggedTensor.
+        Represents the sequence to mask.
+
+    Returns:
+        A Dict with 4 keys:
+            tokens: Tensor or RaggedTensor, has the same type and shape of
+                input. Sequence after getting masked.
+            mask_positions: Tensor, or RaggedTensor if `mask_selection_length`
+                is None. The positions of tokens getting masked.
+            mask_ids: Tensor, or RaggedTensor if  `mask_selection_length` is
+                None. The original token ids at masked positions.
+            mask_weights: Tensor, or RaggedTensor if `mask_selection_length` is
+                None. `mask_weights` has the same shape as `mask_positions` and
+                `mask_ids`. Each element in `mask_weights` should be 0 or 1,
+                1 means the corresponding position in `mask_positions` is an
+                actual mask, 0 means it is a pad.
+
+    Examples:
+
+    Basic usage.
+    >>> masker = keras_nlp.preprocessing.MLMMaskGenerator( \
+    ...     vocaublary_size=10, mask_selection_rate=0.2, mask_token_id=0)
+    >>> masker(tf.constant([1, 2, 3, 4, 5]))
+    {'tokens': <tf.Tensor: shape=(5,), dtype=int32,
+        numpy=array([1, 2, 3, 4, 0], dtype=int32)>,
+    'mask_positions': <tf.Tensor: shape=(5,), dtype=int64,
+        numpy=array([4, 0, 0, 0, 0])>,
+    'mask_ids': <tf.Tensor: shape=(5,), dtype=int32,
+        numpy=array([5, 0, 0, 0, 0], dtype=int32)>,
+    'mask_weights': <tf.Tensor: shape=(1, 5), dtype=int64,
+        numpy=array([[1, 0, 0, 0, 0]])>}
+
+    Ragged Input:
+    >>> masker = keras_nlp.preprocessing.MLMMaskGenerator( \
+    ...     vocaublary_size=10, mask_selection_rate=0.2, mask_token_id=0)
+    >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))
+    {'tokens': <tf.RaggedTensor [[1, 4], [0, 2, 3, 0]]>,
+    'mask_positions': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
+        array([[1, 0, 0, 0, 0],
+               [0, 3, 0, 0, 0]])>,
+    'mask_ids': <tf.Tensor: shape=(2, 5), dtype=int32, numpy=
+        array([[2, 0, 0, 0, 0],
+               [1, 4, 0, 0, 0]], dtype=int32)>,
+    'mask_weights': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
+        array([[1, 0, 0, 0, 0],
+               [1, 1, 0, 0, 0]])>}
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        mask_selection_rate,
+        mask_token_id,
+        mask_selection_length=None,
+        unselectable_token_ids=[0],
+        mask_token_rate=0.8,
+        random_token_rate=0.1,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocabulary_size = vocabulary_size
+        self.unselectable_token_ids = unselectable_token_ids
+        self.mask_selection_rate = mask_selection_rate
+        self.mask_selection_length = mask_selection_length
+        self.mask_token_rate = mask_token_rate
+        self.random_token_rate = random_token_rate
+
+        if mask_token_id >= vocabulary_size:
+            raise ValueError(
+                f"Mask token id should be in range [0, vocabulary_size - 1], "
+                f"but received mask_token_id={mask_token_id}."
+            )
+        self.mask_token_id = mask_token_id
+
+        max_selections = self.mask_selection_length
+        if max_selections is None:
+            # Set a large number to remove the `max_selections_per_batch` cap.
+            max_selections = 2**31 - 1
+        self._random_selector = tf_text.RandomItemSelector(
+            max_selections_per_batch=max_selections,
+            selection_rate=self.mask_selection_rate,
+            unselectable_ids=self.unselectable_token_ids,
+        )
+        self._mask_values_chooser = tf_text.MaskValuesChooser(
+            self.vocabulary_size,
+            self.mask_token_id,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+        )
+
+    def call(self, inputs):
+        input_is_ragged = isinstance(inputs, tf.RaggedTensor)
+        input_is_1d = tf.rank(inputs) == 1
+        if input_is_1d:
+            # If inputs is of rank 1, we manually add the batch axis.
+            inputs = inputs[tf.newaxis, :]
+        if not input_is_ragged:
+            # `tf_text.mask_language_model` requires a ragged tensor, so
+            # convert dense to ragged.
+            inputs = tf.RaggedTensor.from_tensor(inputs)
+        (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(
+            inputs,
+            item_selector=self._random_selector,
+            mask_values_chooser=self._mask_values_chooser,
+        )
+
+        if not input_is_ragged:
+            # If we converted the input from dense to ragged, convert back.
+            tokens = tokens.to_tensor()
+
+        mask_weights = tf.ones_like(mask_positions, self.compute_dtype)
+        # If mask_selection_length is set, covert to raggeds to dense.
+        if self.mask_selection_length:
+            target_shape = tf.cast([-1, self.mask_selection_length], tf.int64)
+            mask_positions = mask_positions.to_tensor(shape=target_shape)
+            mask_ids = mask_ids.to_tensor(shape=target_shape)
+            mask_weights = mask_weights.to_tensor(shape=target_shape)
+
+        if input_is_1d:
+            # If inputs is 1D, we format the output to be 1D as well.
+            tokens = tf.squeeze(tokens, axis=0)
+            mask_positions = tf.squeeze(mask_positions, axis=0)
+            mask_ids = tf.squeeze(mask_ids, axis=0)
+            mask_weights = tf.squeeze(mask_weights, axis=0)
+
+        output_dict = {
+            "tokens": tokens,
+            "mask_positions": mask_positions,
+            "mask_ids": mask_ids,
+            "mask_weights": mask_weights,
+        }
+        return output_dict
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "mask_selection_rate": self.mask_selection_rate,
+                "mask_selection_length": self.mask_selection_length,
+                "unselectable_token_ids": self.unselectable_token_ids,
+                "mask_token_id": self.mask_token_id,
+                "mask_token_rate": self.mask_token_rate,
+                "random_token_rate": self.random_token_rate,
+            }
+        )
+        return config


Test patch:
diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py b/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
new file mode 100644
index 00000000..66547984
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
@@ -0,0 +1,234 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+
+from keras_nlp.layers.preprocessing.mlm_mask_generator import MLMMaskGenerator
+
+
+class MLMMaskGeneratorTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.VOCAB = [
+            "[UNK]",
+            "[MASK]",
+            "[RANDOM]",
+            "[CLS]",
+            "[SEP]",
+            "do",
+            "you",
+            "like",
+            "machine",
+            "learning",
+            "welcome",
+            "to",
+            "keras",
+        ]
+        self.mask_token_id = self.VOCAB.index("[MASK]")
+        self.vocabulary_size = len(self.VOCAB)
+
+    def test_mask_ragged_tensor(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_selection_length=5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = tf.ragged.constant([[5, 3, 2], [1, 2, 3, 4, 5]])
+        outputs = mlm_masker(inputs)
+        tokens, mask_positions, mask_ids = (
+            outputs["tokens"],
+            outputs["mask_positions"],
+            outputs["mask_ids"],
+        )
+        self.assertEqual(type(tokens), type(inputs))
+        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(mask_positions.shape, mask_ids.shape)
+
+        # Test all selected tokens are correctly masked.
+        masked_values = tf.gather(
+            tokens,
+            mask_positions,
+            batch_dims=1,
+        )
+        self.assertEqual(tf.reduce_mean(masked_values), self.mask_token_id)
+
+    def test_mask_tensor(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_selection_length=5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = tf.random.uniform(
+            shape=[5, 10],
+            maxval=len(self.VOCAB),
+            dtype=tf.int32,
+        )
+        outputs = mlm_masker(inputs)
+        tokens, mask_positions, mask_ids = (
+            outputs["tokens"],
+            outputs["mask_positions"],
+            outputs["mask_ids"],
+        )
+        self.assertEqual(type(tokens), type(inputs))
+        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(mask_positions.shape, mask_ids.shape)
+        # Test all selected tokens are correctly masked.
+        masked_values = tf.gather(
+            tokens,
+            mask_positions,
+            batch_dims=1,
+        )
+        self.assertEqual(tf.reduce_mean(masked_values), self.mask_token_id)
+
+    def test_mask_1d_input(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_selection_length=5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = tf.constant([1, 2, 3, 4, 5])
+        outputs = mlm_masker(inputs)
+        self.assertAllEqual(outputs["tokens"].shape, inputs.shape)
+
+    def test_number_of_masked_position_as_expected(self):
+        mask_selection_rate = 0.5
+        mask_selection_length = 5
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=mask_selection_rate,
+            mask_token_id=self.mask_token_id,
+            unselectable_token_ids=None,
+        )
+        inputs = tf.ragged.constant(
+            [[0, 1, 2], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4]]
+        )
+        outputs = mlm_masker(inputs)
+        expected_number_of_masked_tokens = tf.cast(
+            tf.math.ceil(
+                tf.cast(inputs.row_lengths(), dtype=tf.float32)
+                * mask_selection_rate,
+            ),
+            dtype=tf.int64,
+        )
+
+        self.assertAllEqual(
+            outputs["mask_positions"].row_lengths(),
+            expected_number_of_masked_tokens,
+        )
+
+        # Cap the number of masked tokens at 0, so we can test if
+        # mask_selection_length takes effect.
+        mask_selection_length = 0
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=mask_selection_rate,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=mask_selection_length,
+        )
+        outputs = mlm_masker(inputs)
+        self.assertEqual(tf.reduce_sum(outputs["mask_positions"]), 0)
+
+    def test_apply_random_token_not_mask(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=0,
+            random_token_rate=1,
+        )
+        inputs = tf.random.uniform(
+            shape=[5, 10],
+            maxval=len(self.VOCAB),
+            dtype=tf.int32,
+        )
+        outputs = mlm_masker(inputs)
+        tokens, mask_positions, mask_ids = (
+            outputs["tokens"],
+            outputs["mask_positions"],
+            outputs["mask_ids"],
+        )
+        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(
+            mask_positions.row_lengths(), mask_ids.row_lengths()
+        )
+        masked_values = tf.gather(
+            tokens,
+            mask_positions,
+            batch_dims=1,
+        )
+        # Verify that selected tokens are replaced by random tokens.
+        self.assertNotEqual(tf.reduce_mean(masked_values), self.mask_token_id)
+
+    def test_invalid_mask_token(self):
+        with self.assertRaisesRegex(ValueError, "Mask token id should be*"):
+            _ = MLMMaskGenerator(
+                vocabulary_size=self.vocabulary_size,
+                mask_selection_rate=0.5,
+                mask_token_id=self.vocabulary_size,
+                mask_selection_length=5,
+            )
+
+    def test_unselectable_tokens(self):
+        unselectable_token_ids = [
+            self.vocabulary_size - 1,
+            self.vocabulary_size - 2,
+        ]
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=1,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=5,
+            unselectable_token_ids=unselectable_token_ids,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = [unselectable_token_ids]
+        outputs = mlm_masker(inputs)
+        # Verify that no token is masked out.
+        self.assertEqual(tf.reduce_sum(outputs["mask_positions"]), 0)
+
+    def test_config(self):
+        unselectable_token_ids = [
+            self.vocabulary_size - 1,
+            self.vocabulary_size - 2,
+        ]
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=5,
+            unselectable_token_ids=unselectable_token_ids,
+        )
+        config = mlm_masker.get_config()
+        expected_config = {
+            "vocabulary_size": self.vocabulary_size,
+            "unselectable_token_ids": unselectable_token_ids,
+        }
+        self.assertDictContainsSubset(expected_config, config)
+
+        # Test cloned mlm_masker can be run.
+        cloned_mlm_masker = MLMMaskGenerator.from_config(config)
+        inputs = tf.ragged.constant(
+            [[0, 1, 2], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4]]
+        )
+        cloned_mlm_masker(inputs)