Repository: keras-team/keras-nlp
Pull request #836 (2023-03-11 07:39:24)

Add compilation defaults for the BertMaskedLM task model 
add fixes at https://github.com/keras-team/keras-nlp/issues/830

soma2000-lang: @ADITYADAS1999 I had already opened a pr 1 hour ago and clearly mentioned it in the isssue.So why did you open the same pr again?
ADITYADAS1999: > @ADITYADAS1999 I had already opened a pr 1 hour ago and clearly mentioned it in the isssue.So why did you open the same pr again?

maybe a typo ! Now it should fine 👍🏻
mattdangerw: Code looks good. Have we tried this out with a few different presets to make sure we are converging generally?
ADITYADAS1999: 
> Code looks good. Have we tried this out with a few different presets to make sure we are converging generally?

that's a bit tricky, in this case we have to working on some cases on convergence 

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index dda0965b..d9c9360c 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -83,6 +83,7 @@ from keras_nlp.models.roberta.roberta_masked_lm_preprocessor import (
 )
 from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tokenizer import RobertaTokenizer
+from keras_nlp.models.t5.t5_tokenizer import T5Tokenizer
 from keras_nlp.models.whisper.whisper_backbone import WhisperBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_classifier import (
diff --git a/keras_nlp/models/bert/bert_masked_lm.py b/keras_nlp/models/bert/bert_masked_lm.py
index d44abf18..a79be436 100644
--- a/keras_nlp/models/bert/bert_masked_lm.py
+++ b/keras_nlp/models/bert/bert_masked_lm.py
@@ -26,6 +26,7 @@ from keras_nlp.models.bert.bert_masked_lm_preprocessor import (
 )
 from keras_nlp.models.bert.bert_presets import backbone_presets
 from keras_nlp.models.task import Task
+from keras_nlp.utils.keras_utils import is_xla_compatible
 from keras_nlp.utils.python_utils import classproperty
 
 
@@ -139,6 +140,12 @@ class BertMaskedLM(Task):
         # All references to `self` below this line
         self.backbone = backbone
         self.preprocessor = preprocessor
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            weighted_metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
 
     @classproperty
     def backbone_cls(cls):
diff --git a/keras_nlp/models/t5/t5_tokenizer.py b/keras_nlp/models/t5/t5_tokenizer.py
new file mode 100644
index 00000000..d6c058c5
--- /dev/null
+++ b/keras_nlp/models/t5/t5_tokenizer.py
@@ -0,0 +1,76 @@
+# Copyright 2023 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.
+
+"""T5 tokenizer."""
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
+
+
+@keras_nlp_export("keras_nlp.models.T5Tokenizer")
+class T5Tokenizer(SentencePieceTokenizer):
+    """T5 tokenizer layer based on SentencePiece.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.SentencePieceTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by
+    T5 models and provides a `from_preset()` method to automatically
+    download a matching vocabulary for a T5 preset.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        proto: Either a `string` path to a SentencePiece proto file, or a
+            `bytes` object with a serialized SentencePiece proto. See the
+            [SentencePiece repository](https://github.com/google/sentencepiece)
+            for more details on the format.
+
+    Examples:
+
+    ```python
+    tokenizer = keras_nlp.models.T5Tokenizer(proto="model.spm")
+
+    # Batched inputs.
+    tokenizer(["the quick brown fox", "the earth is round"])
+
+    # Unbatched inputs.
+    tokenizer("the quick brown fox")
+
+    # Detokenization.
+    tokenizer.detokenize(tf.constant([[2, 14, 2231, 886, 2385, 3]]))
+    ```
+    """
+
+    def __init__(self, proto, **kwargs):
+        super().__init__(proto=proto, **kwargs)
+
+        # Check for necessary special tokens.
+        end_token = "
"
+        pad_token = "<pad>"
+        for token in [pad_token]:
+            if token not in self.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        self.pad_token_id = self.token_to_id(pad_token)
+        self.end_token_id = self.token_to_id(end_token)
+        # T5 uses the same start token as end token, i.e., "<\s>".
+        self.start_token_id = self.end_token_id


Test patch:
diff --git a/keras_nlp/models/albert/albert_backbone_test.py b/keras_nlp/models/albert/albert_backbone_test.py
index 3e8115a9..89543c0c 100644
--- a/keras_nlp/models/albert/albert_backbone_test.py
+++ b/keras_nlp/models/albert/albert_backbone_test.py
@@ -15,6 +15,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -118,3 +119,34 @@ class AlbertBackboneTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllClose(
             model_output["pooled_output"], restored_output["pooled_output"]
         )
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class AlbertBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.model = AlbertBackbone(
+                vocabulary_size=1000,
+                num_layers=2,
+                num_heads=2,
+                num_groups=1,
+                num_inner_repetitions=1,
+                embedding_dim=16,
+                hidden_dim=64,
+                intermediate_dim=128,
+                max_sequence_length=128,
+            )
+
+        self.input_batch = {
+            "token_ids": tf.ones((8, 128), dtype="int32"),
+            "segment_ids": tf.ones((8, 128), dtype="int32"),
+            "padding_mask": tf.ones((8, 128), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.model.compile()
+        self.model.predict(self.input_dataset)
diff --git a/keras_nlp/models/bert/bert_masked_lm_test.py b/keras_nlp/models/bert/bert_masked_lm_test.py
index 0fa10ce3..921bcedf 100644
--- a/keras_nlp/models/bert/bert_masked_lm_test.py
+++ b/keras_nlp/models/bert/bert_masked_lm_test.py
@@ -79,6 +79,9 @@ class BertMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
     def test_valid_call_masked_lm(self):
         self.masked_lm(self.preprocessed_batch)
 
+    def test_bert_masked_lm_fit_default_compile(self):
+        self.masked_lm.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
diff --git a/keras_nlp/models/f_net/f_net_backbone_test.py b/keras_nlp/models/f_net/f_net_backbone_test.py
index 9878569c..08720afa 100644
--- a/keras_nlp/models/f_net/f_net_backbone_test.py
+++ b/keras_nlp/models/f_net/f_net_backbone_test.py
@@ -15,6 +15,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -96,3 +97,29 @@ class FNetBackboneTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllClose(
             model_output["pooled_output"], restored_output["pooled_output"]
         )
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class FNetBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.model = FNetBackbone(
+                vocabulary_size=1000,
+                num_layers=2,
+                hidden_dim=16,
+                intermediate_dim=32,
+                max_sequence_length=128,
+                num_segments=4,
+            )
+        self.input_batch = {
+            "token_ids": tf.ones((8, 128), dtype="int32"),
+            "segment_ids": tf.ones((8, 128), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.model.compile()
+        self.model.predict(self.input_dataset)
diff --git a/keras_nlp/models/gpt2/gpt2_backbone_test.py b/keras_nlp/models/gpt2/gpt2_backbone_test.py
index a4c562a7..4115a982 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone_test.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone_test.py
@@ -15,6 +15,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -94,3 +95,29 @@ class GPT2Test(tf.test.TestCase, parameterized.TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class GPT2BackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.model = GPT2Backbone(
+                vocabulary_size=1000,
+                num_layers=2,
+                num_heads=2,
+                hidden_dim=64,
+                intermediate_dim=128,
+                max_sequence_length=128,
+            )
+        self.input_batch = {
+            "token_ids": tf.ones((8, 128), dtype="int32"),
+            "padding_mask": tf.ones((8, 128), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.model.compile()
+        self.model.predict(self.input_dataset)
diff --git a/keras_nlp/models/opt/opt_backbone_test.py b/keras_nlp/models/opt/opt_backbone_test.py
index 4283dc13..e18b217b 100644
--- a/keras_nlp/models/opt/opt_backbone_test.py
+++ b/keras_nlp/models/opt/opt_backbone_test.py
@@ -15,6 +15,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -94,3 +95,29 @@ class OPTTest(tf.test.TestCase, parameterized.TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class OPTBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.model = OPTBackbone(
+                vocabulary_size=1000,
+                num_layers=2,
+                num_heads=2,
+                hidden_dim=32,
+                intermediate_dim=128,
+                max_sequence_length=128,
+            )
+        self.input_batch = {
+            "token_ids": tf.ones((8, 128), dtype="int32"),
+            "padding_mask": tf.ones((8, 128), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.model.compile()
+        self.model.predict(self.input_dataset)
diff --git a/keras_nlp/models/roberta/roberta_backbone_test.py b/keras_nlp/models/roberta/roberta_backbone_test.py
index 99104066..2247f065 100644
--- a/keras_nlp/models/roberta/roberta_backbone_test.py
+++ b/keras_nlp/models/roberta/roberta_backbone_test.py
@@ -16,6 +16,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -23,7 +24,7 @@ from tensorflow import keras
 from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 
 
-class RobertaTest(tf.test.TestCase, parameterized.TestCase):
+class RobertaBackboneTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.model = RobertaBackbone(
             vocabulary_size=1000,
@@ -99,3 +100,29 @@ class RobertaTest(tf.test.TestCase, parameterized.TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class RobertaBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.model = RobertaBackbone(
+                vocabulary_size=1000,
+                num_layers=2,
+                num_heads=2,
+                hidden_dim=64,
+                intermediate_dim=128,
+                max_sequence_length=128,
+            )
+        self.input_batch = {
+            "token_ids": tf.ones((8, 128), dtype="int32"),
+            "padding_mask": tf.ones((8, 128), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.model.compile()
+        self.model.predict(self.input_dataset)
diff --git a/keras_nlp/models/t5/t5_tokenizer_test.py b/keras_nlp/models/t5/t5_tokenizer_test.py
new file mode 100644
index 00000000..dfaa700d
--- /dev/null
+++ b/keras_nlp/models/t5/t5_tokenizer_test.py
@@ -0,0 +1,102 @@
+# Copyright 2023 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.
+
+"""Tests for T5 tokenizer."""
+
+import io
+import os
+
+import sentencepiece
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.t5.t5_tokenizer import T5Tokenizer
+
+
+class T5TokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=11,
+            model_type="WORD",
+            bos_id=-1,
+            pad_id=0,
+            eos_id=1,
+            unk_id=2,
+            pad_piece="<pad>",
+            eos_piece="
",
+            unk_piece="<unk>",
+            user_defined_symbols="[MASK]",
+        )
+        self.proto = bytes_io.getvalue()
+
+        self.tokenizer = T5Tokenizer(proto=self.proto)
+
+    def test_tokenize(self):
+        input_data = "the quick brown fox"
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [4, 9, 5, 7])
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant(["the quick brown fox", "the earth is round"])
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [[4, 9, 5, 7], [4, 6, 8, 10]])
+
+    def test_detokenize(self):
+        input_data = tf.constant([[4, 9, 5, 7]])
+        output = self.tokenizer.detokenize(input_data)
+        self.assertEqual(output, tf.constant(["the quick brown fox"]))
+
+    def test_vocabulary_size(self):
+        tokenizer = T5Tokenizer(proto=self.proto)
+        self.assertEqual(tokenizer.vocabulary_size(), 11)
+
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            T5Tokenizer(proto=bytes_io.getvalue())
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        input_data = tf.constant(["the quick brown fox"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
index d9554a9a..2d1d327d 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
@@ -16,6 +16,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -23,7 +24,7 @@ from tensorflow import keras
 from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 
 
-class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
+class XLMRobertaBackboneTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.model = XLMRobertaBackbone(
             vocabulary_size=1000,
@@ -99,3 +100,29 @@ class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class XLMRobertaBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.model = XLMRobertaBackbone(
+                vocabulary_size=1000,
+                num_layers=2,
+                num_heads=2,
+                hidden_dim=64,
+                intermediate_dim=128,
+                max_sequence_length=128,
+            )
+        self.input_batch = {
+            "token_ids": tf.ones((8, 128), dtype="int32"),
+            "padding_mask": tf.ones((8, 128), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.model.compile()
+        self.model.predict(self.input_dataset)