Unexport models that are not in the 0.6 release (#1125)

mattdangerw · mattdangerw · commit 0d3b3b221ab7 · 2023-09-21T12:40:19.000-07:00
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!
diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
@@ -14,7 +14,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.reversible_embedding import ReversibleEmbedding
@@ -29,7 +28,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
@@ -14,7 +14,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
@@ -26,7 +26,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
@@ -14,7 +14,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
@@ -26,7 +25,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@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
@@ -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.bart.bart_preprocessor import BartPreprocessor
 from keras_nlp.models.bart.bart_presets import backbone_presets
@@ -28,7 +28,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
@@ -14,13 +14,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
@@ -12,7 +12,6 @@
 # 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.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
@@ -23,7 +22,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
@@ -15,7 +15,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,
@@ -26,7 +26,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):
     """GPT-NeoX 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
@@ -12,7 +12,7 @@
 # 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.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
@@ -23,7 +23,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
@@ -12,11 +12,11 @@
 # 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.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
@@ -12,7 +12,6 @@
 # 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.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
@@ -22,7 +21,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
@@ -12,11 +12,11 @@
 # 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.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
@@ -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
@@ -14,7 +14,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
@@ -38,7 +37,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
@@ -17,7 +17,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
@@ -34,7 +33,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
@@ -15,7 +15,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
@@ -28,7 +28,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.
 
diff --git a/keras_nlp/models/xlnet/relative_attention.py b/keras_nlp/models/xlnet/relative_attention.py
@@ -15,7 +15,6 @@
 import math
 import string
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 
@@ -76,7 +75,7 @@ def _rel_shift(x, klen=-1):
     return x
 
 
-@keras_nlp_export("keras_nlp.layers.TwoStreamRelativeAttention")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
     """Two-stream relative self-attention for XLNet.
 
diff --git a/keras_nlp/models/xlnet/xlnet_backbone.py b/keras_nlp/models/xlnet/xlnet_backbone.py
@@ -12,7 +12,6 @@
 # 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.xlnet.xlnet_content_and_query_embedding import (
@@ -23,7 +22,7 @@
 from keras_nlp.models.xlnet.xlnet_encoder import XLNetSegmentMatrixLayer
 
 
-@keras_nlp_export("keras_nlp.models.XLNetBackbone")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class XLNetBackbone(Backbone):
     """XLNet encoder network.
 
diff --git a/keras_nlp/models/xlnet/xlnet_encoder.py b/keras_nlp/models/xlnet/xlnet_encoder.py
@@ -12,7 +12,6 @@
 # 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.backend import ops
 from keras_nlp.models.xlnet.relative_attention import TwoStreamRelativeAttention
@@ -22,7 +21,7 @@ def xlnet_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras_nlp_export("keras_nlp.layers.XLNetEncoder")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class XLNetEncoder(keras.layers.Layer):
     """
     XLNet Encoder.
