From a6b43a1abde5d75044bdbb8c46662fe56fcca49a Mon Sep 17 00:00:00 2001
From: Matt Watson <1389937+mattdangerw@users.noreply.github.com>
Date: Mon, 10 Jul 2023 16:21:03 -0700
Subject: [PATCH] Unexport models that are not in the 0.6 release (#1125)

This change will not affect the master branch, where they will
continued to be exported. For the release branch we will not yet
export or document these models.

Though all of these models are coming soon!
---
 keras_nlp/models/bart/bart_backbone.py                        | 3 +--
 keras_nlp/models/bart/bart_preprocessor.py                    | 4 ++--
 keras_nlp/models/bart/bart_seq_2_seq_lm.py                    | 3 +--
 keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py       | 4 ++--
 keras_nlp/models/bart/bart_tokenizer.py                       | 4 ++--
 keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py              | 3 +--
 .../models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py      | 4 ++--
 keras_nlp/models/gpt_neo_x/gpt_neo_x_preprocessor.py          | 4 ++--
 keras_nlp/models/gpt_neo_x/gpt_neo_x_tokenizer.py             | 4 ++--
 keras_nlp/models/t5/t5_backbone.py                            | 3 +--
 keras_nlp/models/t5/t5_tokenizer.py                           | 4 ++--
 keras_nlp/models/whisper/whisper_audio_feature_extractor.py   | 3 +--
 keras_nlp/models/whisper/whisper_backbone.py                  | 3 +--
 keras_nlp/models/whisper/whisper_preprocessor.py              | 3 +--
 keras_nlp/models/whisper/whisper_tokenizer.py                 | 4 ++--
 15 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
index 594f2467a0..73fd08e2ed 100644
--- a/keras_nlp/models/bart/bart_backbone.py
+++ b/keras_nlp/models/bart/bart_backbone.py
@@ -16,7 +16,6 @@
 
 import copy
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
 from keras_nlp.layers.modeling.transformer_decoder import TransformerDecoder
@@ -30,7 +29,7 @@ def bart_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras_nlp_export("keras_nlp.models.BartBackbone")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class BartBackbone(Backbone):
     """BART encoder-decoder network.
 
diff --git a/keras_nlp/models/bart/bart_preprocessor.py b/keras_nlp/models/bart/bart_preprocessor.py
index 04626ddb9c..db0fb51ac4 100644
--- a/keras_nlp/models/bart/bart_preprocessor.py
+++ b/keras_nlp/models/bart/bart_preprocessor.py
@@ -15,7 +15,7 @@
 
 import copy
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.models.bart.bart_presets import backbone_presets
 from keras_nlp.models.bart.bart_tokenizer import BartTokenizer
@@ -27,7 +27,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras_nlp_export("keras_nlp.models.BartPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class BartPreprocessor(Preprocessor):
     """A BART preprocessing layer which tokenizes and packs inputs.
 
diff --git a/keras_nlp/models/bart/bart_seq_2_seq_lm.py b/keras_nlp/models/bart/bart_seq_2_seq_lm.py
index 17c02e8e80..f784270bab 100644
--- a/keras_nlp/models/bart/bart_seq_2_seq_lm.py
+++ b/keras_nlp/models/bart/bart_seq_2_seq_lm.py
@@ -15,7 +15,6 @@
 
 import copy
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.models.bart.bart_backbone import BartBackbone
@@ -41,7 +40,7 @@ def compute_output_shape(self, input_shape):
         return (input_shape[0],) + (self.embedding.embeddings.shape[0],)
 
 
-@keras_nlp_export("keras_nlp.models.BartSeq2SeqLM")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class BartSeq2SeqLM(GenerativeTask):
     """An end-to-end BART model for seq2seq language modeling.
 
diff --git a/keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py b/keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py
index 0eb9a6cb44..5b98cf9b8b 100644
--- a/keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py
+++ b/keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py
@@ -19,7 +19,7 @@
 import tensorflow as tf
 from absl import logging
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.models.bart.bart_preprocessor import BartPreprocessor
 from keras_nlp.models.bart.bart_presets import backbone_presets
@@ -30,7 +30,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras_nlp_export("keras_nlp.models.BartSeq2SeqLMPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class BartSeq2SeqLMPreprocessor(BartPreprocessor):
     """BART Seq2Seq LM preprocessor.
 
diff --git a/keras_nlp/models/bart/bart_tokenizer.py b/keras_nlp/models/bart/bart_tokenizer.py
index 0737ef03b3..f1f8935c8e 100644
--- a/keras_nlp/models/bart/bart_tokenizer.py
+++ b/keras_nlp/models/bart/bart_tokenizer.py
@@ -16,13 +16,13 @@
 
 import copy
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.models.bart.bart_presets import backbone_presets
 from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras_nlp_export("keras_nlp.models.BartTokenizer")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class BartTokenizer(BytePairTokenizer):
     """A BART tokenizer using Byte-Pair Encoding subword segmentation.
 
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
index 3a3ce2f3e1..8256daefce 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
@@ -11,7 +11,6 @@
 # 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.
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.gpt_neo_x.gpt_neo_x_decoder import GPTNeoXDecoder
@@ -22,7 +21,7 @@ def _gpt_neo_x_kernel_initializer(stddev=0.02):
     return keras.initializers.RandomNormal(stddev=stddev)
 
 
-@keras_nlp_export("keras_nlp.models.GPTNeoXBackbone")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class GPTNeoXBackbone(Backbone):
     """GPT-NeoX core network with hyperparameters.
 
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py
index cf93f81fba..0bca9b07ce 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py
@@ -17,7 +17,7 @@
 import tensorflow as tf
 from absl import logging
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.models.gpt_neo_x.gpt_neo_x_preprocessor import (
     GPTNeoXPreprocessor,
@@ -28,7 +28,7 @@
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras_nlp_export("keras_nlp.models.GPTNeoXCausalLMPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class GPTNeoXCausalLMPreprocessor(GPTNeoXPreprocessor):
     """GPTNeoX Causal LM preprocessor.
 
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_preprocessor.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_preprocessor.py
index b20fb8157d..beb74c16f9 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_preprocessor.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_preprocessor.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """GPTNeoX preprocessor layer."""
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.models.gpt_neo_x.gpt_neo_x_tokenizer import GPTNeoXTokenizer
 from keras_nlp.models.preprocessor import Preprocessor
@@ -24,7 +24,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras_nlp_export("keras_nlp.models.GPTNeoXPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class GPTNeoXPreprocessor(Preprocessor):
     """GPTNeoX preprocessing layer which tokenizes and packs inputs.
 
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_tokenizer.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_tokenizer.py
index e34ade14fe..35e4454908 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_tokenizer.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_tokenizer.py
@@ -13,11 +13,11 @@
 # limitations under the License.
 """GptNeoX tokenizer"""
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
-@keras_nlp_export("keras_nlp.models.GPTNeoXTokenizer")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class GPTNeoXTokenizer(BytePairTokenizer):
     """A GPTNeoX tokenizer using Byte-Pair Encoding subword segmentation.
 
diff --git a/keras_nlp/models/t5/t5_backbone.py b/keras_nlp/models/t5/t5_backbone.py
index b23e0becd4..9423362365 100644
--- a/keras_nlp/models/t5/t5_backbone.py
+++ b/keras_nlp/models/t5/t5_backbone.py
@@ -14,7 +14,6 @@
 
 """T5 backbone model."""
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.t5.t5_layer_norm import T5LayerNorm
@@ -23,7 +22,7 @@
 from keras_nlp.utils.tensor_utils import assert_tf_backend
 
 
-@keras_nlp_export("keras_nlp.models.T5Backbone")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class T5Backbone(Backbone):
     """T5 encoder-decoder backbone model.
 
diff --git a/keras_nlp/models/t5/t5_tokenizer.py b/keras_nlp/models/t5/t5_tokenizer.py
index 5eb2437b5d..007eed0a3e 100644
--- a/keras_nlp/models/t5/t5_tokenizer.py
+++ b/keras_nlp/models/t5/t5_tokenizer.py
@@ -14,11 +14,11 @@
 
 """T5 tokenizer."""
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
 
 
-@keras_nlp_export("keras_nlp.models.T5Tokenizer")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class T5Tokenizer(SentencePieceTokenizer):
     """T5 tokenizer layer based on SentencePiece.
 
diff --git a/keras_nlp/models/whisper/whisper_audio_feature_extractor.py b/keras_nlp/models/whisper/whisper_audio_feature_extractor.py
index 73b15a3afd..843e4b91cb 100644
--- a/keras_nlp/models/whisper/whisper_audio_feature_extractor.py
+++ b/keras_nlp/models/whisper/whisper_audio_feature_extractor.py
@@ -17,14 +17,13 @@
 import numpy as np
 import tensorflow as tf
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.models.whisper.whisper_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
 
-@keras_nlp_export("keras_nlp.models.WhisperAudioFeatureExtractor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class WhisperAudioFeatureExtractor(keras.layers.Layer):
     """
     Whisper audio feature extractor layer.
diff --git a/keras_nlp/models/whisper/whisper_backbone.py b/keras_nlp/models/whisper/whisper_backbone.py
index 1bd104d5ff..9271117f31 100644
--- a/keras_nlp/models/whisper/whisper_backbone.py
+++ b/keras_nlp/models/whisper/whisper_backbone.py
@@ -15,7 +15,6 @@
 
 import copy
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
@@ -39,7 +38,7 @@ def call(self, x):
         return ops.pad(x, [[0, 0], [1, 1], [0, 0]])
 
 
-@keras_nlp_export("keras_nlp.models.WhisperBackbone")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class WhisperBackbone(Backbone):
     """A Whisper encoder-decoder network for speech.
 
diff --git a/keras_nlp/models/whisper/whisper_preprocessor.py b/keras_nlp/models/whisper/whisper_preprocessor.py
index 62d853698a..05a85a48a0 100644
--- a/keras_nlp/models/whisper/whisper_preprocessor.py
+++ b/keras_nlp/models/whisper/whisper_preprocessor.py
@@ -18,7 +18,6 @@
 import tensorflow as tf
 from absl import logging
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.models.preprocessor import Preprocessor
@@ -35,7 +34,7 @@
 from keras_nlp.utils.python_utils import format_docstring
 
 
-@keras_nlp_export("keras_nlp.models.WhisperPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class WhisperPreprocessor(Preprocessor):
     """A Whisper preprocessing layer which handles audio and text input.
 
diff --git a/keras_nlp/models/whisper/whisper_tokenizer.py b/keras_nlp/models/whisper/whisper_tokenizer.py
index 37a5e6078a..d372039d4d 100644
--- a/keras_nlp/models/whisper/whisper_tokenizer.py
+++ b/keras_nlp/models/whisper/whisper_tokenizer.py
@@ -16,7 +16,7 @@
 import copy
 import json
 
-from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
 from keras_nlp.models.whisper.whisper_presets import backbone_presets
 from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_nlp.utils.python_utils import classproperty
@@ -29,7 +29,7 @@ def _load_dict(dict_or_path):
     return dict_or_path
 
 
-@keras_nlp_export("keras_nlp.models.WhisperTokenizer")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class WhisperTokenizer(BytePairTokenizer):
     """Whisper text tokenizer using Byte-Pair Encoding subword segmentation.