Add a common test case

keras_nlp/layers/cached_multi_head_attention_test.py

@@ -20,9 +20,10 @@
 from keras_nlp.layers.cached_multi_head_attention import (
     CachedMultiHeadAttention,
 )
+from keras_nlp.tests.test_case import TestCase
 
 
-class CachedMultiHeadAttentionTest(tf.test.TestCase, parameterized.TestCase):
+class CachedMultiHeadAttentionTest(TestCase):
     def test_valid_call(self):
         layer = CachedMultiHeadAttention(num_heads=2, key_dim=4)
         x = tf.random.uniform(shape=[2, 2, 8])

keras_nlp/layers/f_net_encoder_test.py

@@ -20,9 +20,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers import f_net_encoder
+from keras_nlp.tests.test_case import TestCase
 
 
-class FNetEncoderTest(tf.test.TestCase, parameterized.TestCase):
+class FNetEncoderTest(TestCase):
     def test_valid_call(self):
         encoder = f_net_encoder.FNetEncoder(intermediate_dim=4)
         model = keras.Sequential(

keras_nlp/layers/masked_lm_head_test.py

@@ -20,9 +20,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers import masked_lm_head
+from keras_nlp.tests.test_case import TestCase
 
 
-class MaskedLMHeadTest(tf.test.TestCase, parameterized.TestCase):
+class MaskedLMHeadTest(TestCase):
     def test_valid_call(self):
         head = masked_lm_head.MaskedLMHead(
             vocabulary_size=100,

keras_nlp/layers/masked_lm_mask_generator_test.py

@@ -15,9 +15,10 @@
 import tensorflow as tf
 
 from keras_nlp.layers.masked_lm_mask_generator import MaskedLMMaskGenerator
+from keras_nlp.tests.test_case import TestCase
 
 
-class MaskedLMMaskGeneratorTest(tf.test.TestCase):
+class MaskedLMMaskGeneratorTest(TestCase):
     def setUp(self):
         super().setUp()
         self.VOCAB = [

keras_nlp/layers/multi_segment_packer_test.py

@@ -20,9 +20,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
+from keras_nlp.tests.test_case import TestCase
 
 
-class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
+class MultiSegmentPackerTest(TestCase):
     def test_trim_single_input_ints(self):
         input_data = tf.range(3, 10)
         packer = MultiSegmentPacker(8, start_value=1, end_value=2)

keras_nlp/layers/position_embedding_test.py

@@ -20,6 +20,7 @@
 from tensorflow import keras
 
 from keras_nlp.layers import position_embedding
+from keras_nlp.tests.test_case import TestCase
 
 
 def custom_init(shape, dtype=None):
@@ -29,7 +30,7 @@ def custom_init(shape, dtype=None):
     return tf.reshape(tf.range(count, dtype=dtype), shape)
 
 
-class PositionEmbeddingLayerTest(tf.test.TestCase, parameterized.TestCase):
+class PositionEmbeddingLayerTest(TestCase):
     def test_static_layer_output_shape(self):
         # Create a 3-dimensional input (the first dimension is implicit).
         sequence_length = 21

keras_nlp/layers/random_deletion_test.py

@@ -17,9 +17,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers.random_deletion import RandomDeletion
+from keras_nlp.tests.test_case import TestCase
 
 
-class RandomDeletionTest(tf.test.TestCase):
+class RandomDeletionTest(TestCase):
     def test_shape_and_output_from_word_deletion(self):
         keras.utils.set_random_seed(1337)
         inputs = ["Hey I like", "Keras and Tensorflow"]

keras_nlp/layers/random_swap_test.py

@@ -17,9 +17,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers.random_swap import RandomSwap
+from keras_nlp.tests.test_case import TestCase
 
 
-class RandomSwapTest(tf.test.TestCase):
+class RandomSwapTest(TestCase):
     def test_shape_and_output_from_word_swap(self):
         keras.utils.set_random_seed(1337)
         inputs = ["Hey I like", "Keras and Tensorflow"]

keras_nlp/layers/sine_position_encoding_test.py

@@ -17,9 +17,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers import sine_position_encoding
+from keras_nlp.tests.test_case import TestCase
 
 
-class SinePositionEncodingTest(tf.test.TestCase):
+class SinePositionEncodingTest(TestCase):
     def test_valid_call(self):
         pos_encoding = sine_position_encoding.SinePositionEncoding()
         model = keras.Sequential(

keras_nlp/layers/start_end_packer_test.py

@@ -13,15 +13,14 @@
 # limitations under the License.
 
 """Tests for Start End Packer layer."""
-
-
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.layers.start_end_packer import StartEndPacker
+from keras_nlp.tests.test_case import TestCase
 
 
-class StartEndPackerTest(tf.test.TestCase):
+class StartEndPackerTest(TestCase):
     def test_dense_input(self):
         input_data = tf.constant([5, 6, 7])
         start_end_packer = StartEndPacker(sequence_length=5)

keras_nlp/layers/token_and_position_embedding_test.py

@@ -22,9 +22,10 @@
 from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
+from keras_nlp.tests.test_case import TestCase
 
 
-class TokenAndPositionEmbeddingTest(tf.test.TestCase, parameterized.TestCase):
+class TokenAndPositionEmbeddingTest(TestCase):
     def test_get_config_and_from_config(self):
         token_and_position_embed = TokenAndPositionEmbedding(
             vocabulary_size=5,

keras_nlp/layers/transformer_decoder_test.py

@@ -21,9 +21,10 @@
 from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
 
 from keras_nlp.layers import transformer_decoder
+from keras_nlp.tests.test_case import TestCase
 
 
-class TransformerDecoderTest(tf.test.TestCase, parameterized.TestCase):
+class TransformerDecoderTest(TestCase):
     @parameterized.named_parameters(
         ("without_norm_first", False),
         ("with_norm_first", True),

keras_nlp/layers/transformer_encoder_test.py

@@ -20,9 +20,10 @@
 from tensorflow import keras
 
 from keras_nlp.layers import transformer_encoder
+from keras_nlp.tests.test_case import TestCase
 
 
-class TransformerEncoderTest(tf.test.TestCase, parameterized.TestCase):
+class TransformerEncoderTest(TestCase):
     @parameterized.named_parameters(
         ("without_norm_first", False),
         ("with_norm_first", True),

keras_nlp/layers/transformer_layer_utils_test.py

@@ -15,9 +15,10 @@
 import tensorflow as tf
 
 import keras_nlp.layers.transformer_layer_utils as utils
+from keras_nlp.tests.test_case import TestCase
 
 
-class TransformerEncoderTest(tf.test.TestCase):
+class TransformerEncoderTest(TestCase):
     def test_compute_causal_mask(self):
         mask = utils.compute_causal_mask(1, 2, 2)
         self.assertTrue((mask.numpy() == [[1, 0], [1, 1]]).all())

keras_nlp/metrics/bleu_test.py

@@ -13,15 +13,15 @@
 # limitations under the License.
 
 """Tests for Bleu."""
-
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.metrics.bleu import Bleu
+from keras_nlp.tests.test_case import TestCase
 from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
 
 
-class BleuTest(tf.test.TestCase):
+class BleuTest(TestCase):
     def test_initialization(self):
         bleu = Bleu()
         result = bleu.result()

keras_nlp/metrics/edit_distance_test.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 
 """Tests for EditDistance."""
-
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.metrics.edit_distance import EditDistance
+from keras_nlp.tests.test_case import TestCase
 
 
-class EditDistanceTest(tf.test.TestCase):
+class EditDistanceTest(TestCase):
     def test_initialization(self):
         edit_distance = EditDistance()
         result = edit_distance.result()

keras_nlp/metrics/perplexity_test.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 """Tests for Perplexity."""
-
 import tensorflow as tf
 
 from keras_nlp.metrics.perplexity import Perplexity
+from keras_nlp.tests.test_case import TestCase
 
 
-class PerplexityTest(tf.test.TestCase):
+class PerplexityTest(TestCase):
     def test_vars_after_initializing_class(self):
         perplexity = Perplexity()
         self.assertEqual(perplexity.result(), 0.0)

keras_nlp/metrics/rouge_l_test.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 """Tests for RougeL."""
-
 import tensorflow as tf
 
 from keras_nlp.metrics.rouge_l import RougeL
+from keras_nlp.tests.test_case import TestCase
 
 
-class RougeLTest(tf.test.TestCase):
+class RougeLTest(TestCase):
     def test_initialization(self):
         rouge = RougeL()
         result = rouge.result()

keras_nlp/metrics/rouge_n_test.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 
 """Tests for RougeN."""
-
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.metrics.rouge_n import RougeN
+from keras_nlp.tests.test_case import TestCase
 
 
-class RougeNTest(tf.test.TestCase):
+class RougeNTest(TestCase):
     def test_initialization(self):
         rouge = RougeN()
         result = rouge.result()

keras_nlp/models/albert/albert_backbone_test.py

@@ -21,9 +21,10 @@
 from tensorflow import keras
 
 from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.tests.test_case import TestCase
 
 
-class AlbertBackboneTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertBackboneTest(TestCase):
     def setUp(self):
         self.backbone = AlbertBackbone(
             vocabulary_size=10,
@@ -111,7 +112,7 @@ def test_saved_model(self, save_format, filename):
 
 @pytest.mark.tpu
 @pytest.mark.usefixtures("tpu_test_class")
-class AlbertBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertBackboneTPUTest(TestCase):
     def setUp(self):
         with self.tpu_strategy.scope():
             self.backbone = AlbertBackbone(

keras_nlp/models/albert/albert_classifier_test.py

@@ -26,9 +26,10 @@
 from keras_nlp.models.albert.albert_classifier import AlbertClassifier
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_nlp.tests.test_case import TestCase
 
 
-class AlbertClassifierTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertClassifierTest(TestCase):
     def setUp(self):
         # Setup model
 

keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py

@@ -26,9 +26,10 @@
     AlbertMaskedLMPreprocessor,
 )
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_nlp.tests.test_case import TestCase
 
 
-class AlbertMaskedLMPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertMaskedLMPreprocessorTest(TestCase):
     def setUp(self):
         vocab_data = tf.data.Dataset.from_tensor_slices(
             ["the quick brown fox", "the earth is round"]

keras_nlp/models/albert/albert_masked_lm_test.py

@@ -28,9 +28,10 @@
     AlbertMaskedLMPreprocessor,
 )
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_nlp.tests.test_case import TestCase
 
 
-class AlbertMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertMaskedLMTest(TestCase):
     def setUp(self):
         # Setup model.
         vocab_data = tf.data.Dataset.from_tensor_slices(

keras_nlp/models/albert/albert_preprocessor_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Tests for ALBERT preprocessor layer."""
-
 import io
 import os
 
@@ -25,9 +24,10 @@
 
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_nlp.tests.test_case import TestCase
 
 
-class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertPreprocessorTest(TestCase):
     def setUp(self):
         bytes_io = io.BytesIO()
         vocab_data = tf.data.Dataset.from_tensor_slices(

keras_nlp/models/albert/albert_presets_test.py

@@ -21,10 +21,11 @@
 from keras_nlp.models.albert.albert_classifier import AlbertClassifier
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_nlp.tests.test_case import TestCase
 
 
 @pytest.mark.large
-class AlbertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertPresetSmokeTest(TestCase):
     """
     A smoke test for ALBERT presets we run continuously.
     This only tests the smallest weights we have available. Run with:
@@ -121,7 +122,7 @@ def test_unknown_preset_error(self, cls, kwargs):
 
 
 @pytest.mark.extra_large
-class AlbertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertPresetFullTest(TestCase):
     """
     Test the full enumeration of our preset.
     This tests every ALBERT preset and is only run manually.

keras_nlp/models/albert/albert_tokenizer_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Tests for ALBERT tokenizer."""
-
 import io
 import os
 
@@ -24,9 +23,10 @@
 from tensorflow import keras
 
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_nlp.tests.test_case import TestCase
 
 
-class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
+class AlbertTokenizerTest(TestCase):
     def setUp(self):
         bytes_io = io.BytesIO()
         vocab_data = tf.data.Dataset.from_tensor_slices(