diff --git a/keras_nlp/api/models/__init__.py b/keras_nlp/api/models/__init__.py index 5795c85532..e3b52d0bea 100644 --- 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 index 53f33eccef..a0dfca4e28 100644 --- 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 index 6ce531f1dd..9418e4309e 100644 --- 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 index fb3569ef22..9e3f8a5c0d 100644 --- 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 index 7d00459008..fdcc14fe93 100644 --- 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 index 58ff85f709..c1d179adb0 100644 --- 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 index 79cc6047be..eb3d1e994b 100644 --- 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 index f01b1082fb..4326acf481 100644 --- 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 index 2443e7c241..beebfb2f60 100644 --- 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 index e270ae8579..46e146775c 100644 --- 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 index 7db1b0f8ff..1deeb7f837 100644 --- 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 index 411a343f98..2f5eb12451 100644 --- 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 index 3ef73f58d8..6e4a5955e0 100644 --- 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 index 432fa4aebf..a9337039ed 100644 --- 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 index e6f2a7d443..ae1669e45f 100644 --- 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.