keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- keras_hub/__init__.py +52 -0
- keras_hub/api/__init__.py +27 -0
- keras_hub/api/layers/__init__.py +47 -0
- keras_hub/api/metrics/__init__.py +24 -0
- keras_hub/api/models/__init__.py +249 -0
- keras_hub/api/samplers/__init__.py +29 -0
- keras_hub/api/tokenizers/__init__.py +35 -0
- keras_hub/src/__init__.py +13 -0
- keras_hub/src/api_export.py +53 -0
- keras_hub/src/layers/__init__.py +13 -0
- keras_hub/src/layers/modeling/__init__.py +13 -0
- keras_hub/src/layers/modeling/alibi_bias.py +143 -0
- keras_hub/src/layers/modeling/cached_multi_head_attention.py +137 -0
- keras_hub/src/layers/modeling/f_net_encoder.py +200 -0
- keras_hub/src/layers/modeling/masked_lm_head.py +239 -0
- keras_hub/src/layers/modeling/position_embedding.py +123 -0
- keras_hub/src/layers/modeling/reversible_embedding.py +311 -0
- keras_hub/src/layers/modeling/rotary_embedding.py +169 -0
- keras_hub/src/layers/modeling/sine_position_encoding.py +108 -0
- keras_hub/src/layers/modeling/token_and_position_embedding.py +150 -0
- keras_hub/src/layers/modeling/transformer_decoder.py +496 -0
- keras_hub/src/layers/modeling/transformer_encoder.py +262 -0
- keras_hub/src/layers/modeling/transformer_layer_utils.py +106 -0
- keras_hub/src/layers/preprocessing/__init__.py +13 -0
- keras_hub/src/layers/preprocessing/masked_lm_mask_generator.py +220 -0
- keras_hub/src/layers/preprocessing/multi_segment_packer.py +319 -0
- keras_hub/src/layers/preprocessing/preprocessing_layer.py +62 -0
- keras_hub/src/layers/preprocessing/random_deletion.py +271 -0
- keras_hub/src/layers/preprocessing/random_swap.py +267 -0
- keras_hub/src/layers/preprocessing/start_end_packer.py +219 -0
- keras_hub/src/metrics/__init__.py +13 -0
- keras_hub/src/metrics/bleu.py +394 -0
- keras_hub/src/metrics/edit_distance.py +197 -0
- keras_hub/src/metrics/perplexity.py +181 -0
- keras_hub/src/metrics/rouge_base.py +204 -0
- keras_hub/src/metrics/rouge_l.py +97 -0
- keras_hub/src/metrics/rouge_n.py +125 -0
- keras_hub/src/models/__init__.py +13 -0
- keras_hub/src/models/albert/__init__.py +20 -0
- keras_hub/src/models/albert/albert_backbone.py +267 -0
- keras_hub/src/models/albert/albert_classifier.py +202 -0
- keras_hub/src/models/albert/albert_masked_lm.py +129 -0
- keras_hub/src/models/albert/albert_masked_lm_preprocessor.py +194 -0
- keras_hub/src/models/albert/albert_preprocessor.py +206 -0
- keras_hub/src/models/albert/albert_presets.py +70 -0
- keras_hub/src/models/albert/albert_tokenizer.py +119 -0
- keras_hub/src/models/backbone.py +311 -0
- keras_hub/src/models/bart/__init__.py +20 -0
- keras_hub/src/models/bart/bart_backbone.py +261 -0
- keras_hub/src/models/bart/bart_preprocessor.py +276 -0
- keras_hub/src/models/bart/bart_presets.py +74 -0
- keras_hub/src/models/bart/bart_seq_2_seq_lm.py +490 -0
- keras_hub/src/models/bart/bart_seq_2_seq_lm_preprocessor.py +262 -0
- keras_hub/src/models/bart/bart_tokenizer.py +124 -0
- keras_hub/src/models/bert/__init__.py +23 -0
- keras_hub/src/models/bert/bert_backbone.py +227 -0
- keras_hub/src/models/bert/bert_classifier.py +183 -0
- keras_hub/src/models/bert/bert_masked_lm.py +131 -0
- keras_hub/src/models/bert/bert_masked_lm_preprocessor.py +198 -0
- keras_hub/src/models/bert/bert_preprocessor.py +184 -0
- keras_hub/src/models/bert/bert_presets.py +147 -0
- keras_hub/src/models/bert/bert_tokenizer.py +112 -0
- keras_hub/src/models/bloom/__init__.py +20 -0
- keras_hub/src/models/bloom/bloom_attention.py +186 -0
- keras_hub/src/models/bloom/bloom_backbone.py +173 -0
- keras_hub/src/models/bloom/bloom_causal_lm.py +298 -0
- keras_hub/src/models/bloom/bloom_causal_lm_preprocessor.py +176 -0
- keras_hub/src/models/bloom/bloom_decoder.py +206 -0
- keras_hub/src/models/bloom/bloom_preprocessor.py +185 -0
- keras_hub/src/models/bloom/bloom_presets.py +121 -0
- keras_hub/src/models/bloom/bloom_tokenizer.py +116 -0
- keras_hub/src/models/causal_lm.py +383 -0
- keras_hub/src/models/classifier.py +109 -0
- keras_hub/src/models/csp_darknet/__init__.py +13 -0
- keras_hub/src/models/csp_darknet/csp_darknet_backbone.py +410 -0
- keras_hub/src/models/csp_darknet/csp_darknet_image_classifier.py +133 -0
- keras_hub/src/models/deberta_v3/__init__.py +24 -0
- keras_hub/src/models/deberta_v3/deberta_v3_backbone.py +210 -0
- keras_hub/src/models/deberta_v3/deberta_v3_classifier.py +228 -0
- keras_hub/src/models/deberta_v3/deberta_v3_masked_lm.py +135 -0
- keras_hub/src/models/deberta_v3/deberta_v3_masked_lm_preprocessor.py +191 -0
- keras_hub/src/models/deberta_v3/deberta_v3_preprocessor.py +206 -0
- keras_hub/src/models/deberta_v3/deberta_v3_presets.py +82 -0
- keras_hub/src/models/deberta_v3/deberta_v3_tokenizer.py +155 -0
- keras_hub/src/models/deberta_v3/disentangled_attention_encoder.py +227 -0
- keras_hub/src/models/deberta_v3/disentangled_self_attention.py +412 -0
- keras_hub/src/models/deberta_v3/relative_embedding.py +94 -0
- keras_hub/src/models/densenet/__init__.py +13 -0
- keras_hub/src/models/densenet/densenet_backbone.py +210 -0
- keras_hub/src/models/densenet/densenet_image_classifier.py +131 -0
- keras_hub/src/models/distil_bert/__init__.py +26 -0
- keras_hub/src/models/distil_bert/distil_bert_backbone.py +187 -0
- keras_hub/src/models/distil_bert/distil_bert_classifier.py +208 -0
- keras_hub/src/models/distil_bert/distil_bert_masked_lm.py +137 -0
- keras_hub/src/models/distil_bert/distil_bert_masked_lm_preprocessor.py +194 -0
- keras_hub/src/models/distil_bert/distil_bert_preprocessor.py +175 -0
- keras_hub/src/models/distil_bert/distil_bert_presets.py +57 -0
- keras_hub/src/models/distil_bert/distil_bert_tokenizer.py +114 -0
- keras_hub/src/models/electra/__init__.py +20 -0
- keras_hub/src/models/electra/electra_backbone.py +247 -0
- keras_hub/src/models/electra/electra_preprocessor.py +154 -0
- keras_hub/src/models/electra/electra_presets.py +95 -0
- keras_hub/src/models/electra/electra_tokenizer.py +104 -0
- keras_hub/src/models/f_net/__init__.py +20 -0
- keras_hub/src/models/f_net/f_net_backbone.py +236 -0
- keras_hub/src/models/f_net/f_net_classifier.py +154 -0
- keras_hub/src/models/f_net/f_net_masked_lm.py +132 -0
- keras_hub/src/models/f_net/f_net_masked_lm_preprocessor.py +196 -0
- keras_hub/src/models/f_net/f_net_preprocessor.py +177 -0
- keras_hub/src/models/f_net/f_net_presets.py +43 -0
- keras_hub/src/models/f_net/f_net_tokenizer.py +95 -0
- keras_hub/src/models/falcon/__init__.py +20 -0
- keras_hub/src/models/falcon/falcon_attention.py +156 -0
- keras_hub/src/models/falcon/falcon_backbone.py +164 -0
- keras_hub/src/models/falcon/falcon_causal_lm.py +291 -0
- keras_hub/src/models/falcon/falcon_causal_lm_preprocessor.py +173 -0
- keras_hub/src/models/falcon/falcon_preprocessor.py +187 -0
- keras_hub/src/models/falcon/falcon_presets.py +30 -0
- keras_hub/src/models/falcon/falcon_tokenizer.py +110 -0
- keras_hub/src/models/falcon/falcon_transformer_decoder.py +255 -0
- keras_hub/src/models/feature_pyramid_backbone.py +73 -0
- keras_hub/src/models/gemma/__init__.py +20 -0
- keras_hub/src/models/gemma/gemma_attention.py +250 -0
- keras_hub/src/models/gemma/gemma_backbone.py +316 -0
- keras_hub/src/models/gemma/gemma_causal_lm.py +448 -0
- keras_hub/src/models/gemma/gemma_causal_lm_preprocessor.py +167 -0
- keras_hub/src/models/gemma/gemma_decoder_block.py +241 -0
- keras_hub/src/models/gemma/gemma_preprocessor.py +191 -0
- keras_hub/src/models/gemma/gemma_presets.py +248 -0
- keras_hub/src/models/gemma/gemma_tokenizer.py +103 -0
- keras_hub/src/models/gemma/rms_normalization.py +40 -0
- keras_hub/src/models/gpt2/__init__.py +20 -0
- keras_hub/src/models/gpt2/gpt2_backbone.py +199 -0
- keras_hub/src/models/gpt2/gpt2_causal_lm.py +437 -0
- keras_hub/src/models/gpt2/gpt2_causal_lm_preprocessor.py +173 -0
- keras_hub/src/models/gpt2/gpt2_preprocessor.py +187 -0
- keras_hub/src/models/gpt2/gpt2_presets.py +82 -0
- keras_hub/src/models/gpt2/gpt2_tokenizer.py +110 -0
- keras_hub/src/models/gpt_neo_x/__init__.py +13 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py +251 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_backbone.py +175 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_causal_lm.py +201 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py +141 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_decoder.py +258 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_preprocessor.py +145 -0
- keras_hub/src/models/gpt_neo_x/gpt_neo_x_tokenizer.py +88 -0
- keras_hub/src/models/image_classifier.py +90 -0
- keras_hub/src/models/llama/__init__.py +20 -0
- keras_hub/src/models/llama/llama_attention.py +225 -0
- keras_hub/src/models/llama/llama_backbone.py +188 -0
- keras_hub/src/models/llama/llama_causal_lm.py +327 -0
- keras_hub/src/models/llama/llama_causal_lm_preprocessor.py +170 -0
- keras_hub/src/models/llama/llama_decoder.py +246 -0
- keras_hub/src/models/llama/llama_layernorm.py +48 -0
- keras_hub/src/models/llama/llama_preprocessor.py +189 -0
- keras_hub/src/models/llama/llama_presets.py +80 -0
- keras_hub/src/models/llama/llama_tokenizer.py +84 -0
- keras_hub/src/models/llama3/__init__.py +20 -0
- keras_hub/src/models/llama3/llama3_backbone.py +84 -0
- keras_hub/src/models/llama3/llama3_causal_lm.py +46 -0
- keras_hub/src/models/llama3/llama3_causal_lm_preprocessor.py +173 -0
- keras_hub/src/models/llama3/llama3_preprocessor.py +21 -0
- keras_hub/src/models/llama3/llama3_presets.py +69 -0
- keras_hub/src/models/llama3/llama3_tokenizer.py +63 -0
- keras_hub/src/models/masked_lm.py +101 -0
- keras_hub/src/models/mistral/__init__.py +20 -0
- keras_hub/src/models/mistral/mistral_attention.py +238 -0
- keras_hub/src/models/mistral/mistral_backbone.py +203 -0
- keras_hub/src/models/mistral/mistral_causal_lm.py +328 -0
- keras_hub/src/models/mistral/mistral_causal_lm_preprocessor.py +175 -0
- keras_hub/src/models/mistral/mistral_layer_norm.py +48 -0
- keras_hub/src/models/mistral/mistral_preprocessor.py +190 -0
- keras_hub/src/models/mistral/mistral_presets.py +48 -0
- keras_hub/src/models/mistral/mistral_tokenizer.py +82 -0
- keras_hub/src/models/mistral/mistral_transformer_decoder.py +265 -0
- keras_hub/src/models/mix_transformer/__init__.py +13 -0
- keras_hub/src/models/mix_transformer/mix_transformer_backbone.py +181 -0
- keras_hub/src/models/mix_transformer/mix_transformer_classifier.py +133 -0
- keras_hub/src/models/mix_transformer/mix_transformer_layers.py +300 -0
- keras_hub/src/models/opt/__init__.py +20 -0
- keras_hub/src/models/opt/opt_backbone.py +173 -0
- keras_hub/src/models/opt/opt_causal_lm.py +301 -0
- keras_hub/src/models/opt/opt_causal_lm_preprocessor.py +177 -0
- keras_hub/src/models/opt/opt_preprocessor.py +188 -0
- keras_hub/src/models/opt/opt_presets.py +72 -0
- keras_hub/src/models/opt/opt_tokenizer.py +116 -0
- keras_hub/src/models/pali_gemma/__init__.py +23 -0
- keras_hub/src/models/pali_gemma/pali_gemma_backbone.py +277 -0
- keras_hub/src/models/pali_gemma/pali_gemma_causal_lm.py +313 -0
- keras_hub/src/models/pali_gemma/pali_gemma_causal_lm_preprocessor.py +147 -0
- keras_hub/src/models/pali_gemma/pali_gemma_decoder_block.py +160 -0
- keras_hub/src/models/pali_gemma/pali_gemma_presets.py +78 -0
- keras_hub/src/models/pali_gemma/pali_gemma_tokenizer.py +79 -0
- keras_hub/src/models/pali_gemma/pali_gemma_vit.py +566 -0
- keras_hub/src/models/phi3/__init__.py +20 -0
- keras_hub/src/models/phi3/phi3_attention.py +260 -0
- keras_hub/src/models/phi3/phi3_backbone.py +224 -0
- keras_hub/src/models/phi3/phi3_causal_lm.py +218 -0
- keras_hub/src/models/phi3/phi3_causal_lm_preprocessor.py +173 -0
- keras_hub/src/models/phi3/phi3_decoder.py +260 -0
- keras_hub/src/models/phi3/phi3_layernorm.py +48 -0
- keras_hub/src/models/phi3/phi3_preprocessor.py +190 -0
- keras_hub/src/models/phi3/phi3_presets.py +50 -0
- keras_hub/src/models/phi3/phi3_rotary_embedding.py +137 -0
- keras_hub/src/models/phi3/phi3_tokenizer.py +94 -0
- keras_hub/src/models/preprocessor.py +207 -0
- keras_hub/src/models/resnet/__init__.py +13 -0
- keras_hub/src/models/resnet/resnet_backbone.py +612 -0
- keras_hub/src/models/resnet/resnet_image_classifier.py +136 -0
- keras_hub/src/models/roberta/__init__.py +20 -0
- keras_hub/src/models/roberta/roberta_backbone.py +184 -0
- keras_hub/src/models/roberta/roberta_classifier.py +209 -0
- keras_hub/src/models/roberta/roberta_masked_lm.py +136 -0
- keras_hub/src/models/roberta/roberta_masked_lm_preprocessor.py +198 -0
- keras_hub/src/models/roberta/roberta_preprocessor.py +192 -0
- keras_hub/src/models/roberta/roberta_presets.py +43 -0
- keras_hub/src/models/roberta/roberta_tokenizer.py +132 -0
- keras_hub/src/models/seq_2_seq_lm.py +54 -0
- keras_hub/src/models/t5/__init__.py +20 -0
- keras_hub/src/models/t5/t5_backbone.py +261 -0
- keras_hub/src/models/t5/t5_layer_norm.py +35 -0
- keras_hub/src/models/t5/t5_multi_head_attention.py +324 -0
- keras_hub/src/models/t5/t5_presets.py +95 -0
- keras_hub/src/models/t5/t5_tokenizer.py +100 -0
- keras_hub/src/models/t5/t5_transformer_layer.py +178 -0
- keras_hub/src/models/task.py +419 -0
- keras_hub/src/models/vgg/__init__.py +13 -0
- keras_hub/src/models/vgg/vgg_backbone.py +158 -0
- keras_hub/src/models/vgg/vgg_image_classifier.py +124 -0
- keras_hub/src/models/vit_det/__init__.py +13 -0
- keras_hub/src/models/vit_det/vit_det_backbone.py +204 -0
- keras_hub/src/models/vit_det/vit_layers.py +565 -0
- keras_hub/src/models/whisper/__init__.py +20 -0
- keras_hub/src/models/whisper/whisper_audio_feature_extractor.py +260 -0
- keras_hub/src/models/whisper/whisper_backbone.py +305 -0
- keras_hub/src/models/whisper/whisper_cached_multi_head_attention.py +153 -0
- keras_hub/src/models/whisper/whisper_decoder.py +141 -0
- keras_hub/src/models/whisper/whisper_encoder.py +106 -0
- keras_hub/src/models/whisper/whisper_preprocessor.py +326 -0
- keras_hub/src/models/whisper/whisper_presets.py +148 -0
- keras_hub/src/models/whisper/whisper_tokenizer.py +163 -0
- keras_hub/src/models/xlm_roberta/__init__.py +26 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_backbone.py +81 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_classifier.py +225 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm.py +141 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm_preprocessor.py +195 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_preprocessor.py +205 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_presets.py +43 -0
- keras_hub/src/models/xlm_roberta/xlm_roberta_tokenizer.py +191 -0
- keras_hub/src/models/xlnet/__init__.py +13 -0
- keras_hub/src/models/xlnet/relative_attention.py +459 -0
- keras_hub/src/models/xlnet/xlnet_backbone.py +222 -0
- keras_hub/src/models/xlnet/xlnet_content_and_query_embedding.py +133 -0
- keras_hub/src/models/xlnet/xlnet_encoder.py +378 -0
- keras_hub/src/samplers/__init__.py +13 -0
- keras_hub/src/samplers/beam_sampler.py +207 -0
- keras_hub/src/samplers/contrastive_sampler.py +231 -0
- keras_hub/src/samplers/greedy_sampler.py +50 -0
- keras_hub/src/samplers/random_sampler.py +77 -0
- keras_hub/src/samplers/sampler.py +237 -0
- keras_hub/src/samplers/serialization.py +97 -0
- keras_hub/src/samplers/top_k_sampler.py +92 -0
- keras_hub/src/samplers/top_p_sampler.py +113 -0
- keras_hub/src/tests/__init__.py +13 -0
- keras_hub/src/tests/test_case.py +608 -0
- keras_hub/src/tokenizers/__init__.py +13 -0
- keras_hub/src/tokenizers/byte_pair_tokenizer.py +638 -0
- keras_hub/src/tokenizers/byte_tokenizer.py +299 -0
- keras_hub/src/tokenizers/sentence_piece_tokenizer.py +267 -0
- keras_hub/src/tokenizers/sentence_piece_tokenizer_trainer.py +150 -0
- keras_hub/src/tokenizers/tokenizer.py +235 -0
- keras_hub/src/tokenizers/unicode_codepoint_tokenizer.py +355 -0
- keras_hub/src/tokenizers/word_piece_tokenizer.py +544 -0
- keras_hub/src/tokenizers/word_piece_tokenizer_trainer.py +176 -0
- keras_hub/src/utils/__init__.py +13 -0
- keras_hub/src/utils/keras_utils.py +130 -0
- keras_hub/src/utils/pipeline_model.py +293 -0
- keras_hub/src/utils/preset_utils.py +621 -0
- keras_hub/src/utils/python_utils.py +21 -0
- keras_hub/src/utils/tensor_utils.py +206 -0
- keras_hub/src/utils/timm/__init__.py +13 -0
- keras_hub/src/utils/timm/convert.py +37 -0
- keras_hub/src/utils/timm/convert_resnet.py +171 -0
- keras_hub/src/utils/transformers/__init__.py +13 -0
- keras_hub/src/utils/transformers/convert.py +101 -0
- keras_hub/src/utils/transformers/convert_bert.py +173 -0
- keras_hub/src/utils/transformers/convert_distilbert.py +184 -0
- keras_hub/src/utils/transformers/convert_gemma.py +187 -0
- keras_hub/src/utils/transformers/convert_gpt2.py +186 -0
- keras_hub/src/utils/transformers/convert_llama3.py +136 -0
- keras_hub/src/utils/transformers/convert_pali_gemma.py +303 -0
- keras_hub/src/utils/transformers/safetensor_utils.py +97 -0
- keras_hub/src/version_utils.py +23 -0
- keras_hub_nightly-0.15.0.dev20240823171555.dist-info/METADATA +34 -0
- keras_hub_nightly-0.15.0.dev20240823171555.dist-info/RECORD +297 -0
- keras_hub_nightly-0.15.0.dev20240823171555.dist-info/WHEEL +5 -0
- keras_hub_nightly-0.15.0.dev20240823171555.dist-info/top_level.txt +1 -0
@@ -0,0 +1,94 @@
|
|
1
|
+
# Copyright 2024 The KerasHub Authors
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
|
15
|
+
import keras
|
16
|
+
from keras import ops
|
17
|
+
|
18
|
+
|
19
|
+
class RelativeEmbedding(keras.layers.Layer):
|
20
|
+
"""Relative embedding layer.
|
21
|
+
|
22
|
+
This is an implementation of relative embedding as described in the
|
23
|
+
paper ["DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing"](https://arxiv.org/abs/2111.09543).
|
24
|
+
This layer initializes an embedding matrix (of shape
|
25
|
+
`(2 * batch_size, hidden_dim)`) for relative position encoding. It then
|
26
|
+
applies layer normalization on the embedding matrix and returns the relative
|
27
|
+
embedding matrix.
|
28
|
+
|
29
|
+
Args:
|
30
|
+
hidden_dim: int. The size of the dense embedding.
|
31
|
+
bucket_size: int. The size of the relative position buckets.
|
32
|
+
layer_norm_epsilon: float. Epsilon value to initialize the layer
|
33
|
+
normalization layer.
|
34
|
+
kernel_initializer: string or `keras.initializers` initializer.
|
35
|
+
The kernel initializer for the dense embedding.
|
36
|
+
Defaults to `"glorot_uniform"`.
|
37
|
+
"""
|
38
|
+
|
39
|
+
def __init__(
|
40
|
+
self,
|
41
|
+
hidden_dim,
|
42
|
+
bucket_size,
|
43
|
+
layer_norm_epsilon=1e-05,
|
44
|
+
kernel_initializer="glorot_uniform",
|
45
|
+
**kwargs,
|
46
|
+
):
|
47
|
+
super().__init__(**kwargs)
|
48
|
+
|
49
|
+
self.hidden_dim = hidden_dim
|
50
|
+
self.bucket_size = bucket_size
|
51
|
+
self.layer_norm_epsilon = layer_norm_epsilon
|
52
|
+
self.kernel_initializer = keras.initializers.get(kernel_initializer)
|
53
|
+
|
54
|
+
self.rel_embeddings = self.add_weight(
|
55
|
+
shape=(self.bucket_size * 2, self.hidden_dim),
|
56
|
+
initializer=self.kernel_initializer,
|
57
|
+
name="rel_embedding",
|
58
|
+
)
|
59
|
+
self.layer_norm = keras.layers.LayerNormalization(
|
60
|
+
epsilon=layer_norm_epsilon,
|
61
|
+
dtype=self.dtype_policy,
|
62
|
+
name="rel_embeddings_layer_norm",
|
63
|
+
)
|
64
|
+
|
65
|
+
def call(self, inputs):
|
66
|
+
batch_size = ops.shape(inputs)[0]
|
67
|
+
|
68
|
+
rel_embeddings = ops.expand_dims(
|
69
|
+
ops.convert_to_tensor(self.rel_embeddings), axis=0
|
70
|
+
)
|
71
|
+
rel_embeddings = self.layer_norm(rel_embeddings)
|
72
|
+
|
73
|
+
# Repeat `rel_embeddings` along axis = 0 `batch_size` times. The
|
74
|
+
# resultant shape is `(batch_size, bucket_size * 2, hidden_dim)`.
|
75
|
+
rel_embeddings = ops.repeat(rel_embeddings, repeats=batch_size, axis=0)
|
76
|
+
|
77
|
+
return rel_embeddings
|
78
|
+
|
79
|
+
def get_config(self):
|
80
|
+
config = super().get_config()
|
81
|
+
config.update(
|
82
|
+
{
|
83
|
+
"hidden_dim": self.hidden_dim,
|
84
|
+
"bucket_size": self.bucket_size,
|
85
|
+
"layer_norm_epsilon": self.layer_norm_epsilon,
|
86
|
+
"kernel_initializer": keras.initializers.serialize(
|
87
|
+
self.kernel_initializer
|
88
|
+
),
|
89
|
+
}
|
90
|
+
)
|
91
|
+
return config
|
92
|
+
|
93
|
+
def compute_output_shape(self, input_shape):
|
94
|
+
return (input_shape[0],) + (self.bucket_size * 2, self.hidden_dim)
|
@@ -0,0 +1,13 @@
|
|
1
|
+
# Copyright 2024 The KerasHub Authors
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
@@ -0,0 +1,210 @@
|
|
1
|
+
# Copyright 2024 The KerasHub Authors
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
import keras
|
15
|
+
|
16
|
+
from keras_hub.src.api_export import keras_hub_export
|
17
|
+
from keras_hub.src.models.backbone import Backbone
|
18
|
+
|
19
|
+
BN_AXIS = 3
|
20
|
+
BN_EPSILON = 1.001e-5
|
21
|
+
|
22
|
+
|
23
|
+
@keras_hub_export("keras_hub.models.DenseNetBackbone")
|
24
|
+
class DenseNetBackbone(Backbone):
|
25
|
+
"""Instantiates the DenseNet architecture.
|
26
|
+
|
27
|
+
This class implements a DenseNet backbone as described in
|
28
|
+
[Densely Connected Convolutional Networks (CVPR 2017)](
|
29
|
+
https://arxiv.org/abs/1608.06993
|
30
|
+
).
|
31
|
+
|
32
|
+
Args:
|
33
|
+
stackwise_num_repeats: list of ints, number of repeated convolutional
|
34
|
+
blocks per dense block.
|
35
|
+
include_rescaling: bool, whether to rescale the inputs. If set
|
36
|
+
to `True`, inputs will be passed through a `Rescaling(1/255.0)`
|
37
|
+
layer. Defaults to `True`.
|
38
|
+
image_shape: optional shape tuple, defaults to (224, 224, 3).
|
39
|
+
compression_ratio: float, compression rate at transition layers,
|
40
|
+
defaults to 0.5.
|
41
|
+
growth_rate: int, number of filters added by each dense block,
|
42
|
+
defaults to 32
|
43
|
+
|
44
|
+
Examples:
|
45
|
+
```python
|
46
|
+
input_data = np.ones(shape=(8, 224, 224, 3))
|
47
|
+
|
48
|
+
# Pretrained backbone
|
49
|
+
model = keras_hub.models.DenseNetBackbone.from_preset("densenet121_imagenet")
|
50
|
+
model(input_data)
|
51
|
+
|
52
|
+
# Randomly initialized backbone with a custom config
|
53
|
+
model = keras_hub.models.DenseNetBackbone(
|
54
|
+
stackwise_num_repeats=[6, 12, 24, 16],
|
55
|
+
include_rescaling=False,
|
56
|
+
)
|
57
|
+
model(input_data)
|
58
|
+
```
|
59
|
+
"""
|
60
|
+
|
61
|
+
def __init__(
|
62
|
+
self,
|
63
|
+
stackwise_num_repeats,
|
64
|
+
include_rescaling=True,
|
65
|
+
image_shape=(224, 224, 3),
|
66
|
+
compression_ratio=0.5,
|
67
|
+
growth_rate=32,
|
68
|
+
**kwargs,
|
69
|
+
):
|
70
|
+
# === Functional Model ===
|
71
|
+
image_input = keras.layers.Input(shape=image_shape)
|
72
|
+
|
73
|
+
x = image_input
|
74
|
+
if include_rescaling:
|
75
|
+
x = keras.layers.Rescaling(1 / 255.0)(x)
|
76
|
+
|
77
|
+
x = keras.layers.Conv2D(
|
78
|
+
64, 7, strides=2, use_bias=False, padding="same", name="conv1_conv"
|
79
|
+
)(x)
|
80
|
+
x = keras.layers.BatchNormalization(
|
81
|
+
axis=BN_AXIS, epsilon=BN_EPSILON, name="conv1_bn"
|
82
|
+
)(x)
|
83
|
+
x = keras.layers.Activation("relu", name="conv1_relu")(x)
|
84
|
+
x = keras.layers.MaxPooling2D(
|
85
|
+
3, strides=2, padding="same", name="pool1"
|
86
|
+
)(x)
|
87
|
+
|
88
|
+
for stack_index in range(len(stackwise_num_repeats) - 1):
|
89
|
+
index = stack_index + 2
|
90
|
+
x = apply_dense_block(
|
91
|
+
x,
|
92
|
+
stackwise_num_repeats[stack_index],
|
93
|
+
growth_rate,
|
94
|
+
name=f"conv{index}",
|
95
|
+
)
|
96
|
+
x = apply_transition_block(
|
97
|
+
x, compression_ratio, name=f"pool{index}"
|
98
|
+
)
|
99
|
+
|
100
|
+
x = apply_dense_block(
|
101
|
+
x,
|
102
|
+
stackwise_num_repeats[-1],
|
103
|
+
growth_rate,
|
104
|
+
name=f"conv{len(stackwise_num_repeats) + 1}",
|
105
|
+
)
|
106
|
+
|
107
|
+
x = keras.layers.BatchNormalization(
|
108
|
+
axis=BN_AXIS, epsilon=BN_EPSILON, name="bn"
|
109
|
+
)(x)
|
110
|
+
x = keras.layers.Activation("relu", name="relu")(x)
|
111
|
+
|
112
|
+
super().__init__(inputs=image_input, outputs=x, **kwargs)
|
113
|
+
|
114
|
+
# === Config ===
|
115
|
+
self.stackwise_num_repeats = stackwise_num_repeats
|
116
|
+
self.include_rescaling = include_rescaling
|
117
|
+
self.compression_ratio = compression_ratio
|
118
|
+
self.growth_rate = growth_rate
|
119
|
+
self.image_shape = image_shape
|
120
|
+
|
121
|
+
def get_config(self):
|
122
|
+
config = super().get_config()
|
123
|
+
config.update(
|
124
|
+
{
|
125
|
+
"stackwise_num_repeats": self.stackwise_num_repeats,
|
126
|
+
"include_rescaling": self.include_rescaling,
|
127
|
+
"compression_ratio": self.compression_ratio,
|
128
|
+
"growth_rate": self.growth_rate,
|
129
|
+
"image_shape": self.image_shape,
|
130
|
+
}
|
131
|
+
)
|
132
|
+
return config
|
133
|
+
|
134
|
+
|
135
|
+
def apply_dense_block(x, num_repeats, growth_rate, name=None):
|
136
|
+
"""A dense block.
|
137
|
+
|
138
|
+
Args:
|
139
|
+
x: input tensor.
|
140
|
+
num_repeats: int, number of repeated convolutional blocks.
|
141
|
+
growth_rate: int, number of filters added by each dense block.
|
142
|
+
name: string, block label.
|
143
|
+
"""
|
144
|
+
if name is None:
|
145
|
+
name = f"dense_block_{keras.backend.get_uid('dense_block')}"
|
146
|
+
|
147
|
+
for i in range(num_repeats):
|
148
|
+
x = apply_conv_block(x, growth_rate, name=f"{name}_block_{i}")
|
149
|
+
return x
|
150
|
+
|
151
|
+
|
152
|
+
def apply_transition_block(x, compression_ratio, name=None):
|
153
|
+
"""A transition block.
|
154
|
+
|
155
|
+
Args:
|
156
|
+
x: input tensor.
|
157
|
+
compression_ratio: float, compression rate at transition layers.
|
158
|
+
name: string, block label.
|
159
|
+
"""
|
160
|
+
if name is None:
|
161
|
+
name = f"transition_block_{keras.backend.get_uid('transition_block')}"
|
162
|
+
|
163
|
+
x = keras.layers.BatchNormalization(
|
164
|
+
axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_bn"
|
165
|
+
)(x)
|
166
|
+
x = keras.layers.Activation("relu", name=f"{name}_relu")(x)
|
167
|
+
x = keras.layers.Conv2D(
|
168
|
+
int(x.shape[BN_AXIS] * compression_ratio),
|
169
|
+
1,
|
170
|
+
use_bias=False,
|
171
|
+
name=f"{name}_conv",
|
172
|
+
)(x)
|
173
|
+
x = keras.layers.AveragePooling2D(2, strides=2, name=f"{name}_pool")(x)
|
174
|
+
return x
|
175
|
+
|
176
|
+
|
177
|
+
def apply_conv_block(x, growth_rate, name=None):
|
178
|
+
"""A building block for a dense block.
|
179
|
+
|
180
|
+
Args:
|
181
|
+
x: input tensor.
|
182
|
+
growth_rate: int, number of filters added by each dense block.
|
183
|
+
name: string, block label.
|
184
|
+
"""
|
185
|
+
if name is None:
|
186
|
+
name = f"conv_block_{keras.backend.get_uid('conv_block')}"
|
187
|
+
|
188
|
+
shortcut = x
|
189
|
+
x = keras.layers.BatchNormalization(
|
190
|
+
axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_0_bn"
|
191
|
+
)(x)
|
192
|
+
x = keras.layers.Activation("relu", name=f"{name}_0_relu")(x)
|
193
|
+
x = keras.layers.Conv2D(
|
194
|
+
4 * growth_rate, 1, use_bias=False, name=f"{name}_1_conv"
|
195
|
+
)(x)
|
196
|
+
x = keras.layers.BatchNormalization(
|
197
|
+
axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_1_bn"
|
198
|
+
)(x)
|
199
|
+
x = keras.layers.Activation("relu", name=f"{name}_1_relu")(x)
|
200
|
+
x = keras.layers.Conv2D(
|
201
|
+
growth_rate,
|
202
|
+
3,
|
203
|
+
padding="same",
|
204
|
+
use_bias=False,
|
205
|
+
name=f"{name}_2_conv",
|
206
|
+
)(x)
|
207
|
+
x = keras.layers.Concatenate(axis=BN_AXIS, name=f"{name}_concat")(
|
208
|
+
[shortcut, x]
|
209
|
+
)
|
210
|
+
return x
|
@@ -0,0 +1,131 @@
|
|
1
|
+
# Copyright 2024 The KerasHub Authors
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
import keras
|
15
|
+
|
16
|
+
from keras_hub.src.api_export import keras_hub_export
|
17
|
+
from keras_hub.src.models.densenet.densenet_backbone import DenseNetBackbone
|
18
|
+
from keras_hub.src.models.image_classifier import ImageClassifier
|
19
|
+
|
20
|
+
|
21
|
+
@keras_hub_export("keras_hub.models.DenseNetImageClassifier")
|
22
|
+
class DenseNetImageClassifier(ImageClassifier):
|
23
|
+
"""DenseNet image classifier task model.
|
24
|
+
|
25
|
+
To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
|
26
|
+
where `x` is a tensor and `y` is a integer from `[0, num_classes)`.
|
27
|
+
All `ImageClassifier` tasks include a `from_preset()` constructor which can
|
28
|
+
be used to load a pre-trained config and weights.
|
29
|
+
|
30
|
+
Args:
|
31
|
+
backbone: A `keras_hub.models.DenseNetBackbone` instance.
|
32
|
+
num_classes: int. The number of classes to predict.
|
33
|
+
activation: `None`, str or callable. The activation function to use on
|
34
|
+
the `Dense` layer. Set `activation=None` to return the output
|
35
|
+
logits. Defaults to `"softmax"`.
|
36
|
+
|
37
|
+
Examples:
|
38
|
+
|
39
|
+
Call `predict()` to run inference.
|
40
|
+
```python
|
41
|
+
# Load preset and train
|
42
|
+
images = np.ones((2, 224, 224, 3), dtype="float32")
|
43
|
+
classifier = keras_hub.models.DenseNetImageClassifier.from_preset(
|
44
|
+
"densenet121_imagenet")
|
45
|
+
classifier.predict(images)
|
46
|
+
```
|
47
|
+
|
48
|
+
Call `fit()` on a single batch.
|
49
|
+
```python
|
50
|
+
# Load preset and train
|
51
|
+
images = np.ones((2, 224, 224, 3), dtype="float32")
|
52
|
+
labels = [0, 3]
|
53
|
+
classifier = keras_hub.models.DenseNetImageClassifier.from_preset(
|
54
|
+
"densenet121_imagenet")
|
55
|
+
classifier.fit(x=images, y=labels, batch_size=2)
|
56
|
+
```
|
57
|
+
|
58
|
+
Call `fit()` with custom loss, optimizer and backbone.
|
59
|
+
```python
|
60
|
+
classifier = keras_hub.models.DenseNetImageClassifier.from_preset(
|
61
|
+
"densenet121_imagenet")
|
62
|
+
classifier.compile(
|
63
|
+
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
|
64
|
+
optimizer=keras.optimizers.Adam(5e-5),
|
65
|
+
)
|
66
|
+
classifier.backbone.trainable = False
|
67
|
+
classifier.fit(x=images, y=labels, batch_size=2)
|
68
|
+
```
|
69
|
+
|
70
|
+
Custom backbone.
|
71
|
+
```python
|
72
|
+
images = np.ones((2, 224, 224, 3), dtype="float32")
|
73
|
+
labels = [0, 3]
|
74
|
+
backbone = keras_hub.models.DenseNetBackbone(
|
75
|
+
stackwise_num_filters=[128, 256, 512, 1024],
|
76
|
+
stackwise_depth=[3, 9, 9, 3],
|
77
|
+
include_rescaling=False,
|
78
|
+
block_type="basic_block",
|
79
|
+
image_shape = (224, 224, 3),
|
80
|
+
)
|
81
|
+
classifier = keras_hub.models.DenseNetImageClassifier(
|
82
|
+
backbone=backbone,
|
83
|
+
num_classes=4,
|
84
|
+
)
|
85
|
+
classifier.fit(x=images, y=labels, batch_size=2)
|
86
|
+
```
|
87
|
+
"""
|
88
|
+
|
89
|
+
backbone_cls = DenseNetBackbone
|
90
|
+
|
91
|
+
def __init__(
|
92
|
+
self,
|
93
|
+
backbone,
|
94
|
+
num_classes,
|
95
|
+
activation="softmax",
|
96
|
+
preprocessor=None, # adding this dummy arg for saved model test
|
97
|
+
# TODO: once preprocessor flow is figured out, this needs to be updated
|
98
|
+
**kwargs,
|
99
|
+
):
|
100
|
+
# === Layers ===
|
101
|
+
self.backbone = backbone
|
102
|
+
self.output_dense = keras.layers.Dense(
|
103
|
+
num_classes,
|
104
|
+
activation=activation,
|
105
|
+
name="predictions",
|
106
|
+
)
|
107
|
+
|
108
|
+
# === Functional Model ===
|
109
|
+
inputs = self.backbone.input
|
110
|
+
x = self.backbone(inputs)
|
111
|
+
outputs = self.output_dense(x)
|
112
|
+
super().__init__(
|
113
|
+
inputs=inputs,
|
114
|
+
outputs=outputs,
|
115
|
+
**kwargs,
|
116
|
+
)
|
117
|
+
|
118
|
+
# === Config ===
|
119
|
+
self.num_classes = num_classes
|
120
|
+
self.activation = activation
|
121
|
+
|
122
|
+
def get_config(self):
|
123
|
+
# Backbone serialized in `super`
|
124
|
+
config = super().get_config()
|
125
|
+
config.update(
|
126
|
+
{
|
127
|
+
"num_classes": self.num_classes,
|
128
|
+
"activation": self.activation,
|
129
|
+
}
|
130
|
+
)
|
131
|
+
return config
|
@@ -0,0 +1,26 @@
|
|
1
|
+
# Copyright 2024 The KerasHub Authors
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
|
15
|
+
from keras_hub.src.models.distil_bert.distil_bert_backbone import (
|
16
|
+
DistilBertBackbone,
|
17
|
+
)
|
18
|
+
from keras_hub.src.models.distil_bert.distil_bert_presets import (
|
19
|
+
backbone_presets,
|
20
|
+
)
|
21
|
+
from keras_hub.src.models.distil_bert.distil_bert_tokenizer import (
|
22
|
+
DistilBertTokenizer,
|
23
|
+
)
|
24
|
+
from keras_hub.src.utils.preset_utils import register_presets
|
25
|
+
|
26
|
+
register_presets(backbone_presets, (DistilBertBackbone, DistilBertTokenizer))
|
@@ -0,0 +1,187 @@
|
|
1
|
+
# Copyright 2024 The KerasHub Authors
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
|
15
|
+
|
16
|
+
import keras
|
17
|
+
|
18
|
+
from keras_hub.src.api_export import keras_hub_export
|
19
|
+
from keras_hub.src.layers.modeling.token_and_position_embedding import (
|
20
|
+
TokenAndPositionEmbedding,
|
21
|
+
)
|
22
|
+
from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
|
23
|
+
from keras_hub.src.models.backbone import Backbone
|
24
|
+
|
25
|
+
|
26
|
+
def distilbert_kernel_initializer(stddev=0.02):
|
27
|
+
return keras.initializers.TruncatedNormal(stddev=stddev)
|
28
|
+
|
29
|
+
|
30
|
+
@keras_hub_export("keras_hub.models.DistilBertBackbone")
|
31
|
+
class DistilBertBackbone(Backbone):
|
32
|
+
"""A DistilBERT encoder network.
|
33
|
+
|
34
|
+
This network implements a bi-directional Transformer-based encoder as
|
35
|
+
described in ["DistilBERT, a distilled version of BERT: smaller, faster,
|
36
|
+
cheaper and lighter"](https://arxiv.org/abs/1910.01108). It includes the
|
37
|
+
embedding lookups and transformer layers, but not the masked language model
|
38
|
+
or classification task networks.
|
39
|
+
|
40
|
+
The default constructor gives a fully customizable, randomly initialized
|
41
|
+
DistilBERT encoder with any number of layers, heads, and embedding
|
42
|
+
dimensions. To load preset architectures and weights, use the
|
43
|
+
`from_preset()` constructor.
|
44
|
+
|
45
|
+
Disclaimer: Pre-trained models are provided on an "as is" basis, without
|
46
|
+
warranties or conditions of any kind. The underlying model is provided by a
|
47
|
+
third party and subject to a separate license, available
|
48
|
+
[here](https://github.com/huggingface/transformers).
|
49
|
+
|
50
|
+
Args:
|
51
|
+
vocabulary_size: int. The size of the token vocabulary.
|
52
|
+
num_layers: int. The number of transformer layers.
|
53
|
+
num_heads: int. The number of attention heads for each transformer.
|
54
|
+
The hidden size must be divisible by the number of attention heads.
|
55
|
+
hidden_dim: int. The size of the transformer encoding and pooler layers.
|
56
|
+
intermediate_dim: int. The output dimension of the first Dense layer in
|
57
|
+
a two-layer feedforward network for each transformer.
|
58
|
+
dropout: float. Dropout probability for the Transformer encoder.
|
59
|
+
max_sequence_length: int. The maximum sequence length that this encoder
|
60
|
+
can consume. If None, `max_sequence_length` uses the value from
|
61
|
+
sequence length. This determines the variable shape for positional
|
62
|
+
embeddings.
|
63
|
+
dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
|
64
|
+
for model computations and weights. Note that some computations,
|
65
|
+
such as softmax and layer normalization, will always be done at
|
66
|
+
float32 precision regardless of dtype.
|
67
|
+
|
68
|
+
Examples:
|
69
|
+
```python
|
70
|
+
input_data = {
|
71
|
+
"token_ids": np.ones(shape=(1, 12), dtype="int32"),
|
72
|
+
"padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
|
73
|
+
}
|
74
|
+
|
75
|
+
# Pretrained DistilBERT encoder.
|
76
|
+
model = keras_hub.models.DistilBertBackbone.from_preset(
|
77
|
+
"distil_bert_base_en_uncased"
|
78
|
+
)
|
79
|
+
model(input_data)
|
80
|
+
|
81
|
+
# Randomly initialized DistilBERT encoder with custom config.
|
82
|
+
model = keras_hub.models.DistilBertBackbone(
|
83
|
+
vocabulary_size=30552,
|
84
|
+
num_layers=4,
|
85
|
+
num_heads=4,
|
86
|
+
hidden_dim=256,
|
87
|
+
intermediate_dim=512,
|
88
|
+
max_sequence_length=128,
|
89
|
+
)
|
90
|
+
model(input_data)
|
91
|
+
```
|
92
|
+
"""
|
93
|
+
|
94
|
+
def __init__(
|
95
|
+
self,
|
96
|
+
vocabulary_size,
|
97
|
+
num_layers,
|
98
|
+
num_heads,
|
99
|
+
hidden_dim,
|
100
|
+
intermediate_dim,
|
101
|
+
dropout=0.1,
|
102
|
+
max_sequence_length=512,
|
103
|
+
dtype=None,
|
104
|
+
**kwargs,
|
105
|
+
):
|
106
|
+
# === Layers ===
|
107
|
+
self.embeddings = TokenAndPositionEmbedding(
|
108
|
+
vocabulary_size=vocabulary_size,
|
109
|
+
sequence_length=max_sequence_length,
|
110
|
+
embedding_dim=hidden_dim,
|
111
|
+
embeddings_initializer=distilbert_kernel_initializer(),
|
112
|
+
dtype=dtype,
|
113
|
+
name="token_and_position_embedding",
|
114
|
+
)
|
115
|
+
# Keep the token_embedding property for consistency across models.
|
116
|
+
self.token_embedding = self.embeddings.token_embedding
|
117
|
+
self.embeddings_layer_norm = keras.layers.LayerNormalization(
|
118
|
+
axis=-1,
|
119
|
+
epsilon=1e-12,
|
120
|
+
dtype=dtype,
|
121
|
+
name="embeddings_layer_norm",
|
122
|
+
)
|
123
|
+
self.embeddings_dropout = keras.layers.Dropout(
|
124
|
+
dropout,
|
125
|
+
dtype=dtype,
|
126
|
+
name="embeddings_dropout",
|
127
|
+
)
|
128
|
+
self.transformer_layers = []
|
129
|
+
for i in range(num_layers):
|
130
|
+
layer = TransformerEncoder(
|
131
|
+
num_heads=num_heads,
|
132
|
+
intermediate_dim=intermediate_dim,
|
133
|
+
activation="gelu",
|
134
|
+
dropout=dropout,
|
135
|
+
layer_norm_epsilon=1e-12,
|
136
|
+
kernel_initializer=distilbert_kernel_initializer(),
|
137
|
+
dtype=dtype,
|
138
|
+
name=f"transformer_layer_{i}",
|
139
|
+
)
|
140
|
+
self.transformer_layers.append(layer)
|
141
|
+
|
142
|
+
# === Functional Model ===
|
143
|
+
token_id_input = keras.Input(
|
144
|
+
shape=(None,), dtype="int32", name="token_ids"
|
145
|
+
)
|
146
|
+
padding_mask_input = keras.Input(
|
147
|
+
shape=(None,), dtype="int32", name="padding_mask"
|
148
|
+
)
|
149
|
+
x = self.embeddings(token_id_input)
|
150
|
+
x = self.embeddings_layer_norm(x)
|
151
|
+
x = self.embeddings_dropout(x)
|
152
|
+
for transformer_layer in self.transformer_layers:
|
153
|
+
x = transformer_layer(x, padding_mask=padding_mask_input)
|
154
|
+
super().__init__(
|
155
|
+
inputs={
|
156
|
+
"token_ids": token_id_input,
|
157
|
+
"padding_mask": padding_mask_input,
|
158
|
+
},
|
159
|
+
outputs=x,
|
160
|
+
dtype=dtype,
|
161
|
+
**kwargs,
|
162
|
+
)
|
163
|
+
|
164
|
+
# === Config ===
|
165
|
+
self.vocabulary_size = vocabulary_size
|
166
|
+
self.num_layers = num_layers
|
167
|
+
self.num_heads = num_heads
|
168
|
+
self.hidden_dim = hidden_dim
|
169
|
+
self.intermediate_dim = intermediate_dim
|
170
|
+
self.dropout = dropout
|
171
|
+
self.max_sequence_length = max_sequence_length
|
172
|
+
self.cls_token_index = 0
|
173
|
+
|
174
|
+
def get_config(self):
|
175
|
+
config = super().get_config()
|
176
|
+
config.update(
|
177
|
+
{
|
178
|
+
"vocabulary_size": self.vocabulary_size,
|
179
|
+
"num_layers": self.num_layers,
|
180
|
+
"num_heads": self.num_heads,
|
181
|
+
"hidden_dim": self.hidden_dim,
|
182
|
+
"intermediate_dim": self.intermediate_dim,
|
183
|
+
"dropout": self.dropout,
|
184
|
+
"max_sequence_length": self.max_sequence_length,
|
185
|
+
}
|
186
|
+
)
|
187
|
+
return config
|