Repository: keras-team/keras-nlp
Pull request #327 (2022-08-29 23:07:24)

Clean up model input names for consistency
This proposes a few changes to the naming our our model and layer inputs
and outputs.

1) Rename `input_ids` -> `token_ids` for bert/roberta.
   Everything is an "input", including the segment id input, so I don't
   think input is a helpful naming prefix in this case.
2) Rename `input_mask`  -> `padding_mask` for bert/roberta.
   This matches the name of the variable for the transformer
   encoder/decoder argument.
3) Rename `tokens` -> `token_ids` for MLMMaskGenerator.
   This layer only operates in id space, so I think token_ids is more
   descriptive and consistent with above.

jbischof: Looks good in general but I'm confused by the `padding_mask` label. Aren't we also masking the holdout words in pretraining and not just the padding?

Does the `examples/bert/` code still work? I see it still uses `input_ids` (here)[https://github.com/keras-team/keras-nlp/blob/master/examples/bert/bert_train.py#L316].

Finally, one other naming inconsistency you could mop up here [sine_position_encoding](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/layers/sine_position_encoding.py#L66).
mattdangerw: @jbischof Thanks! And good call I will rename in `examples/bert` to match.

Re `padding_mask` and the MLM masked words, that's actually a different mask I think. There's actually a `[PAD]` token and a `[MASK]` token. During pretraining, `[PAD]` will be used only to pad sequences, and `[MASK]` will be used for the MLM labels.

So I do think it is accurate to use `padding_mask`. Other options are `mask` or `token_mask`. I am actually fine with any of those. I was just choosing `padding_mask` because that the name we chose for our TransformerEncoder/TransformerDecoder https://keras.io/api/keras_nlp/layers/transformer_encoder/#call-method.

Patch:
diff --git a/keras_nlp/layers/mlm_mask_generator.py b/keras_nlp/layers/mlm_mask_generator.py
index f2c01321..01eac38a 100644
--- a/keras_nlp/layers/mlm_mask_generator.py
+++ b/keras_nlp/layers/mlm_mask_generator.py
@@ -59,10 +59,10 @@ class MLMMaskGenerator(keras.layers.Layer):
 
     Returns:
         A Dict with 4 keys:
-            tokens: Tensor or RaggedTensor, has the same type and shape of
+            token_ids: 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.
+                is None. The positions of token_ids 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
@@ -139,7 +139,7 @@ class MLMMaskGenerator(keras.layers.Layer):
             # convert dense to ragged.
             inputs = tf.RaggedTensor.from_tensor(inputs)
 
-        (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(
+        (token_ids, mask_positions, mask_ids,) = tf_text.mask_language_model(
             inputs,
             item_selector=self._random_selector,
             mask_values_chooser=self._mask_values_chooser,
@@ -147,7 +147,7 @@ class MLMMaskGenerator(keras.layers.Layer):
 
         if not input_is_ragged:
             # If we converted the input from dense to ragged, convert back.
-            tokens = tokens.to_tensor()
+            token_ids = token_ids.to_tensor()
 
         mask_weights = tf.ones_like(mask_positions, self.compute_dtype)
         # If mask_selection_length is set, covert to raggeds to dense.
@@ -159,18 +159,17 @@ class MLMMaskGenerator(keras.layers.Layer):
 
         if input_is_1d:
             # If inputs is 1D, we format the output to be 1D as well.
-            tokens = tf.squeeze(tokens, axis=0)
+            token_ids = tf.squeeze(token_ids, 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,
+        return {
+            "token_ids": token_ids,
             "mask_positions": mask_positions,
             "mask_ids": mask_ids,
             "mask_weights": mask_weights,
         }
-        return output_dict
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
index 362bf1e8..25e672e8 100644
--- a/keras_nlp/models/bert.py
+++ b/keras_nlp/models/bert.py
@@ -93,13 +93,13 @@ class BertCustom(keras.Model):
 
     # Call encoder on the inputs
     input_data = {
-        "input_ids": tf.random.uniform(
+        "token_ids": tf.random.uniform(
             shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
         ),
         "segment_ids": tf.constant(
             [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
-        "input_mask": tf.constant(
+        "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
     }
@@ -128,13 +128,13 @@ class BertCustom(keras.Model):
         cls_token_index = 0
         # Inputs
         token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="input_ids"
+            shape=(None,), dtype="int32", name="token_ids"
         )
         segment_id_input = keras.Input(
             shape=(None,), dtype="int32", name="segment_ids"
         )
-        input_mask = keras.Input(
-            shape=(None,), dtype="int32", name="input_mask"
+        padding_mask = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
         )
 
         # Embed tokens, positions, and segment ids.
@@ -183,7 +183,7 @@ class BertCustom(keras.Model):
                 dropout=dropout,
                 kernel_initializer=_bert_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, padding_mask=input_mask)
+            )(x, padding_mask=padding_mask)
 
         # Construct the two Bert outputs. The pooled output is a dense layer on
         # top of the [CLS] token.
@@ -198,9 +198,9 @@ class BertCustom(keras.Model):
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
-                "input_ids": token_id_input,
+                "token_ids": token_id_input,
                 "segment_ids": segment_id_input,
-                "input_mask": input_mask,
+                "padding_mask": padding_mask,
             },
             outputs={
                 "sequence_output": sequence_output,
@@ -264,13 +264,13 @@ class BertClassifier(keras.Model):
 
     # Call classifier on the inputs.
     input_data = {
-        "input_ids": tf.random.uniform(
+        "token_ids": tf.random.uniform(
             shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
         ),
         "segment_ids": tf.constant(
             [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
-        "input_mask": tf.constant(
+        "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
     }
@@ -329,11 +329,11 @@ MODEL_DOCSTRING = """Bi-directional Transformer-based encoder network (Bert)
 
     # Call encoder on the inputs.
     input_data = {{
-        "input_ids": tf.random.uniform(
+        "token_ids": tf.random.uniform(
             shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
         ),
         "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
-        "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
     }}
     output = model(input_data)
 
diff --git a/keras_nlp/models/roberta.py b/keras_nlp/models/roberta.py
index e85c44d9..fa46833f 100644
--- a/keras_nlp/models/roberta.py
+++ b/keras_nlp/models/roberta.py
@@ -67,9 +67,9 @@ class RobertaCustom(keras.Model):
 
     # Call encoder on the inputs.
     input_data = {
-        "input_ids": tf.random.uniform(
+        "token_ids": tf.random.uniform(
             shape=(1, 12), dtype=tf.int64, maxval=50265),
-        "input_mask": tf.constant(
+        "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
     output = model(input_data)
@@ -93,10 +93,10 @@ class RobertaCustom(keras.Model):
         cls_token_index = 0
         # Inputs
         token_id_input = keras.Input(
-            shape=(None,), dtype=tf.int32, name="input_ids"
+            shape=(None,), dtype=tf.int32, name="token_ids"
         )
-        input_mask = keras.Input(
-            shape=(None,), dtype=tf.int32, name="input_mask"
+        padding_mask = keras.Input(
+            shape=(None,), dtype=tf.int32, name="padding_mask"
         )
 
         # Embed tokens and positions.
@@ -130,13 +130,13 @@ class RobertaCustom(keras.Model):
                 dropout=dropout,
                 kernel_initializer=_roberta_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, padding_mask=input_mask)
+            )(x, padding_mask=padding_mask)
 
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
-                "input_ids": token_id_input,
-                "input_mask": input_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask,
             },
             outputs={
                 "sequence_output": x,
@@ -197,9 +197,9 @@ class RobertaClassifier(keras.Model):
 
     # Call classifier on the inputs.
     input_data = {
-        "input_ids": tf.random.uniform(
+        "token_ids": tf.random.uniform(
             shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size),
-        "input_mask": tf.constant(
+        "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
     classifier = keras_nlp.models.RobertaClassifier(model, 4)
@@ -265,9 +265,9 @@ def RobertaBase(vocabulary_size, name=None, trainable=True):
 
     # Call encoder on the inputs.
     input_data = {
-        "input_ids": tf.random.uniform(
+        "token_ids": tf.random.uniform(
             shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size),
-        "input_mask": tf.ones((1, 512)),
+        "padding_mask": tf.ones((1, 512)),
     }
     output = model(input_data)
     ```


Test patch:
diff --git a/keras_nlp/layers/mlm_mask_generator_test.py b/keras_nlp/layers/mlm_mask_generator_test.py
index 342fa6f0..4445d948 100644
--- a/keras_nlp/layers/mlm_mask_generator_test.py
+++ b/keras_nlp/layers/mlm_mask_generator_test.py
@@ -49,18 +49,18 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
         )
         inputs = tf.ragged.constant([[5, 3, 2], [1, 2, 3, 4, 5]])
         outputs = mlm_masker(inputs)
-        tokens, mask_positions, mask_ids = (
-            outputs["tokens"],
+        token_ids, mask_positions, mask_ids = (
+            outputs["token_ids"],
             outputs["mask_positions"],
             outputs["mask_ids"],
         )
-        self.assertEqual(type(tokens), type(inputs))
-        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertEqual(type(token_ids), type(inputs))
+        self.assertAllEqual(token_ids.shape, inputs.shape)
         self.assertAllEqual(mask_positions.shape, mask_ids.shape)
 
-        # Test all selected tokens are correctly masked.
+        # Test all selected token_ids are correctly masked.
         masked_values = tf.gather(
-            tokens,
+            token_ids,
             mask_positions,
             batch_dims=1,
         )
@@ -81,17 +81,17 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
             dtype=tf.int32,
         )
         outputs = mlm_masker(inputs)
-        tokens, mask_positions, mask_ids = (
-            outputs["tokens"],
+        token_ids, mask_positions, mask_ids = (
+            outputs["token_ids"],
             outputs["mask_positions"],
             outputs["mask_ids"],
         )
-        self.assertEqual(type(tokens), type(inputs))
-        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertEqual(type(token_ids), type(inputs))
+        self.assertAllEqual(token_ids.shape, inputs.shape)
         self.assertAllEqual(mask_positions.shape, mask_ids.shape)
-        # Test all selected tokens are correctly masked.
+        # Test all selected token_ids are correctly masked.
         masked_values = tf.gather(
-            tokens,
+            token_ids,
             mask_positions,
             batch_dims=1,
         )
@@ -108,7 +108,7 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
         )
         inputs = tf.constant([1, 2, 3, 4, 5])
         outputs = mlm_masker(inputs)
-        self.assertAllEqual(outputs["tokens"].shape, inputs.shape)
+        self.assertAllEqual(outputs["token_ids"].shape, inputs.shape)
 
     def test_number_of_masked_position_as_expected(self):
         mask_selection_rate = 0.5
@@ -162,21 +162,21 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
             dtype=tf.int32,
         )
         outputs = mlm_masker(inputs)
-        tokens, mask_positions, mask_ids = (
-            outputs["tokens"],
+        tokens_ids, mask_positions, mask_ids = (
+            outputs["token_ids"],
             outputs["mask_positions"],
             outputs["mask_ids"],
         )
-        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(tokens_ids.shape, inputs.shape)
         self.assertAllEqual(
             mask_positions.row_lengths(), mask_ids.row_lengths()
         )
         masked_values = tf.gather(
-            tokens,
+            tokens_ids,
             mask_positions,
             batch_dims=1,
         )
-        # Verify that selected tokens are replaced by random tokens.
+        # Verify that selected tokens_ids are replaced by random tokens_ids.
         self.assertNotEqual(tf.reduce_mean(masked_values), self.mask_token_id)
 
     def test_invalid_mask_token(self):
@@ -262,6 +262,6 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
         batch_first = ds.batch(8).map(mlm_masker)
         batch_second = ds.map(mlm_masker).batch(8)
         self.assertEqual(
-            batch_first.take(1).get_single_element()["tokens"].shape,
-            batch_second.take(1).get_single_element()["tokens"].shape,
+            batch_first.take(1).get_single_element()["token_ids"].shape,
+            batch_second.take(1).get_single_element()["token_ids"].shape,
         )
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
index f8aaca4f..e86d92af 100644
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -34,13 +34,13 @@ class BertTest(tf.test.TestCase):
         )
         self.batch_size = 8
         self.input_data = {
-            "input_ids": tf.ones(
+            "token_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
             "segment_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.ones(
+            "padding_mask": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
         }
@@ -51,13 +51,13 @@ class BertTest(tf.test.TestCase):
     def test_variable_sequence_length_call_bert(self):
         for seq_length in (25, 50, 75):
             input_data = {
-                "input_ids": tf.ones(
+                "token_ids": tf.ones(
                     (self.batch_size, seq_length), dtype="int32"
                 ),
                 "segment_ids": tf.ones(
                     (self.batch_size, seq_length), dtype="int32"
                 ),
-                "input_mask": tf.ones(
+                "padding_mask": tf.ones(
                     (self.batch_size, seq_length), dtype="int32"
                 ),
             }
@@ -70,13 +70,13 @@ class BertTest(tf.test.TestCase):
     def test_valid_call_bert_base(self):
         model = bert.BertBase(vocabulary_size=1000, name="encoder")
         input_data = {
-            "input_ids": tf.ones(
+            "token_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
             "segment_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.ones(
+            "padding_mask": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
         }
diff --git a/keras_nlp/models/roberta_test.py b/keras_nlp/models/roberta_test.py
index 6dee0fa8..bfc40d3a 100644
--- a/keras_nlp/models/roberta_test.py
+++ b/keras_nlp/models/roberta_test.py
@@ -34,10 +34,10 @@ class RobertaTest(tf.test.TestCase):
         )
         self.batch_size = 8
         self.input_data = {
-            "input_ids": tf.ones(
+            "token_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.ones(
+            "padding_mask": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
         }
@@ -54,10 +54,10 @@ class RobertaTest(tf.test.TestCase):
     def test_valid_call_roberta_base(self):
         model = roberta.RobertaBase(vocabulary_size=1000, name="encoder")
         input_data = {
-            "input_ids": tf.ones(
+            "token_ids": tf.ones(
                 (self.batch_size, model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.ones(
+            "padding_mask": tf.ones(
                 (self.batch_size, model.max_sequence_length), dtype="int32"
             ),
         }
@@ -66,10 +66,10 @@ class RobertaTest(tf.test.TestCase):
     def test_variable_sequence_length_call_roberta(self):
         for seq_length in (25, 50, 75):
             input_data = {
-                "input_ids": tf.ones(
+                "token_ids": tf.ones(
                     (self.batch_size, seq_length), dtype="int32"
                 ),
-                "input_mask": tf.ones(
+                "padding_mask": tf.ones(
                     (self.batch_size, seq_length), dtype="int32"
                 ),
             }
diff --git a/network_tests/test_load_ckpts.py b/network_tests/test_load_ckpts.py
index 2d00ef75..ddb251f9 100644
--- a/network_tests/test_load_ckpts.py
+++ b/network_tests/test_load_ckpts.py
@@ -22,21 +22,21 @@ class BertCkptTest(tf.test.TestCase):
     def test_load_bert_base_uncased(self):
         model = keras_nlp.models.BertBase(weights="uncased_en")
         input_data = {
-            "input_ids": tf.random.uniform(
+            "token_ids": tf.random.uniform(
                 shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
             ),
             "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
-            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+            "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
         }
         model(input_data)
 
     def test_load_bert_base_cased(self):
         model = keras_nlp.models.BertBase(weights="cased_en")
         input_data = {
-            "input_ids": tf.random.uniform(
+            "token_ids": tf.random.uniform(
                 shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
             ),
             "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
-            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+            "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
         }
         model(input_data)