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.
Files changed (297) hide show
  1. keras_hub/__init__.py +52 -0
  2. keras_hub/api/__init__.py +27 -0
  3. keras_hub/api/layers/__init__.py +47 -0
  4. keras_hub/api/metrics/__init__.py +24 -0
  5. keras_hub/api/models/__init__.py +249 -0
  6. keras_hub/api/samplers/__init__.py +29 -0
  7. keras_hub/api/tokenizers/__init__.py +35 -0
  8. keras_hub/src/__init__.py +13 -0
  9. keras_hub/src/api_export.py +53 -0
  10. keras_hub/src/layers/__init__.py +13 -0
  11. keras_hub/src/layers/modeling/__init__.py +13 -0
  12. keras_hub/src/layers/modeling/alibi_bias.py +143 -0
  13. keras_hub/src/layers/modeling/cached_multi_head_attention.py +137 -0
  14. keras_hub/src/layers/modeling/f_net_encoder.py +200 -0
  15. keras_hub/src/layers/modeling/masked_lm_head.py +239 -0
  16. keras_hub/src/layers/modeling/position_embedding.py +123 -0
  17. keras_hub/src/layers/modeling/reversible_embedding.py +311 -0
  18. keras_hub/src/layers/modeling/rotary_embedding.py +169 -0
  19. keras_hub/src/layers/modeling/sine_position_encoding.py +108 -0
  20. keras_hub/src/layers/modeling/token_and_position_embedding.py +150 -0
  21. keras_hub/src/layers/modeling/transformer_decoder.py +496 -0
  22. keras_hub/src/layers/modeling/transformer_encoder.py +262 -0
  23. keras_hub/src/layers/modeling/transformer_layer_utils.py +106 -0
  24. keras_hub/src/layers/preprocessing/__init__.py +13 -0
  25. keras_hub/src/layers/preprocessing/masked_lm_mask_generator.py +220 -0
  26. keras_hub/src/layers/preprocessing/multi_segment_packer.py +319 -0
  27. keras_hub/src/layers/preprocessing/preprocessing_layer.py +62 -0
  28. keras_hub/src/layers/preprocessing/random_deletion.py +271 -0
  29. keras_hub/src/layers/preprocessing/random_swap.py +267 -0
  30. keras_hub/src/layers/preprocessing/start_end_packer.py +219 -0
  31. keras_hub/src/metrics/__init__.py +13 -0
  32. keras_hub/src/metrics/bleu.py +394 -0
  33. keras_hub/src/metrics/edit_distance.py +197 -0
  34. keras_hub/src/metrics/perplexity.py +181 -0
  35. keras_hub/src/metrics/rouge_base.py +204 -0
  36. keras_hub/src/metrics/rouge_l.py +97 -0
  37. keras_hub/src/metrics/rouge_n.py +125 -0
  38. keras_hub/src/models/__init__.py +13 -0
  39. keras_hub/src/models/albert/__init__.py +20 -0
  40. keras_hub/src/models/albert/albert_backbone.py +267 -0
  41. keras_hub/src/models/albert/albert_classifier.py +202 -0
  42. keras_hub/src/models/albert/albert_masked_lm.py +129 -0
  43. keras_hub/src/models/albert/albert_masked_lm_preprocessor.py +194 -0
  44. keras_hub/src/models/albert/albert_preprocessor.py +206 -0
  45. keras_hub/src/models/albert/albert_presets.py +70 -0
  46. keras_hub/src/models/albert/albert_tokenizer.py +119 -0
  47. keras_hub/src/models/backbone.py +311 -0
  48. keras_hub/src/models/bart/__init__.py +20 -0
  49. keras_hub/src/models/bart/bart_backbone.py +261 -0
  50. keras_hub/src/models/bart/bart_preprocessor.py +276 -0
  51. keras_hub/src/models/bart/bart_presets.py +74 -0
  52. keras_hub/src/models/bart/bart_seq_2_seq_lm.py +490 -0
  53. keras_hub/src/models/bart/bart_seq_2_seq_lm_preprocessor.py +262 -0
  54. keras_hub/src/models/bart/bart_tokenizer.py +124 -0
  55. keras_hub/src/models/bert/__init__.py +23 -0
  56. keras_hub/src/models/bert/bert_backbone.py +227 -0
  57. keras_hub/src/models/bert/bert_classifier.py +183 -0
  58. keras_hub/src/models/bert/bert_masked_lm.py +131 -0
  59. keras_hub/src/models/bert/bert_masked_lm_preprocessor.py +198 -0
  60. keras_hub/src/models/bert/bert_preprocessor.py +184 -0
  61. keras_hub/src/models/bert/bert_presets.py +147 -0
  62. keras_hub/src/models/bert/bert_tokenizer.py +112 -0
  63. keras_hub/src/models/bloom/__init__.py +20 -0
  64. keras_hub/src/models/bloom/bloom_attention.py +186 -0
  65. keras_hub/src/models/bloom/bloom_backbone.py +173 -0
  66. keras_hub/src/models/bloom/bloom_causal_lm.py +298 -0
  67. keras_hub/src/models/bloom/bloom_causal_lm_preprocessor.py +176 -0
  68. keras_hub/src/models/bloom/bloom_decoder.py +206 -0
  69. keras_hub/src/models/bloom/bloom_preprocessor.py +185 -0
  70. keras_hub/src/models/bloom/bloom_presets.py +121 -0
  71. keras_hub/src/models/bloom/bloom_tokenizer.py +116 -0
  72. keras_hub/src/models/causal_lm.py +383 -0
  73. keras_hub/src/models/classifier.py +109 -0
  74. keras_hub/src/models/csp_darknet/__init__.py +13 -0
  75. keras_hub/src/models/csp_darknet/csp_darknet_backbone.py +410 -0
  76. keras_hub/src/models/csp_darknet/csp_darknet_image_classifier.py +133 -0
  77. keras_hub/src/models/deberta_v3/__init__.py +24 -0
  78. keras_hub/src/models/deberta_v3/deberta_v3_backbone.py +210 -0
  79. keras_hub/src/models/deberta_v3/deberta_v3_classifier.py +228 -0
  80. keras_hub/src/models/deberta_v3/deberta_v3_masked_lm.py +135 -0
  81. keras_hub/src/models/deberta_v3/deberta_v3_masked_lm_preprocessor.py +191 -0
  82. keras_hub/src/models/deberta_v3/deberta_v3_preprocessor.py +206 -0
  83. keras_hub/src/models/deberta_v3/deberta_v3_presets.py +82 -0
  84. keras_hub/src/models/deberta_v3/deberta_v3_tokenizer.py +155 -0
  85. keras_hub/src/models/deberta_v3/disentangled_attention_encoder.py +227 -0
  86. keras_hub/src/models/deberta_v3/disentangled_self_attention.py +412 -0
  87. keras_hub/src/models/deberta_v3/relative_embedding.py +94 -0
  88. keras_hub/src/models/densenet/__init__.py +13 -0
  89. keras_hub/src/models/densenet/densenet_backbone.py +210 -0
  90. keras_hub/src/models/densenet/densenet_image_classifier.py +131 -0
  91. keras_hub/src/models/distil_bert/__init__.py +26 -0
  92. keras_hub/src/models/distil_bert/distil_bert_backbone.py +187 -0
  93. keras_hub/src/models/distil_bert/distil_bert_classifier.py +208 -0
  94. keras_hub/src/models/distil_bert/distil_bert_masked_lm.py +137 -0
  95. keras_hub/src/models/distil_bert/distil_bert_masked_lm_preprocessor.py +194 -0
  96. keras_hub/src/models/distil_bert/distil_bert_preprocessor.py +175 -0
  97. keras_hub/src/models/distil_bert/distil_bert_presets.py +57 -0
  98. keras_hub/src/models/distil_bert/distil_bert_tokenizer.py +114 -0
  99. keras_hub/src/models/electra/__init__.py +20 -0
  100. keras_hub/src/models/electra/electra_backbone.py +247 -0
  101. keras_hub/src/models/electra/electra_preprocessor.py +154 -0
  102. keras_hub/src/models/electra/electra_presets.py +95 -0
  103. keras_hub/src/models/electra/electra_tokenizer.py +104 -0
  104. keras_hub/src/models/f_net/__init__.py +20 -0
  105. keras_hub/src/models/f_net/f_net_backbone.py +236 -0
  106. keras_hub/src/models/f_net/f_net_classifier.py +154 -0
  107. keras_hub/src/models/f_net/f_net_masked_lm.py +132 -0
  108. keras_hub/src/models/f_net/f_net_masked_lm_preprocessor.py +196 -0
  109. keras_hub/src/models/f_net/f_net_preprocessor.py +177 -0
  110. keras_hub/src/models/f_net/f_net_presets.py +43 -0
  111. keras_hub/src/models/f_net/f_net_tokenizer.py +95 -0
  112. keras_hub/src/models/falcon/__init__.py +20 -0
  113. keras_hub/src/models/falcon/falcon_attention.py +156 -0
  114. keras_hub/src/models/falcon/falcon_backbone.py +164 -0
  115. keras_hub/src/models/falcon/falcon_causal_lm.py +291 -0
  116. keras_hub/src/models/falcon/falcon_causal_lm_preprocessor.py +173 -0
  117. keras_hub/src/models/falcon/falcon_preprocessor.py +187 -0
  118. keras_hub/src/models/falcon/falcon_presets.py +30 -0
  119. keras_hub/src/models/falcon/falcon_tokenizer.py +110 -0
  120. keras_hub/src/models/falcon/falcon_transformer_decoder.py +255 -0
  121. keras_hub/src/models/feature_pyramid_backbone.py +73 -0
  122. keras_hub/src/models/gemma/__init__.py +20 -0
  123. keras_hub/src/models/gemma/gemma_attention.py +250 -0
  124. keras_hub/src/models/gemma/gemma_backbone.py +316 -0
  125. keras_hub/src/models/gemma/gemma_causal_lm.py +448 -0
  126. keras_hub/src/models/gemma/gemma_causal_lm_preprocessor.py +167 -0
  127. keras_hub/src/models/gemma/gemma_decoder_block.py +241 -0
  128. keras_hub/src/models/gemma/gemma_preprocessor.py +191 -0
  129. keras_hub/src/models/gemma/gemma_presets.py +248 -0
  130. keras_hub/src/models/gemma/gemma_tokenizer.py +103 -0
  131. keras_hub/src/models/gemma/rms_normalization.py +40 -0
  132. keras_hub/src/models/gpt2/__init__.py +20 -0
  133. keras_hub/src/models/gpt2/gpt2_backbone.py +199 -0
  134. keras_hub/src/models/gpt2/gpt2_causal_lm.py +437 -0
  135. keras_hub/src/models/gpt2/gpt2_causal_lm_preprocessor.py +173 -0
  136. keras_hub/src/models/gpt2/gpt2_preprocessor.py +187 -0
  137. keras_hub/src/models/gpt2/gpt2_presets.py +82 -0
  138. keras_hub/src/models/gpt2/gpt2_tokenizer.py +110 -0
  139. keras_hub/src/models/gpt_neo_x/__init__.py +13 -0
  140. keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py +251 -0
  141. keras_hub/src/models/gpt_neo_x/gpt_neo_x_backbone.py +175 -0
  142. keras_hub/src/models/gpt_neo_x/gpt_neo_x_causal_lm.py +201 -0
  143. keras_hub/src/models/gpt_neo_x/gpt_neo_x_causal_lm_preprocessor.py +141 -0
  144. keras_hub/src/models/gpt_neo_x/gpt_neo_x_decoder.py +258 -0
  145. keras_hub/src/models/gpt_neo_x/gpt_neo_x_preprocessor.py +145 -0
  146. keras_hub/src/models/gpt_neo_x/gpt_neo_x_tokenizer.py +88 -0
  147. keras_hub/src/models/image_classifier.py +90 -0
  148. keras_hub/src/models/llama/__init__.py +20 -0
  149. keras_hub/src/models/llama/llama_attention.py +225 -0
  150. keras_hub/src/models/llama/llama_backbone.py +188 -0
  151. keras_hub/src/models/llama/llama_causal_lm.py +327 -0
  152. keras_hub/src/models/llama/llama_causal_lm_preprocessor.py +170 -0
  153. keras_hub/src/models/llama/llama_decoder.py +246 -0
  154. keras_hub/src/models/llama/llama_layernorm.py +48 -0
  155. keras_hub/src/models/llama/llama_preprocessor.py +189 -0
  156. keras_hub/src/models/llama/llama_presets.py +80 -0
  157. keras_hub/src/models/llama/llama_tokenizer.py +84 -0
  158. keras_hub/src/models/llama3/__init__.py +20 -0
  159. keras_hub/src/models/llama3/llama3_backbone.py +84 -0
  160. keras_hub/src/models/llama3/llama3_causal_lm.py +46 -0
  161. keras_hub/src/models/llama3/llama3_causal_lm_preprocessor.py +173 -0
  162. keras_hub/src/models/llama3/llama3_preprocessor.py +21 -0
  163. keras_hub/src/models/llama3/llama3_presets.py +69 -0
  164. keras_hub/src/models/llama3/llama3_tokenizer.py +63 -0
  165. keras_hub/src/models/masked_lm.py +101 -0
  166. keras_hub/src/models/mistral/__init__.py +20 -0
  167. keras_hub/src/models/mistral/mistral_attention.py +238 -0
  168. keras_hub/src/models/mistral/mistral_backbone.py +203 -0
  169. keras_hub/src/models/mistral/mistral_causal_lm.py +328 -0
  170. keras_hub/src/models/mistral/mistral_causal_lm_preprocessor.py +175 -0
  171. keras_hub/src/models/mistral/mistral_layer_norm.py +48 -0
  172. keras_hub/src/models/mistral/mistral_preprocessor.py +190 -0
  173. keras_hub/src/models/mistral/mistral_presets.py +48 -0
  174. keras_hub/src/models/mistral/mistral_tokenizer.py +82 -0
  175. keras_hub/src/models/mistral/mistral_transformer_decoder.py +265 -0
  176. keras_hub/src/models/mix_transformer/__init__.py +13 -0
  177. keras_hub/src/models/mix_transformer/mix_transformer_backbone.py +181 -0
  178. keras_hub/src/models/mix_transformer/mix_transformer_classifier.py +133 -0
  179. keras_hub/src/models/mix_transformer/mix_transformer_layers.py +300 -0
  180. keras_hub/src/models/opt/__init__.py +20 -0
  181. keras_hub/src/models/opt/opt_backbone.py +173 -0
  182. keras_hub/src/models/opt/opt_causal_lm.py +301 -0
  183. keras_hub/src/models/opt/opt_causal_lm_preprocessor.py +177 -0
  184. keras_hub/src/models/opt/opt_preprocessor.py +188 -0
  185. keras_hub/src/models/opt/opt_presets.py +72 -0
  186. keras_hub/src/models/opt/opt_tokenizer.py +116 -0
  187. keras_hub/src/models/pali_gemma/__init__.py +23 -0
  188. keras_hub/src/models/pali_gemma/pali_gemma_backbone.py +277 -0
  189. keras_hub/src/models/pali_gemma/pali_gemma_causal_lm.py +313 -0
  190. keras_hub/src/models/pali_gemma/pali_gemma_causal_lm_preprocessor.py +147 -0
  191. keras_hub/src/models/pali_gemma/pali_gemma_decoder_block.py +160 -0
  192. keras_hub/src/models/pali_gemma/pali_gemma_presets.py +78 -0
  193. keras_hub/src/models/pali_gemma/pali_gemma_tokenizer.py +79 -0
  194. keras_hub/src/models/pali_gemma/pali_gemma_vit.py +566 -0
  195. keras_hub/src/models/phi3/__init__.py +20 -0
  196. keras_hub/src/models/phi3/phi3_attention.py +260 -0
  197. keras_hub/src/models/phi3/phi3_backbone.py +224 -0
  198. keras_hub/src/models/phi3/phi3_causal_lm.py +218 -0
  199. keras_hub/src/models/phi3/phi3_causal_lm_preprocessor.py +173 -0
  200. keras_hub/src/models/phi3/phi3_decoder.py +260 -0
  201. keras_hub/src/models/phi3/phi3_layernorm.py +48 -0
  202. keras_hub/src/models/phi3/phi3_preprocessor.py +190 -0
  203. keras_hub/src/models/phi3/phi3_presets.py +50 -0
  204. keras_hub/src/models/phi3/phi3_rotary_embedding.py +137 -0
  205. keras_hub/src/models/phi3/phi3_tokenizer.py +94 -0
  206. keras_hub/src/models/preprocessor.py +207 -0
  207. keras_hub/src/models/resnet/__init__.py +13 -0
  208. keras_hub/src/models/resnet/resnet_backbone.py +612 -0
  209. keras_hub/src/models/resnet/resnet_image_classifier.py +136 -0
  210. keras_hub/src/models/roberta/__init__.py +20 -0
  211. keras_hub/src/models/roberta/roberta_backbone.py +184 -0
  212. keras_hub/src/models/roberta/roberta_classifier.py +209 -0
  213. keras_hub/src/models/roberta/roberta_masked_lm.py +136 -0
  214. keras_hub/src/models/roberta/roberta_masked_lm_preprocessor.py +198 -0
  215. keras_hub/src/models/roberta/roberta_preprocessor.py +192 -0
  216. keras_hub/src/models/roberta/roberta_presets.py +43 -0
  217. keras_hub/src/models/roberta/roberta_tokenizer.py +132 -0
  218. keras_hub/src/models/seq_2_seq_lm.py +54 -0
  219. keras_hub/src/models/t5/__init__.py +20 -0
  220. keras_hub/src/models/t5/t5_backbone.py +261 -0
  221. keras_hub/src/models/t5/t5_layer_norm.py +35 -0
  222. keras_hub/src/models/t5/t5_multi_head_attention.py +324 -0
  223. keras_hub/src/models/t5/t5_presets.py +95 -0
  224. keras_hub/src/models/t5/t5_tokenizer.py +100 -0
  225. keras_hub/src/models/t5/t5_transformer_layer.py +178 -0
  226. keras_hub/src/models/task.py +419 -0
  227. keras_hub/src/models/vgg/__init__.py +13 -0
  228. keras_hub/src/models/vgg/vgg_backbone.py +158 -0
  229. keras_hub/src/models/vgg/vgg_image_classifier.py +124 -0
  230. keras_hub/src/models/vit_det/__init__.py +13 -0
  231. keras_hub/src/models/vit_det/vit_det_backbone.py +204 -0
  232. keras_hub/src/models/vit_det/vit_layers.py +565 -0
  233. keras_hub/src/models/whisper/__init__.py +20 -0
  234. keras_hub/src/models/whisper/whisper_audio_feature_extractor.py +260 -0
  235. keras_hub/src/models/whisper/whisper_backbone.py +305 -0
  236. keras_hub/src/models/whisper/whisper_cached_multi_head_attention.py +153 -0
  237. keras_hub/src/models/whisper/whisper_decoder.py +141 -0
  238. keras_hub/src/models/whisper/whisper_encoder.py +106 -0
  239. keras_hub/src/models/whisper/whisper_preprocessor.py +326 -0
  240. keras_hub/src/models/whisper/whisper_presets.py +148 -0
  241. keras_hub/src/models/whisper/whisper_tokenizer.py +163 -0
  242. keras_hub/src/models/xlm_roberta/__init__.py +26 -0
  243. keras_hub/src/models/xlm_roberta/xlm_roberta_backbone.py +81 -0
  244. keras_hub/src/models/xlm_roberta/xlm_roberta_classifier.py +225 -0
  245. keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm.py +141 -0
  246. keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm_preprocessor.py +195 -0
  247. keras_hub/src/models/xlm_roberta/xlm_roberta_preprocessor.py +205 -0
  248. keras_hub/src/models/xlm_roberta/xlm_roberta_presets.py +43 -0
  249. keras_hub/src/models/xlm_roberta/xlm_roberta_tokenizer.py +191 -0
  250. keras_hub/src/models/xlnet/__init__.py +13 -0
  251. keras_hub/src/models/xlnet/relative_attention.py +459 -0
  252. keras_hub/src/models/xlnet/xlnet_backbone.py +222 -0
  253. keras_hub/src/models/xlnet/xlnet_content_and_query_embedding.py +133 -0
  254. keras_hub/src/models/xlnet/xlnet_encoder.py +378 -0
  255. keras_hub/src/samplers/__init__.py +13 -0
  256. keras_hub/src/samplers/beam_sampler.py +207 -0
  257. keras_hub/src/samplers/contrastive_sampler.py +231 -0
  258. keras_hub/src/samplers/greedy_sampler.py +50 -0
  259. keras_hub/src/samplers/random_sampler.py +77 -0
  260. keras_hub/src/samplers/sampler.py +237 -0
  261. keras_hub/src/samplers/serialization.py +97 -0
  262. keras_hub/src/samplers/top_k_sampler.py +92 -0
  263. keras_hub/src/samplers/top_p_sampler.py +113 -0
  264. keras_hub/src/tests/__init__.py +13 -0
  265. keras_hub/src/tests/test_case.py +608 -0
  266. keras_hub/src/tokenizers/__init__.py +13 -0
  267. keras_hub/src/tokenizers/byte_pair_tokenizer.py +638 -0
  268. keras_hub/src/tokenizers/byte_tokenizer.py +299 -0
  269. keras_hub/src/tokenizers/sentence_piece_tokenizer.py +267 -0
  270. keras_hub/src/tokenizers/sentence_piece_tokenizer_trainer.py +150 -0
  271. keras_hub/src/tokenizers/tokenizer.py +235 -0
  272. keras_hub/src/tokenizers/unicode_codepoint_tokenizer.py +355 -0
  273. keras_hub/src/tokenizers/word_piece_tokenizer.py +544 -0
  274. keras_hub/src/tokenizers/word_piece_tokenizer_trainer.py +176 -0
  275. keras_hub/src/utils/__init__.py +13 -0
  276. keras_hub/src/utils/keras_utils.py +130 -0
  277. keras_hub/src/utils/pipeline_model.py +293 -0
  278. keras_hub/src/utils/preset_utils.py +621 -0
  279. keras_hub/src/utils/python_utils.py +21 -0
  280. keras_hub/src/utils/tensor_utils.py +206 -0
  281. keras_hub/src/utils/timm/__init__.py +13 -0
  282. keras_hub/src/utils/timm/convert.py +37 -0
  283. keras_hub/src/utils/timm/convert_resnet.py +171 -0
  284. keras_hub/src/utils/transformers/__init__.py +13 -0
  285. keras_hub/src/utils/transformers/convert.py +101 -0
  286. keras_hub/src/utils/transformers/convert_bert.py +173 -0
  287. keras_hub/src/utils/transformers/convert_distilbert.py +184 -0
  288. keras_hub/src/utils/transformers/convert_gemma.py +187 -0
  289. keras_hub/src/utils/transformers/convert_gpt2.py +186 -0
  290. keras_hub/src/utils/transformers/convert_llama3.py +136 -0
  291. keras_hub/src/utils/transformers/convert_pali_gemma.py +303 -0
  292. keras_hub/src/utils/transformers/safetensor_utils.py +97 -0
  293. keras_hub/src/version_utils.py +23 -0
  294. keras_hub_nightly-0.15.0.dev20240823171555.dist-info/METADATA +34 -0
  295. keras_hub_nightly-0.15.0.dev20240823171555.dist-info/RECORD +297 -0
  296. keras_hub_nightly-0.15.0.dev20240823171555.dist-info/WHEEL +5 -0
  297. keras_hub_nightly-0.15.0.dev20240823171555.dist-info/top_level.txt +1 -0
