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 2d1d327de2..d91330a0c0 100644 --- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py +++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py @@ -26,72 +26,63 @@ class XLMRobertaBackboneTest(tf.test.TestCase, parameterized.TestCase): def setUp(self): - self.model = XLMRobertaBackbone( - vocabulary_size=1000, + self.backbone = XLMRobertaBackbone( + vocabulary_size=10, num_layers=2, num_heads=2, - hidden_dim=64, - intermediate_dim=128, - max_sequence_length=128, + hidden_dim=2, + intermediate_dim=4, + max_sequence_length=5, ) - self.batch_size = 8 self.input_batch = { - "token_ids": tf.ones( - (self.batch_size, self.model.max_sequence_length), dtype="int32" - ), - "padding_mask": tf.ones( - (self.batch_size, self.model.max_sequence_length), dtype="int32" - ), + "token_ids": tf.ones((2, 5), dtype="int32"), + "padding_mask": tf.ones((2, 5), dtype="int32"), } - self.input_dataset = tf.data.Dataset.from_tensor_slices( self.input_batch ).batch(2) def test_valid_call_xlm_roberta(self): - self.model(self.input_batch) - - # Check default name passed through - self.assertRegexpMatches(self.model.name, "xlm_roberta_backbone") + self.backbone(self.input_batch) - @parameterized.named_parameters( - ("jit_compile_false", False), ("jit_compile_true", True) - ) - def test_xlm_roberta_compile(self, jit_compile): - self.model.compile(jit_compile=jit_compile) - self.model.predict(self.input_batch) + def test_token_embedding(self): + output = self.backbone.token_embedding(self.input_batch["token_ids"]) + self.assertEqual(output.shape, (2, 5, 2)) - @parameterized.named_parameters( - ("jit_compile_false", False), ("jit_compile_true", True) - ) - def test_xlm_roberta_compile_batched_ds(self, jit_compile): - self.model.compile(jit_compile=jit_compile) - self.model.predict(self.input_dataset) + def test_name(self): + # Check default name passed through + self.assertRegexpMatches(self.backbone.name, "xlm_roberta_backbone") def test_variable_sequence_length_call_xlm_roberta(self): - for seq_length in (25, 50, 75): + for seq_length in (2, 3, 4): input_data = { - "token_ids": tf.ones( - (self.batch_size, seq_length), dtype="int32" - ), - "padding_mask": tf.ones( - (self.batch_size, seq_length), dtype="int32" - ), + "token_ids": tf.ones((2, seq_length), dtype="int32"), + "padding_mask": tf.ones((2, seq_length), dtype="int32"), } - output = self.model(input_data) + output = self.backbone(input_data) self.assertAllEqual( - tf.shape(output), - [self.batch_size, seq_length, self.model.hidden_dim], + tf.shape(output), [2, seq_length, self.backbone.hidden_dim] ) + def test_predict(self): + self.backbone.predict(self.input_batch) + self.backbone.predict(self.input_dataset) + + def test_serialization(self): + new_backbone = keras.utils.deserialize_keras_object( + keras.utils.serialize_keras_object(self.backbone) + ) + self.assertEqual(new_backbone.get_config(), self.backbone.get_config()) + @parameterized.named_parameters( ("tf_format", "tf", "model"), ("keras_format", "keras_v3", "model.keras"), ) + @pytest.mark.large # Saving is slow, so mark these large. def test_saved_model(self, save_format, filename): - model_output = self.model(self.input_batch) + model_output = self.backbone(self.input_batch) save_path = os.path.join(self.get_temp_dir(), filename) - self.model.save(save_path, save_format=save_format) + self.backbone.save(save_path, save_format=save_format) restored_model = keras.models.load_model(save_path) # Check we got the real object back. @@ -107,7 +98,7 @@ def test_saved_model(self, save_format, filename): class XLMRobertaBackboneTPUTest(tf.test.TestCase, parameterized.TestCase): def setUp(self): with self.tpu_strategy.scope(): - self.model = XLMRobertaBackbone( + self.backbone = XLMRobertaBackbone( vocabulary_size=1000, num_layers=2, num_heads=2, @@ -124,5 +115,5 @@ def setUp(self): ).batch(2) def test_predict(self): - self.model.compile() - self.model.predict(self.input_dataset) + self.backbone.compile() + self.backbone.predict(self.input_dataset) diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py index c8efb4fbeb..bf47672032 100644 --- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py +++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py @@ -16,6 +16,7 @@ import io import os +import pytest import sentencepiece import tensorflow as tf from absl.testing import parameterized @@ -50,15 +51,15 @@ def setUp(self): ) self.preprocessor = XLMRobertaPreprocessor( tokenizer=XLMRobertaTokenizer(proto=bytes_io.getvalue()), - sequence_length=12, + sequence_length=5, ) self.backbone = XLMRobertaBackbone( - vocabulary_size=1000, + vocabulary_size=10, num_layers=2, num_heads=2, - hidden_dim=64, - intermediate_dim=128, - max_sequence_length=128, + hidden_dim=2, + intermediate_dim=4, + max_sequence_length=self.preprocessor.packer.sequence_length, ) self.classifier = XLMRobertaClassifier( self.backbone, @@ -75,60 +76,48 @@ def setUp(self): [ "the quick brown fox.", "the slow brown fox.", - "the earth is round", - "the earth is spherical", ] ) self.preprocessed_batch = self.preprocessor(self.raw_batch) self.raw_dataset = tf.data.Dataset.from_tensor_slices( - (self.raw_batch, tf.ones((4,))) + (self.raw_batch, tf.ones((2,))) ).batch(2) self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor) def test_valid_call_classifier(self): self.classifier(self.preprocessed_batch) - @parameterized.named_parameters( - ("jit_compile_false", False), ("jit_compile_true", True) - ) - def test_classifier_predict(self, jit_compile): - self.classifier.compile(jit_compile=jit_compile) + def test_classifier_predict(self): self.classifier.predict(self.raw_batch) + self.classifier.preprocessor = None + self.classifier.predict(self.preprocessed_batch) - @parameterized.named_parameters( - ("jit_compile_false", False), ("jit_compile_true", True) - ) - def test_classifier_predict_no_preprocessing(self, jit_compile): - self.classifier_no_preprocessing.compile(jit_compile=jit_compile) - self.classifier_no_preprocessing.predict(self.preprocessed_batch) + def test_classifier_fit(self): + self.classifier.fit(self.raw_dataset) + self.classifier.preprocessor = None + self.classifier.fit(self.preprocessed_dataset) - @parameterized.named_parameters( - ("jit_compile_false", False), ("jit_compile_true", True) - ) - def test_classifier_fit(self, jit_compile): + def test_classifier_fit_no_xla(self): + self.classifier.preprocessor = None self.classifier.compile( - loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), - jit_compile=jit_compile, + loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False), + jit_compile=False, ) - self.classifier.fit(self.raw_dataset) + self.classifier.fit(self.preprocessed_dataset) - @parameterized.named_parameters( - ("jit_compile_false", False), ("jit_compile_true", True) - ) - def test_classifier_fit_no_preprocessing(self, jit_compile): - self.classifier_no_preprocessing.compile( - loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), - jit_compile=jit_compile, + def test_serialization(self): + config = keras.utils.serialize_keras_object(self.classifier) + new_classifier = keras.utils.deserialize_keras_object(config) + self.assertEqual( + new_classifier.get_config(), + self.classifier.get_config(), ) - self.classifier_no_preprocessing.fit(self.preprocessed_dataset) - - def test_xlmroberta_classifier_fit_default_compile(self): - self.classifier.fit(self.raw_dataset) @parameterized.named_parameters( ("tf_format", "tf", "model"), ("keras_format", "keras_v3", "model.keras"), ) + @pytest.mark.large # Saving is slow, so mark these large. def test_saving_model(self, save_format, filename): model_output = self.classifier.predict(self.raw_batch) save_path = os.path.join(self.get_temp_dir(), filename) diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py index 111be152ed..d983f8bcec 100644 --- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py +++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py @@ -17,6 +17,7 @@ import io import os +import pytest import sentencepiece import tensorflow as tf from absl.testing import parameterized @@ -132,10 +133,19 @@ def test_errors_for_2d_list_input(self): with self.assertRaises(ValueError): self.preprocessor(ambiguous_input) + def test_serialization(self): + config = keras.utils.serialize_keras_object(self.preprocessor) + new_preprocessor = keras.utils.deserialize_keras_object(config) + self.assertEqual( + new_preprocessor.get_config(), + self.preprocessor.get_config(), + ) + @parameterized.named_parameters( ("tf_format", "tf", "model"), ("keras_format", "keras_v3", "model.keras"), ) + @pytest.mark.large # Saving is slow, so mark these large. def test_saved_model(self, save_format, filename): input_data = tf.constant(["the quick brown fox"]) inputs = keras.Input(dtype="string", shape=()) diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py index b62034eb49..42e1fe1877 100644 --- a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py +++ b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py @@ -17,6 +17,7 @@ import io import os +import pytest import sentencepiece import tensorflow as tf from absl.testing import parameterized @@ -96,10 +97,19 @@ def test_token_to_id(self): self.assertEqual(self.tokenizer.token_to_id("▁the"), 4) self.assertEqual(self.tokenizer.token_to_id("▁round"), 10) + def test_serialization(self): + config = keras.utils.serialize_keras_object(self.tokenizer) + new_tokenizer = keras.utils.deserialize_keras_object(config) + self.assertEqual( + new_tokenizer.get_config(), + self.tokenizer.get_config(), + ) + @parameterized.named_parameters( ("tf_format", "tf", "model"), ("keras_format", "keras_v3", "model.keras"), ) + @pytest.mark.large # Saving is slow, so mark these large. def test_saved_model(self, save_format, filename): input_data = tf.constant(["the quick brown fox"])