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,260 @@
|
|
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
|
+
from keras import ops
|
16
|
+
|
17
|
+
from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
|
18
|
+
from keras_hub.src.models.phi3.phi3_rotary_embedding import (
|
19
|
+
Phi3SuScaledRotaryEmbedding,
|
20
|
+
)
|
21
|
+
from keras_hub.src.utils.keras_utils import clone_initializer
|
22
|
+
|
23
|
+
|
24
|
+
class Phi3Attention(keras.layers.Layer):
|
25
|
+
"""A cached grounded query attention layer."""
|
26
|
+
|
27
|
+
def __init__(
|
28
|
+
self,
|
29
|
+
num_query_heads,
|
30
|
+
num_key_value_heads,
|
31
|
+
kernel_initializer="glorot_uniform",
|
32
|
+
dropout=0,
|
33
|
+
max_sequence_length=4096,
|
34
|
+
pretraining_sequence_length=4096,
|
35
|
+
rope_max_wavelength=10000,
|
36
|
+
rope_scaling_type=None,
|
37
|
+
rope_scaling_short_factor=None,
|
38
|
+
rope_scaling_long_factor=None,
|
39
|
+
**kwargs,
|
40
|
+
):
|
41
|
+
super().__init__(**kwargs)
|
42
|
+
self.num_query_heads = num_query_heads
|
43
|
+
self.num_key_value_heads = num_key_value_heads
|
44
|
+
self.num_key_value_groups = num_query_heads // num_key_value_heads
|
45
|
+
self.dropout = dropout
|
46
|
+
|
47
|
+
self.max_sequence_length = max_sequence_length
|
48
|
+
self.pretraining_sequence_length = pretraining_sequence_length
|
49
|
+
self.rope_max_wavelength = rope_max_wavelength
|
50
|
+
self.rope_scaling_type = rope_scaling_type
|
51
|
+
self.rope_scaling_short_factor = rope_scaling_short_factor
|
52
|
+
self.rope_scaling_long_factor = rope_scaling_long_factor
|
53
|
+
|
54
|
+
self.kernel_initializer = keras.initializers.get(
|
55
|
+
clone_initializer(kernel_initializer)
|
56
|
+
)
|
57
|
+
|
58
|
+
def build(self, inputs_shape):
|
59
|
+
# Einsum variables:
|
60
|
+
# b = batch size
|
61
|
+
# q = query length
|
62
|
+
# k = key/value length
|
63
|
+
# m = model dim
|
64
|
+
# u = num query heads
|
65
|
+
# v = num key/value heads
|
66
|
+
# h = head dim
|
67
|
+
hidden_dim = inputs_shape[-1]
|
68
|
+
head_dim = hidden_dim // self.num_query_heads
|
69
|
+
self._norm_factor = ops.sqrt(ops.cast(head_dim, self.compute_dtype))
|
70
|
+
|
71
|
+
self.query_dense = keras.layers.EinsumDense(
|
72
|
+
equation="bqm,muh->bquh",
|
73
|
+
output_shape=(None, self.num_query_heads, head_dim),
|
74
|
+
kernel_initializer=self.kernel_initializer,
|
75
|
+
dtype=self.dtype_policy,
|
76
|
+
name="query",
|
77
|
+
)
|
78
|
+
self.query_dense.build(inputs_shape)
|
79
|
+
|
80
|
+
self.key_dense = keras.layers.EinsumDense(
|
81
|
+
equation="bkm,mvh->bkvh",
|
82
|
+
output_shape=(
|
83
|
+
None,
|
84
|
+
self.num_key_value_heads,
|
85
|
+
head_dim,
|
86
|
+
),
|
87
|
+
kernel_initializer=self.kernel_initializer,
|
88
|
+
dtype=self.dtype_policy,
|
89
|
+
name="key",
|
90
|
+
)
|
91
|
+
self.key_dense.build(inputs_shape)
|
92
|
+
|
93
|
+
self.value_dense = keras.layers.EinsumDense(
|
94
|
+
equation="bkm,mvh->bkvh",
|
95
|
+
output_shape=(
|
96
|
+
None,
|
97
|
+
self.num_key_value_heads,
|
98
|
+
head_dim,
|
99
|
+
),
|
100
|
+
kernel_initializer=self.kernel_initializer,
|
101
|
+
dtype=self.dtype_policy,
|
102
|
+
name="value",
|
103
|
+
)
|
104
|
+
self.value_dense.build(inputs_shape)
|
105
|
+
|
106
|
+
self.softmax = keras.layers.Softmax(
|
107
|
+
axis=-1,
|
108
|
+
dtype="float32",
|
109
|
+
name="attention_softmax",
|
110
|
+
)
|
111
|
+
|
112
|
+
self.dropout_layer = keras.layers.Dropout(
|
113
|
+
rate=self.dropout,
|
114
|
+
dtype=self.dtype_policy,
|
115
|
+
)
|
116
|
+
|
117
|
+
self.output_dense = keras.layers.EinsumDense(
|
118
|
+
equation="bquh,uhm->bqm",
|
119
|
+
output_shape=(None, hidden_dim),
|
120
|
+
kernel_initializer=self.kernel_initializer,
|
121
|
+
dtype=self.dtype_policy,
|
122
|
+
name="attention_output",
|
123
|
+
)
|
124
|
+
self.output_dense.build((None, None, self.num_query_heads, head_dim))
|
125
|
+
|
126
|
+
if self.rope_scaling_type is None:
|
127
|
+
self.rotary_embedding_layer = RotaryEmbedding(
|
128
|
+
max_wavelength=self.rope_max_wavelength,
|
129
|
+
dtype=self.dtype_policy,
|
130
|
+
)
|
131
|
+
elif self.rope_scaling_type == "su":
|
132
|
+
if len(self.rope_scaling_short_factor) != head_dim // 2:
|
133
|
+
raise ValueError(
|
134
|
+
"`rope_scaling_short_factor` must be of length "
|
135
|
+
"`hidden_dim//num_query_heads//2`. "
|
136
|
+
"`len(rope_scaling_short_factor)` is "
|
137
|
+
f"{len(self.rope_scaling_short_factor)} "
|
138
|
+
f"while it should be {head_dim // 2}."
|
139
|
+
)
|
140
|
+
if len(self.rope_scaling_long_factor) != head_dim // 2:
|
141
|
+
raise ValueError(
|
142
|
+
"`rope_scaling_long_factor` must be of length "
|
143
|
+
"`hidden_dim//num_query_heads//2`. "
|
144
|
+
"`len(rope_scaling_long_factor)` is "
|
145
|
+
f"{len(self.rope_scaling_long_factor)} "
|
146
|
+
f"while it should be {head_dim // 2}."
|
147
|
+
)
|
148
|
+
self.rotary_embedding_layer = Phi3SuScaledRotaryEmbedding(
|
149
|
+
inverese_freq_short_factor=self.rope_scaling_short_factor,
|
150
|
+
inverese_freq_long_factor=self.rope_scaling_long_factor,
|
151
|
+
max_sequence_length=self.max_sequence_length,
|
152
|
+
pretraining_sequence_length=self.pretraining_sequence_length,
|
153
|
+
max_wavelength=self.rope_max_wavelength,
|
154
|
+
dtype=self.dtype_policy,
|
155
|
+
)
|
156
|
+
else:
|
157
|
+
raise ValueError(
|
158
|
+
'`rope_scaling_type` must be `None` or `"su"`.'
|
159
|
+
"if `None` is choosed, `RotaryEmbedding` will be used."
|
160
|
+
'if `"su"` is choosed, `Phi3SuScaledRotaryEmbedding` will be '
|
161
|
+
"used."
|
162
|
+
)
|
163
|
+
|
164
|
+
self.built = True
|
165
|
+
|
166
|
+
def call(
|
167
|
+
self,
|
168
|
+
hidden_states,
|
169
|
+
attention_mask=None,
|
170
|
+
cache=None,
|
171
|
+
cache_update_index=None,
|
172
|
+
training=None,
|
173
|
+
):
|
174
|
+
start_index = (
|
175
|
+
cache_update_index if cache_update_index is not None else 0
|
176
|
+
)
|
177
|
+
|
178
|
+
query = self.query_dense(hidden_states)
|
179
|
+
key = self.key_dense(hidden_states)
|
180
|
+
value = self.value_dense(hidden_states)
|
181
|
+
|
182
|
+
# Compute RoPE for queries
|
183
|
+
query = self.rotary_embedding_layer(query, start_index=start_index)
|
184
|
+
key = self.rotary_embedding_layer(key, start_index=start_index)
|
185
|
+
|
186
|
+
if cache is not None:
|
187
|
+
key_cache = cache[:, 0, ...]
|
188
|
+
value_cache = cache[:, 1, ...]
|
189
|
+
if cache_update_index is None:
|
190
|
+
key = key_cache
|
191
|
+
value = value_cache
|
192
|
+
else:
|
193
|
+
start = [0, cache_update_index, 0, 0]
|
194
|
+
key = ops.slice_update(key_cache, start, key)
|
195
|
+
value = ops.slice_update(value_cache, start, value)
|
196
|
+
cache = ops.stack((key, value), axis=1)
|
197
|
+
else:
|
198
|
+
if cache_update_index is not None:
|
199
|
+
raise ValueError(
|
200
|
+
"`cache_update_index` should not be set if `cache` is "
|
201
|
+
f"`None`. Received: cache={cache}, "
|
202
|
+
f"cache_update_index={cache_update_index}"
|
203
|
+
)
|
204
|
+
|
205
|
+
# [batch_shape, seq_len, num_key_value_heads, head_dim]
|
206
|
+
# -> [batch_shape, seq_len, num_heads, head_dim]
|
207
|
+
key = ops.repeat(key, repeats=self.num_key_value_groups, axis=2)
|
208
|
+
value = ops.repeat(value, repeats=self.num_key_value_groups, axis=2)
|
209
|
+
|
210
|
+
attention_output = self._compute_attention(
|
211
|
+
query, key, value, attention_mask
|
212
|
+
)
|
213
|
+
|
214
|
+
attention_output = self.dropout_layer(
|
215
|
+
attention_output, training=training
|
216
|
+
)
|
217
|
+
|
218
|
+
attention_output = self.output_dense(attention_output)
|
219
|
+
|
220
|
+
if cache is not None:
|
221
|
+
return attention_output, cache
|
222
|
+
return attention_output
|
223
|
+
|
224
|
+
def _masked_softmax(self, attention_scores, attention_mask=None):
|
225
|
+
if attention_mask is not None:
|
226
|
+
return self.softmax(attention_scores, attention_mask[:, None, :, :])
|
227
|
+
return self.softmax(attention_scores)
|
228
|
+
|
229
|
+
def _compute_attention(self, query, key, value, attention_mask=None):
|
230
|
+
attention_scores = ops.einsum("bquh,bkuh->buqk", query, key)
|
231
|
+
attention_scores = attention_scores / self._norm_factor
|
232
|
+
attention_scores = self._masked_softmax(
|
233
|
+
attention_scores, attention_mask
|
234
|
+
)
|
235
|
+
attention_scores = ops.cast(attention_scores, self.compute_dtype)
|
236
|
+
attention_output = ops.einsum(
|
237
|
+
"buqk,bkuh->bquh", attention_scores, value
|
238
|
+
)
|
239
|
+
|
240
|
+
return attention_output
|
241
|
+
|
242
|
+
def get_config(self):
|
243
|
+
config = super().get_config()
|
244
|
+
config.update(
|
245
|
+
{
|
246
|
+
"num_query_heads": self.num_query_heads,
|
247
|
+
"num_key_value_heads": self.num_key_value_heads,
|
248
|
+
"kernel_initializer": keras.initializers.serialize(
|
249
|
+
self.kernel_initializer
|
250
|
+
),
|
251
|
+
"dropout": self.dropout,
|
252
|
+
"max_sequence_length": self.max_sequence_length,
|
253
|
+
"pretraining_sequence_length": self.pretraining_sequence_length,
|
254
|
+
"rope_max_wavelength": self.rope_max_wavelength,
|
255
|
+
"rope_scaling_type": self.rope_scaling_type,
|
256
|
+
"rope_scaling_short_factor": self.rope_scaling_short_factor,
|
257
|
+
"rope_scaling_long_factor": self.rope_scaling_long_factor,
|
258
|
+
}
|
259
|
+
)
|
260
|
+
return config
|
@@ -0,0 +1,224 @@
|
|
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.layers.modeling.reversible_embedding import (
|
18
|
+
ReversibleEmbedding,
|
19
|
+
)
|
20
|
+
from keras_hub.src.models.backbone import Backbone
|
21
|
+
from keras_hub.src.models.phi3.phi3_decoder import Phi3Decoder
|
22
|
+
from keras_hub.src.models.phi3.phi3_layernorm import Phi3LayerNorm
|
23
|
+
|
24
|
+
|
25
|
+
def _phi3_kernel_initializer(stddev=0.02):
|
26
|
+
return keras.initializers.RandomNormal(stddev=stddev)
|
27
|
+
|
28
|
+
|
29
|
+
@keras_hub_export("keras_hub.models.Phi3Backbone")
|
30
|
+
class Phi3Backbone(Backbone):
|
31
|
+
"""Phi-3 core network with hyperparameters.
|
32
|
+
|
33
|
+
This network implements a Transformer-based decoder network,
|
34
|
+
Phi-3, as described in ["Phi-3 Technical Report"](https://arxiv.org/pdf/2404.14219.pdf).
|
35
|
+
It includes the embedding lookups and transformer layers.
|
36
|
+
|
37
|
+
The default constructor gives a fully customizable, randomly initialized
|
38
|
+
phi-3 model with any number of layers, heads, and embedding
|
39
|
+
dimensions. To load preset architectures and weights, use the `from_preset`
|
40
|
+
constructor.
|
41
|
+
|
42
|
+
Args:
|
43
|
+
vocabulary_size (int): The size of the token vocabulary.
|
44
|
+
num_layers (int): The number of transformer layers.
|
45
|
+
hidden_dim (int): The size of the embeddings and the hidden states of
|
46
|
+
the transformer layers.
|
47
|
+
intermediate_dim (int): The output dimension of the first Dense layer in
|
48
|
+
a three-layer feedforward network for each transformer.
|
49
|
+
num_query_heads (int): The number of query attention heads for each
|
50
|
+
transformer layer.
|
51
|
+
num_key_value_heads (int): The number of key and value attention heads
|
52
|
+
for each transformer layer.
|
53
|
+
layer_norm_epsilon (float, optional): Epsilon for the RMS layernorm
|
54
|
+
layers in the transformer decoder. Defaults to `1e-6`.
|
55
|
+
dropout: (float, optional): Dropout probability for the Transformer
|
56
|
+
decoder.
|
57
|
+
max_sequence_length (int, optional): The maximum sequence length
|
58
|
+
that this model might ever be used with. Defaults to `4096`.
|
59
|
+
pretraining_sequence_length (int, optional): The maximum sequence length
|
60
|
+
that the model was pretrained with. Defaults to `4096`.
|
61
|
+
rope_max_wavelength (int, optional): The maximum angular wavelength of
|
62
|
+
the sine/cosine curves, for rotary embeddings. Defaults to `10000`.
|
63
|
+
rope_scaling_type (str, optional): The type of the rope scaling. Can be
|
64
|
+
either `None` or `"su"`. `None` is for no rope scaling, `"su"` is
|
65
|
+
for SuScaled rope, `"su"` is used when `max_sequence_length` is
|
66
|
+
larger than `original_max_sequence_length`. Defaults to `None`.
|
67
|
+
rope_scaling_short_factor List[float]: List of factors used to adjust
|
68
|
+
rope frequencies when the `rope_scaling_type` is `"su"`. List must
|
69
|
+
be of length `hidden_dim//num_query_heads//2`. It is used when
|
70
|
+
`sequence_length` is smaller than `original_max_sequence_length`.
|
71
|
+
Defaults to `None`.
|
72
|
+
rope_scaling_long_factor List[float]: List of factors used to adjust
|
73
|
+
rope frequencies when the `rope_scaling_type` is `"su"`. List must
|
74
|
+
be of length `hidden_dim//num_query_heads//2`. It is used when
|
75
|
+
`sequence_length` is larger than `original_max_sequence_length`.
|
76
|
+
Defaults to `None`.
|
77
|
+
dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
|
78
|
+
for model computations and weights. Note that some computations,
|
79
|
+
such as softmax and layer normalization, will always be done at
|
80
|
+
float32 precision regardless of dtype.
|
81
|
+
|
82
|
+
Examples:
|
83
|
+
|
84
|
+
```python
|
85
|
+
input_data = {
|
86
|
+
"token_ids": np.ones(shape=(1, 12), dtype="int32"),
|
87
|
+
"padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
|
88
|
+
}
|
89
|
+
|
90
|
+
# Pretrained Phi3 decoder.
|
91
|
+
model = keras_hub.models.Phi3Backbone.from_preset(
|
92
|
+
"phi3_mini_4k_instruct_en"
|
93
|
+
)
|
94
|
+
model(input_data)
|
95
|
+
|
96
|
+
# Randomly initialized Phi3 decoder with custom config.
|
97
|
+
model = keras_hub.models.Phi3Backbone(
|
98
|
+
vocabulary_size=10,
|
99
|
+
num_layers=2,
|
100
|
+
hidden_dim=512,
|
101
|
+
intermediate_dim=1024,
|
102
|
+
num_query_heads=32,
|
103
|
+
num_key_value_heads=8,
|
104
|
+
layer_norm_epsilon=1e-6,
|
105
|
+
dtype="float32"
|
106
|
+
)
|
107
|
+
model(input_data)
|
108
|
+
```
|
109
|
+
"""
|
110
|
+
|
111
|
+
def __init__(
|
112
|
+
self,
|
113
|
+
vocabulary_size,
|
114
|
+
num_layers,
|
115
|
+
hidden_dim,
|
116
|
+
intermediate_dim,
|
117
|
+
num_query_heads,
|
118
|
+
num_key_value_heads,
|
119
|
+
layer_norm_epsilon=1e-6,
|
120
|
+
dropout=0.0,
|
121
|
+
max_sequence_length=4096,
|
122
|
+
pretraining_sequence_length=4096,
|
123
|
+
rope_max_wavelength=10000,
|
124
|
+
rope_scaling_type=None,
|
125
|
+
rope_scaling_short_factor=None,
|
126
|
+
rope_scaling_long_factor=None,
|
127
|
+
dtype=None,
|
128
|
+
**kwargs,
|
129
|
+
):
|
130
|
+
# === Layers ===
|
131
|
+
self.token_embedding = ReversibleEmbedding(
|
132
|
+
input_dim=vocabulary_size,
|
133
|
+
output_dim=hidden_dim,
|
134
|
+
tie_weights=False,
|
135
|
+
embeddings_initializer=_phi3_kernel_initializer(stddev=0.01),
|
136
|
+
dtype=dtype,
|
137
|
+
name="token_embedding",
|
138
|
+
)
|
139
|
+
self.transformer_layers = []
|
140
|
+
for i in range(num_layers):
|
141
|
+
layer = Phi3Decoder(
|
142
|
+
hidden_dim=hidden_dim,
|
143
|
+
intermediate_dim=intermediate_dim,
|
144
|
+
num_query_heads=num_query_heads,
|
145
|
+
num_key_value_heads=num_key_value_heads,
|
146
|
+
rope_max_wavelength=rope_max_wavelength,
|
147
|
+
layer_norm_epsilon=layer_norm_epsilon,
|
148
|
+
activation="silu",
|
149
|
+
kernel_initializer=_phi3_kernel_initializer(stddev=0.02),
|
150
|
+
dropout=dropout,
|
151
|
+
max_sequence_length=max_sequence_length,
|
152
|
+
pretraining_sequence_length=pretraining_sequence_length,
|
153
|
+
rope_scaling_type=rope_scaling_type,
|
154
|
+
rope_scaling_short_factor=rope_scaling_short_factor,
|
155
|
+
rope_scaling_long_factor=rope_scaling_long_factor,
|
156
|
+
dtype=dtype,
|
157
|
+
name=f"transformer_layer_{i}",
|
158
|
+
)
|
159
|
+
self.transformer_layers.append(layer)
|
160
|
+
self.layer_norm = Phi3LayerNorm(
|
161
|
+
epsilon=layer_norm_epsilon,
|
162
|
+
dtype=dtype,
|
163
|
+
name="sequence_output_layernorm",
|
164
|
+
)
|
165
|
+
|
166
|
+
# === Functional Model ===
|
167
|
+
token_id_input = keras.Input(
|
168
|
+
shape=(None,), dtype="int32", name="token_ids"
|
169
|
+
)
|
170
|
+
padding_mask_input = keras.Input(
|
171
|
+
shape=(None,), dtype="int32", name="padding_mask"
|
172
|
+
)
|
173
|
+
x = self.token_embedding(token_id_input)
|
174
|
+
for transformer_layer in self.transformer_layers:
|
175
|
+
x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
|
176
|
+
sequence_output = self.layer_norm(x)
|
177
|
+
super().__init__(
|
178
|
+
inputs={
|
179
|
+
"token_ids": token_id_input,
|
180
|
+
"padding_mask": padding_mask_input,
|
181
|
+
},
|
182
|
+
outputs=sequence_output,
|
183
|
+
dtype=dtype,
|
184
|
+
**kwargs,
|
185
|
+
)
|
186
|
+
|
187
|
+
# === Config ===
|
188
|
+
self.vocabulary_size = vocabulary_size
|
189
|
+
self.num_layers = num_layers
|
190
|
+
self.num_query_heads = num_query_heads
|
191
|
+
self.num_key_value_heads = num_key_value_heads
|
192
|
+
self.hidden_dim = hidden_dim
|
193
|
+
self.intermediate_dim = intermediate_dim
|
194
|
+
self.rope_scaling_type = rope_scaling_type
|
195
|
+
self.layer_norm_epsilon = layer_norm_epsilon
|
196
|
+
self.dropout = dropout
|
197
|
+
self.max_sequence_length = max_sequence_length
|
198
|
+
self.pretraining_sequence_length = pretraining_sequence_length
|
199
|
+
self.rope_max_wavelength = rope_max_wavelength
|
200
|
+
self.rope_scaling_type = rope_scaling_type
|
201
|
+
self.rope_scaling_short_factor = rope_scaling_short_factor
|
202
|
+
self.rope_scaling_long_factor = rope_scaling_long_factor
|
203
|
+
|
204
|
+
def get_config(self):
|
205
|
+
config = super().get_config()
|
206
|
+
config.update(
|
207
|
+
{
|
208
|
+
"vocabulary_size": self.vocabulary_size,
|
209
|
+
"num_layers": self.num_layers,
|
210
|
+
"num_query_heads": self.num_query_heads,
|
211
|
+
"hidden_dim": self.hidden_dim,
|
212
|
+
"intermediate_dim": self.intermediate_dim,
|
213
|
+
"num_key_value_heads": self.num_key_value_heads,
|
214
|
+
"layer_norm_epsilon": self.layer_norm_epsilon,
|
215
|
+
"dropout": self.dropout,
|
216
|
+
"max_sequence_length": self.max_sequence_length,
|
217
|
+
"pretraining_sequence_length": self.pretraining_sequence_length,
|
218
|
+
"rope_max_wavelength": self.rope_max_wavelength,
|
219
|
+
"rope_scaling_type": self.rope_scaling_type,
|
220
|
+
"rope_scaling_short_factor": self.rope_scaling_short_factor,
|
221
|
+
"rope_scaling_long_factor": self.rope_scaling_long_factor,
|
222
|
+
}
|
223
|
+
)
|
224
|
+
return config
|