Add API exports for models documented on keras.io

keras_nlp/models/albert/albert_backbone.py

@@ -19,6 +19,7 @@
 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_encoder import TransformerEncoder
 from keras_nlp.models.albert.albert_presets import backbone_presets
@@ -30,7 +31,7 @@ def albert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.AlbertBackbone")
 class AlbertBackbone(Backbone):
     """ALBERT encoder network.
 

keras_nlp/models/albert/albert_classifier.py

@@ -17,6 +17,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.albert.albert_backbone import AlbertBackbone
 from keras_nlp.models.albert.albert_backbone import albert_kernel_initializer
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
@@ -26,7 +27,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.AlbertClassifier")
 class AlbertClassifier(Task):
     """An end-to-end ALBERT model for classification tasks
 

keras_nlp/models/albert/albert_masked_lm.py

@@ -18,6 +18,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.masked_lm_head import MaskedLMHead
 from keras_nlp.models.albert.albert_backbone import AlbertBackbone
 from keras_nlp.models.albert.albert_backbone import albert_kernel_initializer
@@ -29,7 +30,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.AlbertMaskedLM")
 class AlbertMaskedLM(Task):
     """An end-to-end ALBERT model for the masked language modeling task.
 

keras_nlp/models/albert/albert_masked_lm_preprocessor.py

@@ -15,14 +15,14 @@
 """ALBERT masked language model preprocessor layer."""
 
 from absl import logging
-from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.masked_lm_mask_generator import MaskedLMMaskGenerator
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.AlbertMaskedLMPreprocessor")
 class AlbertMaskedLMPreprocessor(AlbertPreprocessor):
     """ALBERT preprocessing for the masked language modeling task.
 

keras_nlp/models/albert/albert_preprocessor.py

@@ -15,8 +15,7 @@
 
 import copy
 
-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.albert.albert_presets import backbone_presets
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
@@ -28,7 +27,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.AlbertPreprocessor")
 class AlbertPreprocessor(Preprocessor):
     """An ALBERT preprocessing layer which tokenizes and packs inputs.
 

keras_nlp/models/albert/albert_tokenizer.py

@@ -16,14 +16,13 @@
 
 import copy
 
-from tensorflow import keras
-
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.albert.albert_presets import backbone_presets
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.AlbertTokenizer")
 class AlbertTokenizer(SentencePieceTokenizer):
     """ALBERT tokenizer layer based on SentencePiece.
 

keras_nlp/models/backbone.py