@@ -0,0 +1,188 @@
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.preprocessing.start_end_packer import StartEndPacker
20
+ from keras_hub.src.models.opt.opt_tokenizer import OPTTokenizer
21
+ from keras_hub.src.models.preprocessor import Preprocessor
22
+ from keras_hub.src.utils.keras_utils import (
23
+ convert_inputs_to_list_of_tensor_segments,
24
+ )
25
+
26
+
27
+ @keras_hub_export("keras_hub.models.OPTPreprocessor")
28
+ class OPTPreprocessor(Preprocessor):
29
+ """OPT preprocessing layer which tokenizes and packs inputs.
30
+
31
+ This preprocessing layer will do 2 things:
32
+
33
+ - Tokenize the input using the `tokenizer`.
34
+ - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
35
+ be passed directly to a `keras_hub.models.OPTBackbone`.
36
+
37
+ This layer can be used directly with `tf.data.Dataset.map` to preprocess
38
+ string data in the `(x, y, sample_weight)` format used by
39
+ `keras.Model.fit`.
40
+
41
+ The call method of this layer accepts three arguments, `x`, `y`, and
42
+ `sample_weight`. `x` can be a python string or tensor representing a single
43
+ segment, a list of python strings representing a batch of single segments,
44
+ or a list of tensors representing multiple segments to be packed together.
45
+ `y` and `sample_weight` are both optional, can have any format, and will be
46
+ passed through unaltered.
47
+
48
+ `OPTPreprocessor` forces the input to have only one segment, as OPT is
49
+ mainly used for generation tasks. For tasks having multi-segment inputs
50
+ like "glue/mnli", please use a model designed for classification purposes
51
+ such as BERT or RoBERTa.
52
+
53
+ Args:
54
+ tokenizer: A `keras_hub.models.OPTTokenizer` instance.
55
+ sequence_length: The length of the packed inputs.
56
+ add_start_token: If `True`, the preprocessor will append the tokenizer
57
+ start token to each input sequence.
58
+ add_end_token: If `True`, the preprocessor will append the tokenizer
59
+ end token to each input sequence.
60
+
61
+ Call arguments:
62
+ x: A string, `tf.Tensor` or list of python strings.
63
+ y: Any label data. Will be passed through unaltered.
64
+ sample_weight: Any label weight data. Will be passed through unaltered.
65
+ sequence_length: Pass to override the configured `sequence_length` of
66
+ the layer.
67
+
68
+ Examples:
69
+
70
+ Directly calling the layer on data.
71
+ ```python
72
+ preprocessor = keras_hub.models.OPTPreprocessor.from_preset("opt_125m_en")
73
+
74
+ # Tokenize and pack a single sentence.
75
+ preprocessor("The quick brown fox jumped.")
76
+
77
+ # Tokenize a batch of single sentences.
78
+ preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
79
+
80
+ # Custom vocabulary.
81
+ features = ["a quick fox.", "a fox quick."]
82
+ vocab = {"<|endoftext|>": 0, "a": 4, "Ġquick": 5, "Ġfox": 6}
83
+ merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
84
+ merges += ["Ġ f", "o x", "Ġf ox"]
85
+ tokenizer = keras_hub.models.OPTTokenizer(
86
+ vocabulary=vocab,
87
+ merges=merges,
88
+ )
89
+ preprocessor = keras_hub.models.OPTPreprocessor(tokenizer=tokenizer)
90
+ preprocessor("The quick brown fox jumped.")
91
+ ```
92
+
93
+ Mapping with `tf.data.Dataset`.
94
+ ```python
95
+ preprocessor = keras_hub.models.OPTPreprocessor.from_preset("opt_125m_en")
96
+
97
+ text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
98
+ label = tf.constant([1, 1])
99
+
100
+ # Map labeled single sentences.
101
+ ds = tf.data.Dataset.from_tensor_slices((text, label))
102
+ ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
103
+
104
+ # Map unlabeled single sentences.
105
+ ds = tf.data.Dataset.from_tensor_slices(text)
106
+ ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
107
+ ```
108
+ """
109
+
110
+ tokenizer_cls = OPTTokenizer
111
+
112
+ def __init__(
113
+ self,
114
+ tokenizer,
115
+ sequence_length=1024,
116
+ add_start_token=True,
117
+ add_end_token=True,
118
+ **kwargs,
119
+ ):
120
+ super().__init__(**kwargs)
121
+
122
+ self.tokenizer = tokenizer
123
+ self.packer = None
124
+ self.sequence_length = sequence_length
125
+ self.add_start_token = add_start_token
126
+ self.add_end_token = add_end_token
127
+
128
+ def build(self, input_shape):
129
+ # Defer packer creation to `build()` so that we can be sure tokenizer
130
+ # assets have loaded when restoring a saved model.
131
+ self.packer = StartEndPacker(
132
+ start_value=self.tokenizer.start_token_id,
133
+ end_value=self.tokenizer.end_token_id,
134
+ pad_value=self.tokenizer.pad_token_id,
135
+ sequence_length=self.sequence_length,
136
+ return_padding_mask=True,
137
+ )
138
+ self.built = True
139
+
140
+ def get_config(self):
141
+ config = super().get_config()
142
+ config.update(
143
+ {
144
+ "sequence_length": self.sequence_length,
145
+ "add_start_token": self.add_start_token,
146
+ "add_end_token": self.add_end_token,
147
+ }
148
+ )
149
+ return config
150
+
151
+ def call(
152
+ self,
153
+ x,
154
+ y=None,
155
+ sample_weight=None,
156
+ sequence_length=None,
157
+ ):
158
+ x = convert_inputs_to_list_of_tensor_segments(x)
159
+ if len(x) != 1:
160
+ raise ValueError(
161
+ "OPT requires each input feature to contain only "
162
+ f"one segment, but received {len(x)}. If you are using OPT "
163
+ "for a multi-segment classification task, please refer to "
164
+ "classification models like BERT or RoBERTa."
165
+ )
166
+ sequence_length = sequence_length or self.sequence_length
167
+ token_ids, padding_mask = self.packer(
168
+ self.tokenizer(x[0]),
169
+ sequence_length=sequence_length,
170
+ add_start_value=self.add_start_token,
171
+ add_end_value=self.add_end_token,
172
+ )
173
+ x = {
174
+ "token_ids": token_ids,
175
+ "padding_mask": padding_mask,
176
+ }
177
+ return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
178
+
179
+ @property
180
+ def sequence_length(self):
181
+ """The padded length of model input sequences."""
182
+ return self._sequence_length
183
+
184
+ @sequence_length.setter
185
+ def sequence_length(self, value):
186
+ self._sequence_length = value
187
+ if self.packer is not None:
188
+ self.packer.sequence_length = value
@@ -0,0 +1,72 @@
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
+ """OPT model preset configurations."""
15
+
16
+ # Metadata for loading pretrained model weights.
17
+ backbone_presets = {
18
+ "opt_125m_en": {
19
+ "metadata": {
20
+ "description": (
21
+ "12-layer OPT model where case in maintained. Trained on "
22
+ "BookCorpus, CommonCrawl, Pile, and PushShift.io corpora."
23
+ ),
24
+ "params": 125237760,
25
+ "official_name": "OPT",
26
+ "path": "opt",
27
+ "model_card": "https://github.com/facebookresearch/metaseq/blob/main/projects/OPT/model_card.md",
28
+ },
29
+ "kaggle_handle": "kaggle://keras/opt/keras/opt_125m_en/2",
30
+ },
31
+ # We skip the 350m checkpoint because it does not match the structure of
32
+ # other checkpoints.
33
+ "opt_1.3b_en": {
34
+ "metadata": {
35
+ "description": (
36
+ "24-layer OPT model where case in maintained. Trained on "
37
+ "BookCorpus, CommonCrawl, Pile, and PushShift.io corpora."
38
+ ),
39
+ "params": 1315753984,
40
+ "official_name": "OPT",
41
+ "path": "opt",
42
+ "model_card": "https://github.com/facebookresearch/metaseq/blob/main/projects/OPT/model_card.md",
43
+ },
44
+ "kaggle_handle": "kaggle://keras/opt/keras/opt_1.3b_en/2",
45
+ },
46
+ "opt_2.7b_en": {
47
+ "metadata": {
48
+ "description": (
49
+ "32-layer OPT model where case in maintained. Trained on "
50
+ "BookCorpus, CommonCrawl, Pile, and PushShift.io corpora."
51
+ ),
52
+ "params": 2700000000,
53
+ "official_name": "OPT",
54
+ "path": "opt",
55
+ "model_card": "https://github.com/facebookresearch/metaseq/blob/main/projects/OPT/model_card.md",
56
+ },
57
+ "kaggle_handle": "kaggle://keras/opt/keras/opt_2.7b_en/2",
58
+ },
59
+ "opt_6.7b_en": {
60
+ "metadata": {
61
+ "description": (
62
+ "32-layer OPT model where case in maintained. Trained on "
63
+ "BookCorpus, CommonCrawl, Pile, and PushShift.io corpora."
64
+ ),
65
+ "params": 6700000000,
66
+ "official_name": "OPT",
67
+ "path": "opt",
68
+ "model_card": "https://github.com/facebookresearch/metaseq/blob/main/projects/OPT/model_card.md",
69
+ },
70
+ "kaggle_handle": "kaggle://keras/opt/keras/opt_6.7b_en/2",
71
+ },
72
+ }
@@ -0,0 +1,116 @@
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
+ from keras_hub.src.api_export import keras_hub_export
17
+ from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
18
+
19
+
20
+ @keras_hub_export("keras_hub.models.OPTTokenizer")
21
+ class OPTTokenizer(BytePairTokenizer):
22
+ """An OPT tokenizer using Byte-Pair Encoding subword segmentation.
23
+
24
+ This tokenizer class will tokenize raw strings into integer sequences and
25
+ is based on `keras_hub.tokenizers.BytePairTokenizer`. Unlike the
26
+ underlying tokenizer, it will check for all special tokens needed by OPT
27
+ models and provides a `from_preset()` method to automatically download
28
+ a matching vocabulary for a OPT preset.
29
+
30
+ This tokenizer does not provide truncation or padding of inputs.
31
+
32
+ If input is a batch of strings (rank > 0), the layer will output a
33
+ `tf.RaggedTensor` where the last dimension of the output is ragged.
34
+ If input is a scalar string (rank == 0), the layer will output a dense
35
+ `tf.Tensor` with static shape `[None]`.
36
+
37
+ Args:
38
+ vocabulary: string or dict, maps token to integer ids. If it is a
39
+ string, it should be the file path to a json file.
40
+ merges: string or list, contains the merge rule. If it is a string,
41
+ it should be the file path to merge rules. The merge rule file
42
+ should have one merge rule per line. Every merge rule contains
43
+ merge entities separated by a space.
44
+
45
+ Examples:
46
+ ```python
47
+ # Unbatched input.
48
+ tokenizer = keras_hub.models.OPTTokenizer.from_preset(
49
+ "opt_125m_en",
50
+ )
51
+ tokenizer("The quick brown fox jumped.")
52
+
53
+ # Batched input.
54
+ tokenizer(["The quick brown fox jumped.", "The fox slept."])
55
+
56
+ # Detokenization.
57
+ tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
58
+
59
+ # Custom vocabulary.
60
+ vocab = {"<pad>": 1, "</s>": 2, "Ġquick": 4, "Ġfox": 5}
61
+ merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
62
+ merges += ["Ġ f", "o x", "Ġf ox"]
63
+ tokenizer = keras_hub.models.OPTTokenizer(vocabulary=vocab, merges=merges)
64
+ tokenizer("The quick brown fox jumped.")
65
+ ```
66
+ """
67
+
68
+ def __init__(
69
+ self,
70
+ vocabulary=None,
71
+ merges=None,
72
+ **kwargs,
73
+ ):
74
+ self.start_token = "</s>"
75
+ self.pad_token = "<pad>"
76
+ self.end_token = "</s>"
77
+
78
+ super().__init__(
79
+ vocabulary=vocabulary,
80
+ merges=merges,
81
+ unsplittable_tokens=[
82
+ self.start_token,
83
+ self.pad_token,
84
+ self.end_token,
85
+ ],
86
+ **kwargs,
87
+ )
88
+
89
+ def set_vocabulary_and_merges(self, vocabulary, merges):
90
+ super().set_vocabulary_and_merges(vocabulary, merges)
91
+
92
+ if vocabulary is not None:
93
+ # Check for necessary special tokens.
94
+ for token in [self.start_token, self.pad_token, self.end_token]:
95
+ if token not in self.vocabulary:
96
+ raise ValueError(
97
+ f"Cannot find token `'{token}'` in the provided "
98
+ f"`vocabulary`. Please provide `'{token}'` in your "
99
+ "`vocabulary` or use a pretrained `vocabulary` name."
100
+ )
101
+
102
+ self.start_token_id = self.token_to_id(self.start_token)
103
+ self.pad_token_id = self.token_to_id(self.pad_token)
104
+ self.end_token_id = self.token_to_id(self.end_token)
105
+ else:
106
+ self.start_token_id = None
107
+ self.pad_token_id = None
108
+ self.end_token_id = None
109
+
110
+ def get_config(self):
111
+ config = super().get_config()
112
+ # In the constructor, we pass the list of special tokens to the
113
+ # `unsplittable_tokens` arg of the superclass' constructor. Hence, we
114
+ # delete it from the config here.
115
+ del config["unsplittable_tokens"]
116
+ return config
@@ -0,0 +1,23 @@
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
+ from keras_hub.src.models.pali_gemma.pali_gemma_backbone import (
15
+ PaliGemmaBackbone,
16
+ )
17
+ from keras_hub.src.models.pali_gemma.pali_gemma_presets import backbone_presets
18
+ from keras_hub.src.models.pali_gemma.pali_gemma_tokenizer import (
19
+ PaliGemmaTokenizer,
20
+ )
21
+ from keras_hub.src.utils.preset_utils import register_presets
22
+
23
+ register_presets(backbone_presets, (PaliGemmaBackbone, PaliGemmaTokenizer))
@@ -0,0 +1,277 @@
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.api_export import keras_hub_export
18
+ from keras_hub.src.layers.modeling.reversible_embedding import (
19
+ ReversibleEmbedding,
20
+ )
21
+ from keras_hub.src.models.backbone import Backbone
22
+ from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
23
+ from keras_hub.src.models.pali_gemma.pali_gemma_decoder_block import (
24
+ PaliGemmaDecoderBlock,
25
+ )
26
+ from keras_hub.src.models.pali_gemma.pali_gemma_vit import PaliGemmaVit
27
+
28
+
29
+ @keras_hub_export("keras_hub.models.PaliGemmaBackbone")
30
+ class PaliGemmaBackbone(Backbone):
31
+ """PaliGemma core network with hyperparameters.
32
+
33
+ This backbone implements the mixed-modality PaliGemma architecture. It
34
+ contains a Visual Transformer network, as well as text token embedding
35
+ layer, followed by a backend-agnostic concatenation operation to
36
+ construct a sequence of representations of mixed type embeddings (visual
37
+ and textual). Then, the concatenated sequence is passed through a series
38
+ of Mixed Modality Decoder Blocks. The returned value from calling this model
39
+ represents probabilistic values for output tokens.
40
+
41
+ For a higher-level object for text-generation,
42
+ see `keras_hub.models.PaliGemmaCausalLM`.
43
+
44
+ The default constructor gives a fully customizable, randomly initialized
45
+ PaliGemma model with any number of vit layers, heads, embedding
46
+ dimensions, and equivalent configuration for Paligemma Decoder layers. To
47
+ load preset architectures and weights, use the `from_preset` constructor.
48
+
49
+ Args:
50
+ vocabulary_size: int. The size of the token vocabulary.
51
+ image_size: int. The resolution of the image in both width and height.
52
+ Note: input images must be square.
53
+ num_layers: int. The number of transformer mixed decoder layers.
54
+ num_query_heads: int. The number of heads for the query projections in
55
+ the mixed decoder attention layer.
56
+ num_key_value_heads: int. The number of heads for the key and value
57
+ projections in the mixed decoder attention layers.
58
+ hidden_dim: int. The size of the transformer hidden state at the end
59
+ of each mixed transformer layer.
60
+ intermediate_dim: int. The output dimension of the first Dense layer in
61
+ a two-layer feedforward network for each transformer decoder block.
62
+ head_dim: int. The size of each attention head in the mixed decoder.
63
+ vit_patch_size: int. The size of each square patch in the input image.
64
+ vit_num_heads: int. The number of attention heads for the vision(image)
65
+ transformer encoder.
66
+ vit_hidden_dim: int. The size of the transformer hidden state at the end
67
+ of each vision transformer layer.
68
+ vit_num_layers: int. The number of vision transformer layers.
69
+ vit_intermediate_dim: int. The output dimension of the first Dense layer
70
+ in a two-layer feedforward network for vision transformer.
71
+ vit_pooling: string. The encoded vision embeddings are pooled using the
72
+ specified polling setting. The accepted values are `"map"`, `"gap"`,
73
+ `"0"` or `"none"`. Defaults to `"none"`.
74
+ vit_classifier_activation: activation function. The activation that
75
+ is used for final output classification in the vision transformer.
76
+ vit_name: string. The name used for vision transformer layers.
77
+ include_rescaling: bool. If true, the image input will be rescaled from
78
+ the range `[0, 255]`, to the range `[0, 1]`.
79
+ layer_norm_epsilon: float. The epsilon value user for every layer norm
80
+ in all transformer blocks.
81
+ dropout: float. Dropout probability for the Transformer decoder blocks.
82
+ dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
83
+ for the models computations and weights. Note that some
84
+ computations, such as softmax and layer normalization will always
85
+ be done a float32 precision regardless of dtype.
86
+
87
+ Example:
88
+ ```python
89
+ input_data = {
90
+ "token_ids": np.ones(shape=(1, 12), dtype="int32"),
91
+ "images": np.random.uniform(size=(1, 224, 224, 3)),
92
+ "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
93
+ }
94
+
95
+ # Pretrained PaliGemma decoder.
96
+ model = keras_hub.models.PaliGemmaBackbone.from_preset("pali_gemma_mix_224")
97
+ model(input_data)
98
+
99
+ # Randomly initialized PaliGemma decoder with custom config.
100
+ model = keras_hub.models.PaliGemmaBackbone(
101
+ vocabulary_size=50257,
102
+ images_size=224,
103
+ num_layers=12,
104
+ num_query_heads=12,
105
+ num_key_value_heads=1,
106
+ hidden_dim=768,
107
+ intermediate_dim=3072,
108
+ head_dim=64,
109
+ vit_patch_size=14,
110
+ vit_num_heads=8,
111
+ vit_hidden_dim=768,
112
+ vit_intermediate_dim=3072,
113
+ vit_num_layers=2,
114
+ )
115
+ model(input_data)
116
+ ```
117
+ """
118
+
119
+ def __init__(
120
+ self,
121
+ vocabulary_size,
122
+ image_size,
123
+ num_layers,
124
+ num_query_heads,
125
+ num_key_value_heads,
126
+ hidden_dim,
127
+ intermediate_dim,
128
+ head_dim,
129
+ vit_patch_size,
130
+ vit_num_heads,
131
+ vit_hidden_dim,
132
+ vit_num_layers,
133
+ vit_intermediate_dim=None, # TODO remove default
134
+ vit_pooling=None,
135
+ vit_classifier_activation=None,
136
+ vit_name=None,
137
+ include_rescaling=True,
138
+ layer_norm_epsilon=1e-6,
139
+ dropout=0,
140
+ dtype=None,
141
+ **kwargs,
142
+ ):
143
+ # === Layers ===
144
+ self.token_embedding = ReversibleEmbedding(
145
+ input_dim=vocabulary_size,
146
+ output_dim=hidden_dim,
147
+ tie_weights=True,
148
+ embeddings_initializer=keras.initializers.VarianceScaling(
149
+ scale=1.0,
150
+ mode="fan_in",
151
+ distribution="untruncated_normal",
152
+ seed=None,
153
+ ),
154
+ dtype=dtype,
155
+ name="token_embedding",
156
+ )
157
+ # TODO Remove this. Work around for previous serialization bug.
158
+ vit_intermediate_dim = vit_intermediate_dim or 4304
159
+ self.vit_encoder = PaliGemmaVit(
160
+ image_size=image_size,
161
+ include_rescaling=include_rescaling,
162
+ patch_size=vit_patch_size,
163
+ num_heads=vit_num_heads,
164
+ hidden_dim=vit_hidden_dim,
165
+ num_layers=vit_num_layers,
166
+ intermediate_dim=vit_intermediate_dim,
167
+ pooling=vit_pooling,
168
+ num_classes=hidden_dim,
169
+ classifier_activation=vit_classifier_activation,
170
+ dtype=dtype,
171
+ name=vit_name,
172
+ )
173
+ self.transformer_layers = []
174
+ for i in range(num_layers):
175
+ layer = PaliGemmaDecoderBlock(
176
+ hidden_dim=hidden_dim,
177
+ intermediate_dim=intermediate_dim,
178
+ num_query_heads=num_query_heads,
179
+ head_dim=head_dim,
180
+ num_key_value_heads=num_key_value_heads,
181
+ dropout=dropout,
182
+ dtype=dtype,
183
+ name=f"decoder_block_{i}",
184
+ )
185
+ self.transformer_layers.append(layer)
186
+ self.layer_norm = RMSNormalization(
187
+ epsilon=layer_norm_epsilon,
188
+ dtype=dtype,
189
+ name="final_normalization",
190
+ )
191
+
192
+ # === Functional Model ===
193
+ image_input = self.vit_encoder.inputs[0]
194
+ token_id_input = keras.Input(
195
+ shape=(None,), dtype="int32", name="token_ids"
196
+ )
197
+ padding_mask_input = keras.Input(
198
+ shape=(None,), dtype="int32", name="padding_mask"
199
+ )
200
+ response_mask_input = keras.Input(
201
+ shape=(None,), dtype="int32", name="response_mask"
202
+ )
203
+ img_embeddings = self.vit_encoder(image_input)
204
+ text_embeddings = self.token_embedding(token_id_input)
205
+ text_embeddings = text_embeddings * ops.cast(
206
+ ops.sqrt(hidden_dim), text_embeddings.dtype
207
+ )
208
+ x = ops.concatenate((img_embeddings, text_embeddings), axis=1)
209
+ for transformer_layer in self.transformer_layers:
210
+ x = transformer_layer(
211
+ x,
212
+ padding_mask=padding_mask_input,
213
+ response_mask=response_mask_input,
214
+ )
215
+ sequence_output = self.layer_norm(x)
216
+ super().__init__(
217
+ inputs={
218
+ "images": image_input,
219
+ "token_ids": token_id_input,
220
+ "padding_mask": padding_mask_input,
221
+ "response_mask": response_mask_input,
222
+ },
223
+ outputs=sequence_output,
224
+ dtype=dtype,
225
+ **kwargs,
226
+ )
227
+
228
+ # === Config ===
229
+ self.vocabulary_size = vocabulary_size
230
+ self.image_size = image_size
231
+ self.include_rescaling = include_rescaling
232
+ self.num_layers = num_layers
233
+ self.num_query_heads = num_query_heads
234
+ self.num_key_value_heads = num_key_value_heads
235
+ self.hidden_dim = hidden_dim
236
+ self.intermediate_dim = intermediate_dim
237
+ self.head_dim = head_dim
238
+ self.layer_norm_epsilon = layer_norm_epsilon
239
+ self.dropout = dropout
240
+ # VIT Params
241
+ self.vit_patch_size = vit_patch_size
242
+ self.vit_num_heads = vit_num_heads
243
+ self.vit_hidden_dim = vit_hidden_dim
244
+ self.vit_num_layers = vit_num_layers
245
+ self.vit_intermediate_dim = vit_intermediate_dim
246
+ self.vit_pooling = vit_pooling
247
+ self.vit_classifier_activation = vit_classifier_activation
248
+ self.vit_name = vit_name
249
+ # Keep the image_sequence_length as a backbone property for easy access.
250
+ self.image_sequence_length = self.vit_encoder.image_sequence_length
251
+
252
+ def get_config(self):
253
+ config = super().get_config()
254
+ config.update(
255
+ {
256
+ "vocabulary_size": self.vocabulary_size,
257
+ "image_size": self.image_size,
258
+ "include_rescaling": self.include_rescaling,
259
+ "num_layers": self.num_layers,
260
+ "num_query_heads": self.num_query_heads,
261
+ "num_key_value_heads": self.num_key_value_heads,
262
+ "hidden_dim": self.hidden_dim,
263
+ "intermediate_dim": self.intermediate_dim,
264
+ "head_dim": self.head_dim,
265
+ "layer_norm_epsilon": self.layer_norm_epsilon,
266
+ "dropout": self.dropout,
267
+ "vit_patch_size": self.vit_patch_size,
268
+ "vit_num_heads": self.vit_num_heads,
269
+ "vit_hidden_dim": self.vit_hidden_dim,
270
+ "vit_num_layers": self.vit_num_layers,
271
+ "vit_intermediate_dim": self.vit_intermediate_dim,
272
+ "vit_pooling": self.vit_pooling,
273
+ "vit_classifier_activation": self.vit_classifier_activation,
274
+ "vit_name": self.vit_name,
275
+ }
276
+ )
277
+ return config