keras-team · mattdangerw · May 3, 2024 · May 3, 2024
diff --git a/keras_nlp/api/models/__init__.py b/keras_nlp/api/models/__init__.py
@@ -99,12 +99,6 @@
 )
 from keras_nlp.src.models.f_net.f_net_preprocessor import FNetPreprocessor
 from keras_nlp.src.models.f_net.f_net_tokenizer import FNetTokenizer
-from keras_nlp.src.models.falcon.falcon_backbone import FalconBackbone
-from keras_nlp.src.models.falcon.falcon_causal_lm_preprocessor import (
-    FalconCausalLMPreprocessor,
-)
-from keras_nlp.src.models.falcon.falcon_preprocessor import FalconPreprocessor
-from keras_nlp.src.models.falcon.falcon_tokenizer import FalconTokenizer
 from keras_nlp.src.models.gemma.gemma_backbone import GemmaBackbone
 from keras_nlp.src.models.gemma.gemma_causal_lm import GemmaCausalLM
 from keras_nlp.src.models.gemma.gemma_causal_lm_preprocessor import (
@@ -119,15 +113,6 @@
 )
 from keras_nlp.src.models.gpt2.gpt2_preprocessor import GPT2Preprocessor
 from keras_nlp.src.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
-from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_backbone import GPTNeoXBackbone
-from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_causal_lm import GPTNeoXCausalLM
-from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_causal_lm_preprocessor import (
-    GPTNeoXCausalLMPreprocessor,
-)
-from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_preprocessor import (
-    GPTNeoXPreprocessor,
-)
-from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_tokenizer import GPTNeoXTokenizer
 from keras_nlp.src.models.llama.llama_backbone import LlamaBackbone
 from keras_nlp.src.models.llama.llama_causal_lm import LlamaCausalLM
 from keras_nlp.src.models.llama.llama_causal_lm_preprocessor import (
@@ -164,17 +149,10 @@
 )
 from keras_nlp.src.models.roberta.roberta_tokenizer import RobertaTokenizer
 from keras_nlp.src.models.seq_2_seq_lm import Seq2SeqLM
-from keras_nlp.src.models.t5.t5_backbone import T5Backbone
-from keras_nlp.src.models.t5.t5_tokenizer import T5Tokenizer
 from keras_nlp.src.models.task import Task
 from keras_nlp.src.models.whisper.whisper_audio_feature_extractor import (
     WhisperAudioFeatureExtractor,
 )
-from keras_nlp.src.models.whisper.whisper_backbone import WhisperBackbone
-from keras_nlp.src.models.whisper.whisper_preprocessor import (
-    WhisperPreprocessor,
-)
-from keras_nlp.src.models.whisper.whisper_tokenizer import WhisperTokenizer
 from keras_nlp.src.models.xlm_roberta.xlm_roberta_backbone import (
     XLMRobertaBackbone,
 )

diff --git a/keras_nlp/src/models/falcon/falcon_backbone.py b/keras_nlp/src/models/falcon/falcon_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.src.api_export import keras_nlp_export
 from keras_nlp.src.backend import keras
 from keras_nlp.src.layers.modeling.reversible_embedding import (
     ReversibleEmbedding,
@@ -22,7 +21,7 @@
 )
 
 
-@keras_nlp_export("keras_nlp.models.FalconBackbone")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class FalconBackbone(Backbone):
     """The Falcon core architecure.
 

diff --git a/keras_nlp/src/models/falcon/falcon_causal_lm_preprocessor.py b/keras_nlp/src/models/falcon/falcon_causal_lm_preprocessor.py
@@ -21,7 +21,7 @@
     )
 from absl import logging
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.backend import ops
 from keras_nlp.src.models.falcon.falcon_preprocessor import FalconPreprocessor
 from keras_nlp.src.utils.keras_utils import (
@@ -30,7 +30,7 @@
 from keras_nlp.src.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras_nlp_export("keras_nlp.models.FalconCausalLMPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class FalconCausalLMPreprocessor(FalconPreprocessor):
     """Falcon Causal LM preprocessor.
 

diff --git a/keras_nlp/src/models/falcon/falcon_preprocessor.py b/keras_nlp/src/models/falcon/falcon_preprocessor.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.src.models.falcon.falcon_tokenizer import FalconTokenizer
 from keras_nlp.src.models.preprocessor import Preprocessor
