We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Error with tf.data and MLMMaskGenerator for dense inputs
When using the MLMMaskGenerator for to map over dense, batched inputs in a tf.data.Dataset, we get the following error...

`TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'`

This colab has a reproduction https://colab.research.google.com/gist/mattdangerw/4596df85105ff6e6731128fc79d16bf3/mlmmaskgenerator-bug.ipynb

tf.data + dense, batched inputs might be the most common use case for this layer, so this is an important one to fix.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `keras_nlp/layers/preprocessing/mlm_mask_generator.py`
Content:
```
1 # Copyright 2022 The KerasNLP Authors
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import tensorflow as tf
16 import tensorflow_text as tf_text
17 from tensorflow import keras
18 
19 
20 class MLMMaskGenerator(keras.layers.Layer):
21     """Layer that applies language model masking.
22 
23     This layer is useful for preparing inputs for masked languaged modeling
24     (MLM) tasks. It follows the masking strategy described in the [original BERT
25     paper](https://arxiv.org/abs/1810.04805). Given tokenized text,
26     it randomly selects certain number of tokens for masking. Then for each
27     selected token, it has a chance (configurable) to be replaced by
28     "mask token" or random token, or stay unchanged.
29 
30     Users should use this layer with `tf.data` to generate masks.
31 
32     Args:
33         vocabulary_size: int, the size of the vocabulary.
34         mask_selection_rate: float, the probability of a token is selected for
35             masking.
36         mask_token_id: int. The id of mask token.
37         mask_selection_length: int, defaults to None. Maximum number of tokens
38             selected for  masking in each sequence. If set, the output
39             `mask_positions`, `mask_ids` and `mask_weights` will be padded
40             to dense tensors of length `mask_selection_length`,
41             otherwise the output will be a RaggedTensor.
42         unselectable_token_ids: A list of tokens, defaults to [0] (the default
43             `padding_token_id`).
44         mask_token_rate: float, defaults to 0.8. `mask_token_rate` must be
45             between 0 and 1 which indicates how often the mask_token is
46             substituted for tokens selected for masking.
47         random_token_rate: float, defaults to 0.1. `random_token_rate` must be
48             between 0 and 1 which indicates how often a random token is
49             substituted for tokens selected for masking. Default is 0.1.
50             Note: mask_token_rate + random_token_rate <= 1,  and for
51             (1 - mask_token_rate - random_token_rate), the token will not be
52             changed.
53 
54     Input:
55         A 1D integer tensor of shape [sequence_length] or a 2D integer tensor
56         of shape [batch_size, sequence_length], or a 2D integer RaggedTensor.
57         Represents the sequence to mask.
58 
59     Returns:
60         A Dict with 4 keys:
61             tokens: Tensor or RaggedTensor, has the same type and shape of
62                 input. Sequence after getting masked.
63             mask_positions: Tensor, or RaggedTensor if `mask_selection_length`
64                 is None. The positions of tokens getting masked.
65             mask_ids: Tensor, or RaggedTensor if  `mask_selection_length` is
66                 None. The original token ids at masked positions.
67             mask_weights: Tensor, or RaggedTensor if `mask_selection_length` is
68                 None. `mask_weights` has the same shape as `mask_positions` and
69                 `mask_ids`. Each element in `mask_weights` should be 0 or 1,
70                 1 means the corresponding position in `mask_positions` is an
71                 actual mask, 0 means it is a pad.
72 
73     Examples:
74 
75     Basic usage.
76     >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \
77             vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0, \
78             mask_selection_length=5)
79     >>> masker(tf.constant([1, 2, 3, 4, 5]))
80 
81     Ragged Input:
82     >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \
83             vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0, \
84             mask_selection_length=5)
85     >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))
86     """
87 
88     def __init__(
89         self,
90         vocabulary_size,
91         mask_selection_rate,
92         mask_token_id,
93         mask_selection_length=None,
94         unselectable_token_ids=[0],
95         mask_token_rate=0.8,
96         random_token_rate=0.1,
97         **kwargs,
98     ):
99         super().__init__(**kwargs)
100         self.vocabulary_size = vocabulary_size
101         self.unselectable_token_ids = unselectable_token_ids
102         self.mask_selection_rate = mask_selection_rate
103         self.mask_selection_length = mask_selection_length
104         self.mask_token_rate = mask_token_rate
105         self.random_token_rate = random_token_rate
106 
107         if mask_token_id >= vocabulary_size:
108             raise ValueError(
109                 f"Mask token id should be in range [0, vocabulary_size - 1], "
110                 f"but received mask_token_id={mask_token_id}."
111             )
112         self.mask_token_id = mask_token_id
113 
114         max_selections = self.mask_selection_length
115         if max_selections is None:
116             # Set a large number to remove the `max_selections_per_batch` cap.
117             max_selections = 2**31 - 1
118         self._random_selector = tf_text.RandomItemSelector(
119             max_selections_per_batch=max_selections,
120             selection_rate=self.mask_selection_rate,
121             unselectable_ids=self.unselectable_token_ids,
122         )
123         self._mask_values_chooser = tf_text.MaskValuesChooser(
124             self.vocabulary_size,
125             self.mask_token_id,
126             mask_token_rate=self.mask_token_rate,
127             random_token_rate=self.random_token_rate,
128         )
129 
130     def call(self, inputs):
131         input_is_ragged = isinstance(inputs, tf.RaggedTensor)
132         input_is_1d = tf.rank(inputs) == 1
133         if input_is_1d:
134             # If inputs is of rank 1, we manually add the batch axis.
135             inputs = inputs[tf.newaxis, :]
136         if not input_is_ragged:
137             # `tf_text.mask_language_model` requires a ragged tensor, so
138             # convert dense to ragged.
139             inputs = tf.RaggedTensor.from_tensor(inputs)
140         (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(
141             inputs,
142             item_selector=self._random_selector,
143             mask_values_chooser=self._mask_values_chooser,
144         )
145 
146         if not input_is_ragged:
147             # If we converted the input from dense to ragged, convert back.
148             tokens = tokens.to_tensor()
149 
150         mask_weights = tf.ones_like(mask_positions, self.compute_dtype)
151         # If mask_selection_length is set, covert to raggeds to dense.
152         if self.mask_selection_length:
153             target_shape = tf.cast([-1, self.mask_selection_length], tf.int64)
154             mask_positions = mask_positions.to_tensor(shape=target_shape)
155             mask_ids = mask_ids.to_tensor(shape=target_shape)
156             mask_weights = mask_weights.to_tensor(shape=target_shape)
157 
158         if input_is_1d:
159             # If inputs is 1D, we format the output to be 1D as well.
160             tokens = tf.squeeze(tokens, axis=0)
161             mask_positions = tf.squeeze(mask_positions, axis=0)
162             mask_ids = tf.squeeze(mask_ids, axis=0)
163             mask_weights = tf.squeeze(mask_weights, axis=0)
164 
165         output_dict = {
166             "tokens": tokens,
167             "mask_positions": mask_positions,
168             "mask_ids": mask_ids,
169             "mask_weights": mask_weights,
170         }
171         return output_dict
172 
173     def get_config(self):
174         config = super().get_config()
175         config.update(
176             {
177                 "vocabulary_size": self.vocabulary_size,
178                 "mask_selection_rate": self.mask_selection_rate,
179                 "mask_selection_length": self.mask_selection_length,
180                 "unselectable_token_ids": self.unselectable_token_ids,
181                 "mask_token_id": self.mask_token_id,
182                 "mask_token_rate": self.mask_token_rate,
183                 "random_token_rate": self.random_token_rate,
184             }
185         )
186         return config
187
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator.py b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
--- a/keras_nlp/layers/preprocessing/mlm_mask_generator.py
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
@@ -129,7 +129,7 @@
 
     def call(self, inputs):
         input_is_ragged = isinstance(inputs, tf.RaggedTensor)
-        input_is_1d = tf.rank(inputs) == 1
+        input_is_1d = inputs.shape.rank == 1
         if input_is_1d:
             # If inputs is of rank 1, we manually add the batch axis.
             inputs = inputs[tf.newaxis, :]
@@ -137,6 +137,7 @@
             # `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,