@@ -18,11 +18,12 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.Backbone")
 class Backbone(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

keras_nlp/models/bart/bart_backbone.py

@@ -19,6 +19,7 @@
 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
@@ -31,7 +32,7 @@ def bart_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BartBackbone")
 class BartBackbone(Backbone):
     """BART encoder-decoder network.
 

keras_nlp/models/bart/bart_tokenizer.py

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

keras_nlp/models/bert/bert_backbone.py

@@ -19,6 +19,7 @@
 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_encoder import TransformerEncoder
 from keras_nlp.models.backbone import Backbone
@@ -30,7 +31,7 @@ def bert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BertBackbone")
 class BertBackbone(Backbone):
     """BERT encoder network.
 

keras_nlp/models/bert/bert_classifier.py

@@ -17,6 +17,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_backbone import bert_kernel_initializer
 from keras_nlp.models.bert.bert_preprocessor import BertPreprocessor
@@ -27,7 +28,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BertClassifier")
 class BertClassifier(Task):
     """An end-to-end BERT model for classification tasks
 

keras_nlp/models/bert/bert_masked_lm.py

@@ -17,6 +17,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.masked_lm_head import MaskedLMHead
 from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_backbone import bert_kernel_initializer
@@ -28,7 +29,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BertMaskedLM")
 class BertMaskedLM(Task):
     """An end-to-end BERT model for the masked language modeling task.
 

keras_nlp/models/bert/bert_masked_lm_preprocessor.py

@@ -15,14 +15,14 @@
 """BERT masked language model preprocessor layer."""
 
 from absl import logging
-from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.masked_lm_mask_generator import MaskedLMMaskGenerator
 from keras_nlp.models.bert.bert_preprocessor import BertPreprocessor
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BertMaskedLMPreprocessor")
 class BertMaskedLMPreprocessor(BertPreprocessor):
     """BERT preprocessing for the masked language modeling task.
     This preprocessing layer will prepare inputs for a masked language modeling

keras_nlp/models/bert/bert_preprocessor.py

@@ -15,8 +15,7 @@
 
 import copy
 
-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.bert.bert_presets import backbone_presets
 from keras_nlp.models.bert.bert_presets import classifier_presets
@@ -31,7 +30,7 @@
 PRESET_NAMES = ", ".join(list(backbone_presets) + list(classifier_presets))
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BertPreprocessor")
 class BertPreprocessor(Preprocessor):
     """A BERT preprocessing layer which tokenizes and packs inputs.
 

keras_nlp/models/bert/bert_tokenizer.py

@@ -15,8 +15,7 @@
 
 import copy
 
-from tensorflow import keras
-
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.bert.bert_presets import backbone_presets
 from keras_nlp.models.bert.bert_presets import classifier_presets
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
@@ -25,7 +24,7 @@
 PRESET_NAMES = ", ".join(list(backbone_presets) + list(classifier_presets))
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.BertTokenizer")
 class BertTokenizer(WordPieceTokenizer):
     """A BERT tokenizer using WordPiece subword segmentation.
 

keras_nlp/models/deberta_v3/deberta_v3_backbone.py

@@ -19,6 +19,7 @@
 import tensorflow as tf
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.models.deberta_v3.disentangled_attention_encoder import (
@@ -32,7 +33,7 @@ def deberta_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DebertaV3Backbone")
 class DebertaV3Backbone(Backbone):
     """DeBERTa encoder network.
 

keras_nlp/models/deberta_v3/deberta_v3_classifier.py

@@ -17,6 +17,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import (
     deberta_kernel_initializer,
@@ -30,7 +31,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DebertaV3Classifier")
 class DebertaV3Classifier(Task):
     """An end-to-end DeBERTa model for classification tasks.
 

keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py

@@ -17,6 +17,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.masked_lm_head import MaskedLMHead
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import (
@@ -30,7 +31,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DebertaV3MaskedLM")
 class DebertaV3MaskedLM(Task):
     """An end-to-end DeBERTaV3 model for the masked language modeling task.
 

keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor.py

@@ -15,16 +15,16 @@
 """DeBERTa masked language model preprocessor layer."""
 
 from absl import logging
-from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.masked_lm_mask_generator import MaskedLMMaskGenerator
 from keras_nlp.models.deberta_v3.deberta_v3_preprocessor import (
     DebertaV3Preprocessor,
 )
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DebertaV3MaskedLMPreprocessor")
 class DebertaV3MaskedLMPreprocessor(DebertaV3Preprocessor):
     """DeBERTa preprocessing for the masked language modeling task.
 

keras_nlp/models/deberta_v3/deberta_v3_preprocessor.py

@@ -15,8 +15,7 @@
 
 import copy
 
-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.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.models.deberta_v3.deberta_v3_tokenizer import DebertaV3Tokenizer
@@ -28,7 +27,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DebertaV3Preprocessor")
 class DebertaV3Preprocessor(Preprocessor):
     """A DeBERTa preprocessing layer which tokenizes and packs inputs.
 

keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py

@@ -16,14 +16,13 @@
 
 import copy
 
-from tensorflow import keras
-
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DebertaV3Tokenizer")
 class DebertaV3Tokenizer(SentencePieceTokenizer):
     """DeBERTa tokenizer layer based on SentencePiece.
 

keras_nlp/models/distil_bert/distil_bert_backbone.py

@@ -19,6 +19,7 @@
 import tensorflow as tf
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
@@ -32,7 +33,7 @@ def distilbert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DistilBertBackbone")
 class DistilBertBackbone(Backbone):
     """DistilBERT encoder network.
 

keras_nlp/models/distil_bert/distil_bert_classifier.py

@@ -17,6 +17,7 @@
 
 from tensorflow import keras
 
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_backbone import (
     distilbert_kernel_initializer,
@@ -30,7 +31,7 @@
 from keras_nlp.utils.python_utils import classproperty
 
 
-@keras.utils.register_keras_serializable(package="keras_nlp")
+@keras_nlp_export("keras_nlp.models.DistilBertClassifier")
 class DistilBertClassifier(Task):
     """An end-to-end DistilBERT model for classification tasks.