@@ -23,7 +23,7 @@
 from keras_nlp.src.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras_nlp_export("keras_nlp.models.FalconPreprocessor")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class FalconPreprocessor(Preprocessor):
     """Falcon preprocessing layer which tokenizes and packs inputs.
 

diff --git a/keras_nlp/src/models/falcon/falcon_tokenizer.py b/keras_nlp/src/models/falcon/falcon_tokenizer.py
@@ -13,11 +13,11 @@
 # limitations under the License.
 
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
-@keras_nlp_export("keras_nlp.models.FalconTokenizer")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class FalconTokenizer(BytePairTokenizer):
     """Falcon tokenizer based on BytePairTokenizer.
 

diff --git a/keras_nlp/src/models/gpt_neo_x/gpt_neo_x_backbone.py b/keras_nlp/src/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.src.api_export import keras_nlp_export
 from keras_nlp.src.backend import keras
 from keras_nlp.src.layers.modeling.reversible_embedding import (
     ReversibleEmbedding,
@@ -26,7 +25,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/src/models/gpt_neo_x/gpt_neo_x_causal_lm.py b/keras_nlp/src/models/gpt_neo_x/gpt_neo_x_causal_lm.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.backend import ops
 from keras_nlp.src.models.causal_lm import CausalLM
 from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_backbone import GPTNeoXBackbone
@@ -22,7 +22,7 @@
 from keras_nlp.src.utils.tensor_utils import any_equal
 
 
-@keras_nlp_export("keras_nlp.models.GPTNeoXCausalLM")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 class GPTNeoXCausalLM(CausalLM):
     """An end-to-end GPTNeoX model for causal language modeling.
 

diff --git a/keras_nlp/src/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py b/keras_nlp/src/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py
@@ -21,7 +21,7 @@
     )
 from absl import logging
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.backend import ops
 from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_preprocessor import (
     GPTNeoXPreprocessor,
@@ -32,7 +32,7 @@
 from keras_nlp.src.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/src/models/gpt_neo_x/gpt_neo_x_preprocessor.py b/keras_nlp/src/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.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.src.models.gpt_neo_x.gpt_neo_x_tokenizer import GPTNeoXTokenizer
 from keras_nlp.src.models.preprocessor import Preprocessor
@@ -22,7 +22,7 @@
 from keras_nlp.src.utils.keras_utils import pack_x_y_sample_weight
 
 
-@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/src/models/gpt_neo_x/gpt_neo_x_tokenizer.py b/keras_nlp/src/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.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.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/src/models/t5/t5_backbone.py b/keras_nlp/src/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.src.api_export import keras_nlp_export
 from keras_nlp.src.backend import keras
 from keras_nlp.src.layers.modeling.reversible_embedding import (
     ReversibleEmbedding,
@@ -22,7 +21,7 @@
 from keras_nlp.src.models.t5.t5_transformer_layer import T5TransformerLayer
 
 
-@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/src/models/t5/t5_tokenizer.py b/keras_nlp/src/models/t5/t5_tokenizer.py
@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.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/src/models/whisper/whisper_backbone.py b/keras_nlp/src/models/whisper/whisper_backbone.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 
-from keras_nlp.src.api_export import keras_nlp_export
 from keras_nlp.src.backend import keras
 from keras_nlp.src.backend import ops
 from keras_nlp.src.layers.modeling.position_embedding import PositionEmbedding
@@ -35,7 +34,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/src/models/whisper/whisper_preprocessor.py b/keras_nlp/src/models/whisper/whisper_preprocessor.py
@@ -15,7 +15,6 @@
 
 from absl import logging
 
-from keras_nlp.src.api_export import keras_nlp_export
 from keras_nlp.src.backend import keras
 from keras_nlp.src.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.src.models.preprocessor import Preprocessor
@@ -29,7 +28,7 @@
 from keras_nlp.src.utils.keras_utils import pack_x_y_sample_weight
 
 
-@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/src/models/whisper/whisper_tokenizer.py b/keras_nlp/src/models/whisper/whisper_tokenizer.py
@@ -14,7 +14,7 @@
 
 import json
 
-from keras_nlp.src.api_export import keras_nlp_export
+from keras_nlp.src.backend import keras
 from keras_nlp.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
@@ -25,7 +25,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.