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,190 @@
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
+
17
+ from keras_hub.src.api_export import keras_hub_export
18
+ from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
19
+ from keras_hub.src.models.mistral.mistral_tokenizer import MistralTokenizer
20
+ from keras_hub.src.models.preprocessor import Preprocessor
21
+ from keras_hub.src.utils.keras_utils import (
22
+ convert_inputs_to_list_of_tensor_segments,
23
+ )
24
+
25
+
26
+ @keras_hub_export("keras_hub.models.MistralPreprocessor")
27
+ class MistralPreprocessor(Preprocessor):
28
+ """A Mistral preprocessing layer which tokenizes and packs inputs.
29
+
30
+ This preprocessing layer will do three things:
31
+
32
+ 1. Tokenize any number of input segments using the `tokenizer`.
33
+ 2. Pack the inputs together using a `keras_hub.layers.StartEndPacker`.
34
+ with the appropriate tokens.
35
+ 3. Construct a dictionary with keys `"token_ids"`, and `"padding_mask"`
36
+ that can be passed directly to `keras_hub.models.MistralBackbone`.
37
+
38
+ This layer can be used directly with `tf.data.Dataset.map` to preprocess
39
+ string data in the `(x, y, sample_weight)` format used by
40
+ `keras.Model.fit`.
41
+
42
+ Args:
43
+ tokenizer: A `keras_hub.models.MistralTokenizer` instance.
44
+ sequence_length: The length of the packed inputs.
45
+ add_start_token: If `True`, the preprocessor will prepend the tokenizer
46
+ start token to each input sequence. Default is `True`.
47
+ add_end_token: If `True`, the preprocessor will append the tokenizer
48
+ end token to each input sequence. Default is `False`.
49
+
50
+ Call arguments:
51
+ x: A tensor of single string sequences, or a tuple of multiple
52
+ tensor sequences to be packed together. Inputs may be batched or
53
+ unbatched. For single sequences, raw python inputs will be converted
54
+ to tensors. For multiple sequences, pass tensors directly.
55
+ y: Any label data. Will be passed through unaltered.
56
+ sample_weight: Any label weight data. Will be passed through unaltered.
57
+ sequence_length: Pass to override the configured `sequence_length` of
58
+ the layer.
59
+
60
+ Examples:
61
+
62
+ Directly calling the from_preset().
63
+ ```python
64
+ preprocessor = keras_hub.models.MistralPreprocessor.from_preset(
65
+ "mistral_base_en"
66
+ )
67
+
68
+ # Tokenize and pack a single sentence.
69
+ preprocessor("The quick brown fox jumped.")
70
+
71
+ # Tokenize and a batch of single sentences.
72
+ preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
73
+
74
+ # Preprocess a batch of sentence pairs.
75
+ # When handling multiple sequences, always convert to tensors first!
76
+ first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
77
+ second = tf.constant(["The fox tripped.", "Oh look, a whale."])
78
+ preprocessor((first, second))
79
+ ```
80
+
81
+ Mapping with `tf.data.Dataset`.
82
+ ```python
83
+ preprocessor = keras_hub.models.MistralPreprocessor.from_preset(
84
+ "mistral_base_en"
85
+ )
86
+ first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
87
+ second = tf.constant(["The fox tripped.", "Oh look, a whale."])
88
+ label = tf.constant([1, 1])
89
+
90
+ # Map labeled single sentences.
91
+ ds = tf.data.Dataset.from_tensor_slices((first, label))
92
+ ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
93
+
94
+ # Map unlabeled single sentences.
95
+ ds = tf.data.Dataset.from_tensor_slices(first)
96
+ ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
97
+
98
+ # Map labeled sentence pairs.
99
+ ds = tf.data.Dataset.from_tensor_slices(((first, second), label))
100
+ ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
101
+
102
+ # Map unlabeled sentence pairs.
103
+ ds = tf.data.Dataset.from_tensor_slices((first, second))
104
+
105
+ # Watch out for tf.data's default unpacking of tuples here!
106
+ # Best to invoke the `preprocessor` directly in this case.
107
+ ds = ds.map(
108
+ lambda first, second: preprocessor(x=(first, second)),
109
+ num_parallel_calls=tf.data.AUTOTUNE,
110
+ )
111
+ ```
112
+ """
113
+
114
+ tokenizer_cls = MistralTokenizer
115
+
116
+ def __init__(
117
+ self,
118
+ tokenizer,
119
+ sequence_length=1024,
120
+ add_start_token=True,
121
+ add_end_token=False,
122
+ **kwargs,
123
+ ):
124
+ super().__init__(**kwargs)
125
+ self.tokenizer = tokenizer
126
+ self.packer = None
127
+ self.add_start_token = add_start_token
128
+ self.add_end_token = add_end_token
129
+ self.sequence_length = sequence_length
130
+
131
+ def build(self, input_shape):
132
+ # Defer packer creation to `build()` so that we can be sure tokenizer
133
+ # assets have loaded when restoring a saved model.
134
+ self.packer = StartEndPacker(
135
+ start_value=self.tokenizer.start_token_id,
136
+ end_value=self.tokenizer.end_token_id,
137
+ sequence_length=self.sequence_length,
138
+ return_padding_mask=True,
139
+ )
140
+ self.built = True
141
+
142
+ def get_config(self):
143
+ config = super().get_config()
144
+ config.update(
145
+ {
146
+ "sequence_length": self.sequence_length,
147
+ "add_start_token": self.add_start_token,
148
+ "add_end_token": self.add_end_token,
149
+ }
150
+ )
151
+ return config
152
+
153
+ def call(
154
+ self,
155
+ x,
156
+ y=None,
157
+ sample_weight=None,
158
+ sequence_length=None,
159
+ ):
160
+ x = convert_inputs_to_list_of_tensor_segments(x)
161
+ if len(x) != 1:
162
+ raise ValueError(
163
+ "Mistral requires each input feature to contain only "
164
+ f"one segment, but received {len(x)}. If you are using Mistral"
165
+ " for a multi-segment classification task, please refer to "
166
+ "classification models like BERT or RoBERTa."
167
+ )
168
+ sequence_length = sequence_length or self.sequence_length
169
+ token_ids, padding_mask = self.packer(
170
+ self.tokenizer(x[0]),
171
+ sequence_length=sequence_length,
172
+ add_start_value=self.add_start_token,
173
+ add_end_value=self.add_end_token,
174
+ )
175
+ x = {
176
+ "token_ids": token_ids,
177
+ "padding_mask": padding_mask,
178
+ }
179
+ return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
180
+
181
+ @property
182
+ def sequence_length(self):
183
+ """The padded length of model input sequences."""
184
+ return self._sequence_length
185
+
186
+ @sequence_length.setter
187
+ def sequence_length(self, value):
188
+ self._sequence_length = value
189
+ if self.packer is not None:
190
+ self.packer.sequence_length = value
@@ -0,0 +1,48 @@
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
+ """Mistral model preset configurations."""
15
+
16
+ # Metadata for loading pretrained model weights.
17
+ backbone_presets = {
18
+ "mistral_7b_en": {
19
+ "metadata": {
20
+ "description": "Mistral 7B base model",
21
+ "params": 7241732096,
22
+ "official_name": "Mistral",
23
+ "path": "mistral",
24
+ "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
25
+ },
26
+ "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/6",
27
+ },
28
+ "mistral_instruct_7b_en": {
29
+ "metadata": {
30
+ "description": "Mistral 7B instruct model",
31
+ "params": 7241732096,
32
+ "official_name": "Mistral",
33
+ "path": "mistral",
34
+ "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
35
+ },
36
+ "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/6",
37
+ },
38
+ "mistral_0.2_instruct_7b_en": {
39
+ "metadata": {
40
+ "description": "Mistral 7B instruct Version 0.2 model",
41
+ "params": 7241732096,
42
+ "official_name": "Mistral",
43
+ "path": "mistral",
44
+ "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
45
+ },
46
+ "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.2_instruct_7b_en/1",
47
+ },
48
+ }
@@ -0,0 +1,82 @@
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.api_export import keras_hub_export
16
+ from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
17
+ SentencePieceTokenizer,
18
+ )
19
+
20
+
21
+ @keras_hub_export("keras_hub.models.MistralTokenizer")
22
+ class MistralTokenizer(SentencePieceTokenizer):
23
+ """Mistral tokenizer layer based on SentencePiece.
24
+
25
+ This tokenizer class will tokenize raw strings into integer sequences and
26
+ is based on `keras_hub.tokenizers.SentencePieceTokenizer`. Unlike the
27
+ underlying tokenizer, it will check for all special tokens needed by
28
+ Mistral models and provides a `from_preset()` method to automatically
29
+ download a matching vocabulary for a Mistral preset.
30
+
31
+ This tokenizer does not provide truncation or padding of inputs. It can be
32
+ combined with a `keras_hub.models.MistralPreprocessor` layer for input
33
+ packing.
34
+
35
+ If input is a batch of strings (rank > 0), the layer will output a
36
+ `tf.RaggedTensor` where the last dimension of the output is ragged.
37
+
38
+ If input is a scalar string (rank == 0), the layer will output a dense
39
+ `tf.Tensor` with static shape `[None]`.
40
+
41
+ Args:
42
+ proto: Either a `string` path to a SentencePiece proto file, or a
43
+ `bytes` object with a serialized SentencePiece proto. See the
44
+ [SentencePiece repository](https://github.com/google/sentencepiece)
45
+ for more details on the format.
46
+
47
+ Examples:
48
+ ```python
49
+ # Unbatched input.
50
+ tokenizer = keras_hub.models.MistralTokenizer.from_preset(
51
+ "mistral_7b_en",
52
+ )
53
+ tokenizer("The quick brown fox jumped.")
54
+
55
+ # Batched input.
56
+ tokenizer(["The quick brown fox jumped.", "The fox slept."])
57
+
58
+ # Detokenization.
59
+ tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
60
+ ```
61
+ """
62
+
63
+ def __init__(self, proto, **kwargs):
64
+ self.start_token = "<s>"
65
+ self.end_token = "</s>"
66
+ super().__init__(proto=proto, **kwargs)
67
+
68
+ def set_proto(self, proto):
69
+ super().set_proto(proto)
70
+ if proto is not None:
71
+ for token in [self.start_token, self.end_token]:
72
+ if token not in self.get_vocabulary():
73
+ raise ValueError(
74
+ f"Cannot find token `'{token}'` in the provided "
75
+ f"`vocabulary`. Please provide `'{token}'` in your "
76
+ "`vocabulary` or use a pretrained `vocabulary` name."
77
+ )
78
+ self.start_token_id = self.token_to_id(self.start_token)
79
+ self.end_token_id = self.token_to_id(self.end_token)
80
+ else:
81
+ self.start_token_id = None
82
+ self.end_token_id = None
@@ -0,0 +1,265 @@
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.transformer_layer_utils import (
18
+ compute_causal_mask,
19
+ )
20
+ from keras_hub.src.layers.modeling.transformer_layer_utils import (
21
+ merge_padding_and_attention_mask,
22
+ )
23
+ from keras_hub.src.models.mistral.mistral_attention import (
24
+ CachedMistralAttention,
25
+ )
26
+ from keras_hub.src.models.mistral.mistral_layer_norm import (
27
+ MistralLayerNormalization,
28
+ )
29
+ from keras_hub.src.utils.keras_utils import clone_initializer
30
+
31
+
32
+ class MistralTransformerDecoder(keras.layers.Layer):
33
+ """A Transformer decoder layer for the Mistral backbone."""
34
+
35
+ def __init__(
36
+ self,
37
+ intermediate_dim,
38
+ num_query_heads,
39
+ num_key_value_heads,
40
+ rope_max_wavelength=10000,
41
+ rope_scaling_factor=1.0,
42
+ activation="silu",
43
+ layer_norm_epsilon=1e-5,
44
+ kernel_initializer="glorot_uniform",
45
+ sliding_window=512,
46
+ dropout=0,
47
+ **kwargs,
48
+ ):
49
+ super().__init__(**kwargs)
50
+ self.intermediate_dim = intermediate_dim
51
+ self.num_query_heads = num_query_heads
52
+ self.num_key_value_heads = num_key_value_heads
53
+
54
+ self.rope_max_wavelength = rope_max_wavelength
55
+ self.rope_scaling_factor = rope_scaling_factor
56
+
57
+ self.dropout = dropout
58
+
59
+ self.sliding_window = sliding_window
60
+ self.activation = keras.activations.get(activation)
61
+ self.layer_norm_epsilon = layer_norm_epsilon
62
+ self.kernel_initializer = keras.initializers.get(kernel_initializer)
63
+
64
+ self.supports_masking = True
65
+
66
+ def build(self, decoder_sequence_shape):
67
+ self._decoder_sequence_shape = decoder_sequence_shape
68
+ self.hidden_dim = decoder_sequence_shape[-1]
69
+
70
+ # Self attention layer.
71
+ self._self_attention_layer = CachedMistralAttention(
72
+ num_query_heads=self.num_query_heads,
73
+ num_key_value_heads=self.num_key_value_heads,
74
+ rope_max_wavelength=self.rope_max_wavelength,
75
+ rope_scaling_factor=self.rope_scaling_factor,
76
+ sliding_window=self.sliding_window,
77
+ kernel_initializer=clone_initializer(self.kernel_initializer),
78
+ dropout=self.dropout,
79
+ dtype=self.dtype_policy,
80
+ name="self_attention",
81
+ )
82
+ self._self_attention_layer.build(decoder_sequence_shape)
83
+
84
+ self._self_attention_layernorm = MistralLayerNormalization(
85
+ epsilon=self.layer_norm_epsilon,
86
+ dtype=self.dtype_policy,
87
+ name="self_attention_layernorm",
88
+ )
89
+ self._self_attention_layernorm.build(decoder_sequence_shape)
90
+ self._self_attention_dropout = keras.layers.Dropout(
91
+ rate=self.dropout,
92
+ dtype=self.dtype_policy,
93
+ name="self_attention_dropout",
94
+ )
95
+
96
+ # Feedforward layers.
97
+ self._feedforward_intermediate_dense = keras.layers.Dense(
98
+ self.intermediate_dim,
99
+ kernel_initializer=clone_initializer(self.kernel_initializer),
100
+ use_bias=False,
101
+ dtype=self.dtype_policy,
102
+ name="feedforward_intermediate_dense",
103
+ )
104
+ self._feedforward_intermediate_dense.build(decoder_sequence_shape)
105
+
106
+ self._feedforward_gate_dense = keras.layers.Dense(
107
+ self.intermediate_dim,
108
+ kernel_initializer=clone_initializer(self.kernel_initializer),
109
+ use_bias=False,
110
+ dtype=self.dtype_policy,
111
+ name="feedforward_gate_dense",
112
+ )
113
+ self._feedforward_gate_dense.build(decoder_sequence_shape)
114
+
115
+ self._feedforward_output_dense = keras.layers.Dense(
116
+ self.hidden_dim,
117
+ kernel_initializer=clone_initializer(self.kernel_initializer),
118
+ use_bias=False,
119
+ dtype=self.dtype_policy,
120
+ name="feedforward_output_dense",
121
+ )
122
+
123
+ self._feedforward_output_dense.build(
124
+ self._feedforward_gate_dense.compute_output_shape(
125
+ decoder_sequence_shape
126
+ )
127
+ )
128
+
129
+ self._feedforward_layernorm = MistralLayerNormalization(
130
+ epsilon=self.layer_norm_epsilon,
131
+ dtype=self.dtype_policy,
132
+ name="feedforward_layernorm",
133
+ )
134
+ self._feedforward_layernorm.build(decoder_sequence_shape)
135
+
136
+ self.built = True
137
+
138
+ def call(
139
+ self,
140
+ decoder_sequence,
141
+ decoder_padding_mask=None,
142
+ decoder_attention_mask=None,
143
+ self_attention_cache=None,
144
+ self_attention_cache_update_index=None,
145
+ training=None,
146
+ ):
147
+ self_attention_mask = self._compute_self_attention_mask(
148
+ decoder_sequence=decoder_sequence,
149
+ decoder_padding_mask=decoder_padding_mask,
150
+ decoder_attention_mask=decoder_attention_mask,
151
+ self_attention_cache=self_attention_cache,
152
+ self_attention_cache_update_index=self_attention_cache_update_index,
153
+ )
154
+ residual = decoder_sequence
155
+
156
+ x = self._self_attention_layernorm(decoder_sequence)
157
+
158
+ # Self attention block.
159
+ x = self._self_attention_layer(
160
+ hidden_states=x,
161
+ attention_mask=self_attention_mask,
162
+ cache=self_attention_cache,
163
+ cache_update_index=self_attention_cache_update_index,
164
+ )
165
+
166
+ if self_attention_cache is not None:
167
+ x, self_attention_cache = x
168
+
169
+ x = self._self_attention_dropout(x, training=training)
170
+
171
+ x = x + residual
172
+ residual = x
173
+
174
+ x = self._feedforward_layernorm(x)
175
+ gate_output = self._feedforward_gate_dense(x)
176
+
177
+ # Note that we run the activation function in full 32-bit
178
+ # precision since this is what `torch.nn.functional.silu`
179
+ # does. Internally, `torch.nn.functional.silu` converts the
180
+ # inputs to float32, computes SiLU, and converts the outputs
181
+ # back to compute dtype.
182
+ # CPU Kernel: https://github.com/pytorch/pytorch/blob/35c493f2cf9b623bfdc7e6b34dc1cb39690a7919/aten/src/ATen/native/cpu/Activation.cpp#L1221-L1235 # noqa: E501
183
+ # CUDA Kernel: https://github.com/pytorch/pytorch/blob/35c493f2cf9b623bfdc7e6b34dc1cb39690a7919/aten/src/ATen/native/cuda/ActivationSiluKernel.cu # noqa: E501
184
+ gate_output = ops.cast(gate_output, "float32")
185
+ gate_output = self.activation(gate_output)
186
+ gate_output = ops.cast(gate_output, self.compute_dtype)
187
+
188
+ x = self._feedforward_intermediate_dense(x)
189
+
190
+ x = self._feedforward_output_dense(ops.multiply(x, gate_output))
191
+
192
+ decoder_output = x + residual
193
+
194
+ if self_attention_cache is not None:
195
+ return decoder_output, self_attention_cache
196
+ return decoder_output
197
+
198
+ def _compute_self_attention_mask(
199
+ self,
200
+ decoder_sequence,
201
+ decoder_padding_mask,
202
+ decoder_attention_mask,
203
+ self_attention_cache,
204
+ self_attention_cache_update_index,
205
+ ):
206
+ decoder_mask = merge_padding_and_attention_mask(
207
+ decoder_sequence, decoder_padding_mask, decoder_attention_mask
208
+ )
209
+ batch_size = ops.shape(decoder_sequence)[0]
210
+ input_length = output_length = ops.shape(decoder_sequence)[1]
211
+ # We need to handle a rectangular causal mask when doing cached
212
+ # decoding. For generative inference, `decoder_sequence` will
213
+ # generally be length 1, and `cache` will be the full generation length.
214
+ if self_attention_cache is not None:
215
+ input_length = ops.shape(self_attention_cache)[2]
216
+
217
+ cache_update_index = (
218
+ 0
219
+ if self_attention_cache_update_index is None
220
+ else self_attention_cache_update_index
221
+ )
222
+
223
+ # The lower traingular attention mask
224
+ causal_mask = compute_causal_mask(
225
+ batch_size, input_length, output_length, cache_update_index
226
+ )
227
+
228
+ # Mistral uses a banded attention mask if sliding window is not None
229
+ if self.sliding_window is not None:
230
+ # Below is a workaround for `ops.triu` for Keras 2.
231
+ # TODO(tirthasheshpatel): Use `ops.triu` once Keras 2 support is removed.
232
+ # causal_mask = ops.triu(causal_mask, k=-self.sliding_window)
233
+ i = ops.arange(output_length)[:, None] + cache_update_index
234
+ j = ops.arange(input_length)[None, :]
235
+ causal_mask_upper = ops.cast(i < j + self.sliding_window, "int32")
236
+ causal_mask = ops.minimum(causal_mask, causal_mask_upper)
237
+
238
+ return (
239
+ ops.minimum(decoder_mask, causal_mask)
240
+ if decoder_mask is not None
241
+ else causal_mask
242
+ )
243
+
244
+ def compute_output_shape(self, decoder_sequence_shape):
245
+ return decoder_sequence_shape
246
+
247
+ def get_config(self):
248
+ config = super().get_config()
249
+ config.update(
250
+ {
251
+ "intermediate_dim": self.intermediate_dim,
252
+ "num_query_heads": self.num_query_heads,
253
+ "rope_max_wavelength": self.rope_max_wavelength,
254
+ "rope_scaling_factor": self.rope_scaling_factor,
255
+ "num_key_value_heads": self.num_key_value_heads,
256
+ "sliding_window": self.sliding_window,
257
+ "activation": keras.activations.serialize(self.activation),
258
+ "layer_norm_epsilon": self.layer_norm_epsilon,
259
+ "kernel_initializer": keras.initializers.serialize(
260
+ self.kernel_initializer
261
+ ),
262
+ "dropout": self.dropout,
263
+ }
264
+ )
265
+ return config
@@ -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.