Skip exporting BART, T5 and Whisper for 0.5 release (#1003)

* Version bump to 0.5.0.dev0 (#1002)

* Skip exporting BART, T5 and Whisper

* remove export

* fix version number

* fix version number

keras_nlp/__init__.py

@@ -28,4 +28,4 @@
 from keras_nlp import utils
 
 # This is the global source of truth for the version number.
-__version__ = "0.5.0.dev0"
+__version__ = "0.5.1"

keras_nlp/models/__init__.py

@@ -20,12 +20,6 @@
 )
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
-from keras_nlp.models.bart.bart_backbone import BartBackbone
-from keras_nlp.models.bart.bart_preprocessor import BartPreprocessor
-from keras_nlp.models.bart.bart_seq_2_seq_lm_preprocessor import (
-    BartSeq2SeqLMPreprocessor,
-)
-from keras_nlp.models.bart.bart_tokenizer import BartTokenizer
 from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_classifier import BertClassifier
 from keras_nlp.models.bert.bert_masked_lm import BertMaskedLM
@@ -92,9 +86,6 @@
 )
 from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tokenizer import RobertaTokenizer
-from keras_nlp.models.t5.t5_backbone import T5Backbone
-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 (
     XLMRobertaClassifier,

keras_nlp/models/bart/bart_backbone.py

@@ -19,7 +19,6 @@
 import tensorflow as tf
 from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.transformer_decoder import TransformerDecoder
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
@@ -32,7 +31,7 @@ def bart_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras_nlp_export("keras_nlp.models.BartBackbone")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class BartBackbone(Backbone):
     """BART encoder-decoder network.
 

keras_nlp/models/bart/bart_preprocessor.py

@@ -16,8 +16,8 @@
 import copy
 
 import tensorflow as tf
+from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.models.bart.bart_presets import backbone_presets
 from keras_nlp.models.bart.bart_tokenizer import BartTokenizer
@@ -29,7 +29,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras_nlp_export("keras_nlp.models.BartPreprocessor")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class BartPreprocessor(Preprocessor):
     """A BART preprocessing layer which tokenizes and packs inputs.
 

keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py

@@ -15,13 +15,13 @@
 """BART Seq2Seq LM preprocessor layer."""
 
 from absl import logging
+from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.bart.bart_preprocessor import BartPreprocessor
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras_nlp_export("keras_nlp.models.BartSeq2SeqLMPreprocessor")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class BartSeq2SeqLMPreprocessor(BartPreprocessor):
     """BART Seq2Seq LM preprocessor.
 

keras_nlp/models/bart/bart_tokenizer.py

@@ -16,13 +16,14 @@
 
 import copy
 
-from keras_nlp.api_export import keras_nlp_export
+from tensorflow 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.utils.register_keras_serializable(package="keras_nlp")
 class BartTokenizer(BytePairTokenizer):
     """A BART tokenizer using Byte-Pair Encoding subword segmentation.
 

keras_nlp/models/t5/t5_backbone.py

@@ -17,15 +17,14 @@
 import tensorflow as tf
 from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.transformer_layer_utils import compute_causal_mask
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.t5.t5_layer_norm import T5LayerNorm
 from keras_nlp.models.t5.t5_transformer_layer import T5TransformerLayer
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras_nlp_export("keras_nlp.models.T5Backbone")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class T5Backbone(Backbone):
     """T5 encoder-decoder backbone model.
 

keras_nlp/models/t5/t5_tokenizer.py

@@ -14,11 +14,12 @@
 
 """T5 tokenizer."""
 
-from keras_nlp.api_export import keras_nlp_export
+from tensorflow import keras
+
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
 
 
-@keras_nlp_export("keras_nlp.models.T5Tokenizer")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class T5Tokenizer(SentencePieceTokenizer):
     """T5 tokenizer layer based on SentencePiece.
 

keras_nlp/models/whisper/whisper_backbone.py

@@ -17,7 +17,6 @@
 import tensorflow as tf
 from tensorflow import keras
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
@@ -36,7 +35,7 @@ def whisper_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras_nlp_export("keras_nlp.models.WhisperBackbone")
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class WhisperBackbone(Backbone):
     """A Whisper encoder-decoder network for speech.