{"golden_diff": "diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator.py b/keras_nlp/layers/preprocessing/mlm_mask_generator.py\n--- a/keras_nlp/layers/preprocessing/mlm_mask_generator.py\n+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator.py\n@@ -129,7 +129,7 @@\n \n     def call(self, inputs):\n         input_is_ragged = isinstance(inputs, tf.RaggedTensor)\n-        input_is_1d = tf.rank(inputs) == 1\n+        input_is_1d = inputs.shape.rank == 1\n         if input_is_1d:\n             # If inputs is of rank 1, we manually add the batch axis.\n             inputs = inputs[tf.newaxis, :]\n@@ -137,6 +137,7 @@\n             # `tf_text.mask_language_model` requires a ragged tensor, so\n             # convert dense to ragged.\n             inputs = tf.RaggedTensor.from_tensor(inputs)\n+\n         (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(\n             inputs,\n             item_selector=self._random_selector,\n", "issue": "Error with tf.data and MLMMaskGenerator for dense inputs\nWhen using the MLMMaskGenerator for to map over dense, batched inputs in a tf.data.Dataset, we get the following error...\r\n\r\n`TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'`\r\n\r\nThis colab has a reproduction https://colab.research.google.com/gist/mattdangerw/4596df85105ff6e6731128fc79d16bf3/mlmmaskgenerator-bug.ipynb\r\n\r\ntf.data + dense, batched inputs might be the most common use case for this layer, so this is an important one to fix.\n", "before_files": [{"content": "# Copyright 2022 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tensorflow as tf\nimport tensorflow_text as tf_text\nfrom tensorflow import keras\n\n\nclass MLMMaskGenerator(keras.layers.Layer):\n    \"\"\"Layer that applies language model masking.\n\n    This layer is useful for preparing inputs for masked languaged modeling\n    (MLM) tasks. It follows the masking strategy described in the [original BERT\n    paper](https://arxiv.org/abs/1810.04805). Given tokenized text,\n    it randomly selects certain number of tokens for masking. Then for each\n    selected token, it has a chance (configurable) to be replaced by\n    \"mask token\" or random token, or stay unchanged.\n\n    Users should use this layer with `tf.data` to generate masks.\n\n    Args:\n        vocabulary_size: int, the size of the vocabulary.\n        mask_selection_rate: float, the probability of a token is selected for\n            masking.\n        mask_token_id: int. The id of mask token.\n        mask_selection_length: int, defaults to None. Maximum number of tokens\n            selected for  masking in each sequence. If set, the output\n            `mask_positions`, `mask_ids` and `mask_weights` will be padded\n            to dense tensors of length `mask_selection_length`,\n            otherwise the output will be a RaggedTensor.\n        unselectable_token_ids: A list of tokens, defaults to [0] (the default\n            `padding_token_id`).\n        mask_token_rate: float, defaults to 0.8. `mask_token_rate` must be\n            between 0 and 1 which indicates how often the mask_token is\n            substituted for tokens selected for masking.\n        random_token_rate: float, defaults to 0.1. `random_token_rate` must be\n            between 0 and 1 which indicates how often a random token is\n            substituted for tokens selected for masking. Default is 0.1.\n            Note: mask_token_rate + random_token_rate <= 1,  and for\n            (1 - mask_token_rate - random_token_rate), the token will not be\n            changed.\n\n    Input:\n        A 1D integer tensor of shape [sequence_length] or a 2D integer tensor\n        of shape [batch_size, sequence_length], or a 2D integer RaggedTensor.\n        Represents the sequence to mask.\n\n    Returns:\n        A Dict with 4 keys:\n            tokens: Tensor or RaggedTensor, has the same type and shape of\n                input. Sequence after getting masked.\n            mask_positions: Tensor, or RaggedTensor if `mask_selection_length`\n                is None. The positions of tokens getting masked.\n            mask_ids: Tensor, or RaggedTensor if  `mask_selection_length` is\n                None. The original token ids at masked positions.\n            mask_weights: Tensor, or RaggedTensor if `mask_selection_length` is\n                None. `mask_weights` has the same shape as `mask_positions` and\n                `mask_ids`. Each element in `mask_weights` should be 0 or 1,\n                1 means the corresponding position in `mask_positions` is an\n                actual mask, 0 means it is a pad.\n\n    Examples:\n\n    Basic usage.\n    >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \\\n            vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0, \\\n            mask_selection_length=5)\n    >>> masker(tf.constant([1, 2, 3, 4, 5]))\n\n    Ragged Input:\n    >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \\\n            vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0, \\\n            mask_selection_length=5)\n    >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))\n    \"\"\"\n\n    def __init__(\n        self,\n        vocabulary_size,\n        mask_selection_rate,\n        mask_token_id,\n        mask_selection_length=None,\n        unselectable_token_ids=[0],\n        mask_token_rate=0.8,\n        random_token_rate=0.1,\n        **kwargs,\n    ):\n        super().__init__(**kwargs)\n        self.vocabulary_size = vocabulary_size\n        self.unselectable_token_ids = unselectable_token_ids\n        self.mask_selection_rate = mask_selection_rate\n        self.mask_selection_length = mask_selection_length\n        self.mask_token_rate = mask_token_rate\n        self.random_token_rate = random_token_rate\n\n        if mask_token_id >= vocabulary_size:\n            raise ValueError(\n                f\"Mask token id should be in range [0, vocabulary_size - 1], \"\n                f\"but received mask_token_id={mask_token_id}.\"\n            )\n        self.mask_token_id = mask_token_id\n\n        max_selections = self.mask_selection_length\n        if max_selections is None:\n            # Set a large number to remove the `max_selections_per_batch` cap.\n            max_selections = 2**31 - 1\n        self._random_selector = tf_text.RandomItemSelector(\n            max_selections_per_batch=max_selections,\n            selection_rate=self.mask_selection_rate,\n            unselectable_ids=self.unselectable_token_ids,\n        )\n        self._mask_values_chooser = tf_text.MaskValuesChooser(\n            self.vocabulary_size,\n            self.mask_token_id,\n            mask_token_rate=self.mask_token_rate,\n            random_token_rate=self.random_token_rate,\n        )\n\n    def call(self, inputs):\n        input_is_ragged = isinstance(inputs, tf.RaggedTensor)\n        input_is_1d = tf.rank(inputs) == 1\n        if input_is_1d:\n            # If inputs is of rank 1, we manually add the batch axis.\n            inputs = inputs[tf.newaxis, :]\n        if not input_is_ragged:\n            # `tf_text.mask_language_model` requires a ragged tensor, so\n            # convert dense to ragged.\n            inputs = tf.RaggedTensor.from_tensor(inputs)\n        (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(\n            inputs,\n            item_selector=self._random_selector,\n            mask_values_chooser=self._mask_values_chooser,\n        )\n\n        if not input_is_ragged:\n            # If we converted the input from dense to ragged, convert back.\n            tokens = tokens.to_tensor()\n\n        mask_weights = tf.ones_like(mask_positions, self.compute_dtype)\n        # If mask_selection_length is set, covert to raggeds to dense.\n        if self.mask_selection_length:\n            target_shape = tf.cast([-1, self.mask_selection_length], tf.int64)\n            mask_positions = mask_positions.to_tensor(shape=target_shape)\n            mask_ids = mask_ids.to_tensor(shape=target_shape)\n            mask_weights = mask_weights.to_tensor(shape=target_shape)\n\n        if input_is_1d:\n            # If inputs is 1D, we format the output to be 1D as well.\n            tokens = tf.squeeze(tokens, axis=0)\n            mask_positions = tf.squeeze(mask_positions, axis=0)\n            mask_ids = tf.squeeze(mask_ids, axis=0)\n            mask_weights = tf.squeeze(mask_weights, axis=0)\n\n        output_dict = {\n            \"tokens\": tokens,\n            \"mask_positions\": mask_positions,\n            \"mask_ids\": mask_ids,\n            \"mask_weights\": mask_weights,\n        }\n        return output_dict\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"vocabulary_size\": self.vocabulary_size,\n                \"mask_selection_rate\": self.mask_selection_rate,\n                \"mask_selection_length\": self.mask_selection_length,\n                \"unselectable_token_ids\": self.unselectable_token_ids,\n                \"mask_token_id\": self.mask_token_id,\n                \"mask_token_rate\": self.mask_token_rate,\n                \"random_token_rate\": self.random_token_rate,\n            }\n        )\n        return config\n", "path": "keras_nlp/layers/preprocessing/mlm_mask_generator.py"}], "after_files": [{"content": "# Copyright 2022 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tensorflow as tf\nimport tensorflow_text as tf_text\nfrom tensorflow import keras\n\n\nclass MLMMaskGenerator(keras.layers.Layer):\n    \"\"\"Layer that applies language model masking.\n\n    This layer is useful for preparing inputs for masked languaged modeling\n    (MLM) tasks. It follows the masking strategy described in the [original BERT\n    paper](https://arxiv.org/abs/1810.04805). Given tokenized text,\n    it randomly selects certain number of tokens for masking. Then for each\n    selected token, it has a chance (configurable) to be replaced by\n    \"mask token\" or random token, or stay unchanged.\n\n    Users should use this layer with `tf.data` to generate masks.\n\n    Args:\n        vocabulary_size: int, the size of the vocabulary.\n        mask_selection_rate: float, the probability of a token is selected for\n            masking.\n        mask_token_id: int. The id of mask token.\n        mask_selection_length: int, defaults to None. Maximum number of tokens\n            selected for  masking in each sequence. If set, the output\n            `mask_positions`, `mask_ids` and `mask_weights` will be padded\n            to dense tensors of length `mask_selection_length`,\n            otherwise the output will be a RaggedTensor.\n        unselectable_token_ids: A list of tokens, defaults to [0] (the default\n            `padding_token_id`).\n        mask_token_rate: float, defaults to 0.8. `mask_token_rate` must be\n            between 0 and 1 which indicates how often the mask_token is\n            substituted for tokens selected for masking.\n        random_token_rate: float, defaults to 0.1. `random_token_rate` must be\n            between 0 and 1 which indicates how often a random token is\n            substituted for tokens selected for masking. Default is 0.1.\n            Note: mask_token_rate + random_token_rate <= 1,  and for\n            (1 - mask_token_rate - random_token_rate), the token will not be\n            changed.\n\n    Input:\n        A 1D integer tensor of shape [sequence_length] or a 2D integer tensor\n        of shape [batch_size, sequence_length], or a 2D integer RaggedTensor.\n        Represents the sequence to mask.\n\n    Returns:\n        A Dict with 4 keys:\n            tokens: Tensor or RaggedTensor, has the same type and shape of\n                input. Sequence after getting masked.\n            mask_positions: Tensor, or RaggedTensor if `mask_selection_length`\n                is None. The positions of tokens getting masked.\n            mask_ids: Tensor, or RaggedTensor if  `mask_selection_length` is\n                None. The original token ids at masked positions.\n            mask_weights: Tensor, or RaggedTensor if `mask_selection_length` is\n                None. `mask_weights` has the same shape as `mask_positions` and\n                `mask_ids`. Each element in `mask_weights` should be 0 or 1,\n                1 means the corresponding position in `mask_positions` is an\n                actual mask, 0 means it is a pad.\n\n    Examples:\n\n    Basic usage.\n    >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \\\n            vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0, \\\n            mask_selection_length=5)\n    >>> masker(tf.constant([1, 2, 3, 4, 5]))\n\n    Ragged Input:\n    >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \\\n            vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0, \\\n            mask_selection_length=5)\n    >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))\n    \"\"\"\n\n    def __init__(\n        self,\n        vocabulary_size,\n        mask_selection_rate,\n        mask_token_id,\n        mask_selection_length=None,\n        unselectable_token_ids=[0],\n        mask_token_rate=0.8,\n        random_token_rate=0.1,\n        **kwargs,\n    ):\n        super().__init__(**kwargs)\n        self.vocabulary_size = vocabulary_size\n        self.unselectable_token_ids = unselectable_token_ids\n        self.mask_selection_rate = mask_selection_rate\n        self.mask_selection_length = mask_selection_length\n        self.mask_token_rate = mask_token_rate\n        self.random_token_rate = random_token_rate\n\n        if mask_token_id >= vocabulary_size:\n            raise ValueError(\n                f\"Mask token id should be in range [0, vocabulary_size - 1], \"\n                f\"but received mask_token_id={mask_token_id}.\"\n            )\n        self.mask_token_id = mask_token_id\n\n        max_selections = self.mask_selection_length\n        if max_selections is None:\n            # Set a large number to remove the `max_selections_per_batch` cap.\n            max_selections = 2**31 - 1\n        self._random_selector = tf_text.RandomItemSelector(\n            max_selections_per_batch=max_selections,\n            selection_rate=self.mask_selection_rate,\n            unselectable_ids=self.unselectable_token_ids,\n        )\n        self._mask_values_chooser = tf_text.MaskValuesChooser(\n            self.vocabulary_size,\n            self.mask_token_id,\n            mask_token_rate=self.mask_token_rate,\n            random_token_rate=self.random_token_rate,\n        )\n\n    def call(self, inputs):\n        input_is_ragged = isinstance(inputs, tf.RaggedTensor)\n        input_is_1d = inputs.shape.rank == 1\n        if input_is_1d:\n            # If inputs is of rank 1, we manually add the batch axis.\n            inputs = inputs[tf.newaxis, :]\n        if not input_is_ragged:\n            # `tf_text.mask_language_model` requires a ragged tensor, so\n            # convert dense to ragged.\n            inputs = tf.RaggedTensor.from_tensor(inputs)\n\n        (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(\n            inputs,\n            item_selector=self._random_selector,\n            mask_values_chooser=self._mask_values_chooser,\n        )\n\n        if not input_is_ragged:\n            # If we converted the input from dense to ragged, convert back.\n            tokens = tokens.to_tensor()\n\n        mask_weights = tf.ones_like(mask_positions, self.compute_dtype)\n        # If mask_selection_length is set, covert to raggeds to dense.\n        if self.mask_selection_length:\n            target_shape = tf.cast([-1, self.mask_selection_length], tf.int64)\n            mask_positions = mask_positions.to_tensor(shape=target_shape)\n            mask_ids = mask_ids.to_tensor(shape=target_shape)\n            mask_weights = mask_weights.to_tensor(shape=target_shape)\n\n        if input_is_1d:\n            # If inputs is 1D, we format the output to be 1D as well.\n            tokens = tf.squeeze(tokens, axis=0)\n            mask_positions = tf.squeeze(mask_positions, axis=0)\n            mask_ids = tf.squeeze(mask_ids, axis=0)\n            mask_weights = tf.squeeze(mask_weights, axis=0)\n\n        output_dict = {\n            \"tokens\": tokens,\n            \"mask_positions\": mask_positions,\n            \"mask_ids\": mask_ids,\n            \"mask_weights\": mask_weights,\n        }\n        return output_dict\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"vocabulary_size\": self.vocabulary_size,\n                \"mask_selection_rate\": self.mask_selection_rate,\n                \"mask_selection_length\": self.mask_selection_length,\n                \"unselectable_token_ids\": self.unselectable_token_ids,\n                \"mask_token_id\": self.mask_token_id,\n                \"mask_token_rate\": self.mask_token_rate,\n                \"random_token_rate\": self.random_token_rate,\n            }\n        )\n        return config\n", "path": "keras_nlp/layers/preprocessing/mlm_mask_generator.py"}]}