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,81 @@
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.models.roberta import roberta_backbone
18
+
19
+
20
+ @keras_hub_export("keras_hub.models.XLMRobertaBackbone")
21
+ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
22
+ """An XLM-RoBERTa encoder network.
23
+
24
+ This class implements a bi-directional Transformer-based encoder as
25
+ described in ["Unsupervised Cross-lingual Representation Learning at Scale"](https://arxiv.org/abs/1911.02116).
26
+ It includes the embedding lookups and transformer layers, but it does not
27
+ include the masked language modeling head used during pretraining.
28
+
29
+ The default constructor gives a fully customizable, randomly initialized
30
+ RoBERTa encoder with any number of layers, heads, and embedding dimensions.
31
+ To load preset architectures and weights, use the `from_preset()`
32
+ constructor.
33
+
34
+ Disclaimer: Pre-trained models are provided on an "as is" basis, without
35
+ warranties or conditions of any kind. The underlying model is provided by a
36
+ third party and subject to a separate license, available
37
+ [here](https://github.com/facebookresearch/fairseq).
38
+
39
+ Args:
40
+ vocabulary_size: int. The size of the token vocabulary.
41
+ num_layers: int. The number of transformer layers.
42
+ num_heads: int. The number of attention heads for each transformer.
43
+ The hidden size must be divisible by the number of attention heads.
44
+ hidden_dim: int. The size of the transformer encoding layer.
45
+ intermediate_dim: int. The output dimension of the first Dense layer in
46
+ a two-layer feedforward network for each transformer.
47
+ dropout: float. Dropout probability for the Transformer encoder.
48
+ max_sequence_length: int. The maximum sequence length this encoder can
49
+ consume. The sequence length of the input must be less than
50
+ `max_sequence_length` default value. This determines the variable
51
+ shape for positional embeddings.
52
+ dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
53
+ for model computations and weights. Note that some computations,
54
+ such as softmax and layer normalization, will always be done at
55
+ float32 precision regardless of dtype.
56
+
57
+ Examples:
58
+ ```python
59
+ input_data = {
60
+ "token_ids": np.ones(shape=(1, 12), dtype="int32"),
61
+ "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
62
+ }
63
+
64
+ # Pretrained XLM-R encoder.
65
+ model = keras_hub.models.XLMRobertaBackbone.from_preset(
66
+ "xlm_roberta_base_multi",
67
+ )
68
+ model(input_data)
69
+
70
+ # Randomly initialized XLM-R model with custom config.
71
+ model = keras_hub.models.XLMRobertaBackbone(
72
+ vocabulary_size=250002,
73
+ num_layers=4,
74
+ num_heads=4,
75
+ hidden_dim=256,
76
+ intermediate_dim=512,
77
+ max_sequence_length=128
78
+ )
79
+ model(input_data)
80
+ ```
81
+ """
@@ -0,0 +1,225 @@
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.models.classifier import Classifier
20
+ from keras_hub.src.models.roberta.roberta_backbone import (
21
+ roberta_kernel_initializer,
22
+ )
23
+ from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
24
+ XLMRobertaBackbone,
25
+ )
26
+ from keras_hub.src.models.xlm_roberta.xlm_roberta_preprocessor import (
27
+ XLMRobertaPreprocessor,
28
+ )
29
+
30
+
31
+ @keras_hub_export("keras_hub.models.XLMRobertaClassifier")
32
+ class XLMRobertaClassifier(Classifier):
33
+ """An end-to-end XLM-RoBERTa model for classification tasks.
34
+
35
+ This model attaches a classification head to a
36
+ `keras_hub.model.XLMRobertaBackbone` instance, mapping from the backbone
37
+ outputs to logits suitable for a classification task. For usage of
38
+ this model with pre-trained weights, see the `from_preset()` constructor.
39
+
40
+ This model can optionally be configured with a `preprocessor` layer, in
41
+ which case it will automatically apply preprocessing to raw inputs during
42
+ `fit()`, `predict()`, and `evaluate()`. This is done by default when
43
+ creating the model with `from_preset()`.
44
+
45
+ Disclaimer: Pre-trained models are provided on an "as is" basis, without
46
+ warranties or conditions of any kind. The underlying model is provided by a
47
+ third party and subject to a separate license, available
48
+ [here](https://github.com/facebookresearch/fairseq).
49
+
50
+ Args:
51
+ backbone: A `keras_hub.models.XLMRobertaBackbone` instance.
52
+ num_classes: int. Number of classes to predict.
53
+ preprocessor: A `keras_hub.models.XLMRobertaPreprocessor` or `None`. If
54
+ `None`, this model will not apply preprocessing, and inputs should
55
+ be preprocessed before calling the model.
56
+ activation: Optional `str` or callable. The activation function to use
57
+ on the model outputs. Set `activation="softmax"` to return output
58
+ probabilities. Defaults to `None`.
59
+ hidden_dim: int. The size of the pooler layer.
60
+ dropout: float. The dropout probability value, applied to the pooled
61
+ output, and after the first dense layer.
62
+
63
+ Examples:
64
+
65
+ Raw string data.
66
+ ```python
67
+ features = ["The quick brown fox jumped.", "نسيت الواجب"]
68
+ labels = [0, 3]
69
+
70
+ # Pretrained classifier.
71
+ classifier = keras_hub.models.XLMRobertaClassifier.from_preset(
72
+ "xlm_roberta_base_multi",
73
+ num_classes=4,
74
+ )
75
+ classifier.fit(x=features, y=labels, batch_size=2)
76
+ classifier.predict(x=features, batch_size=2)
77
+
78
+ # Re-compile (e.g., with a new learning rate).
79
+ classifier.compile(
80
+ loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
81
+ optimizer=keras.optimizers.Adam(5e-5),
82
+ jit_compile=True,
83
+ )
84
+ # Access backbone programmatically (e.g., to change `trainable`).
85
+ classifier.backbone.trainable = False
86
+ # Fit again.
87
+ classifier.fit(x=features, y=labels, batch_size=2)
88
+ ```
89
+
90
+ Preprocessed integer data.
91
+ ```python
92
+ features = {
93
+ "token_ids": np.ones(shape=(2, 12), dtype="int32"),
94
+ "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2),
95
+ }
96
+ labels = [0, 3]
97
+
98
+ # Pretrained classifier without preprocessing.
99
+ classifier = keras_hub.models.XLMRobertaClassifier.from_preset(
100
+ "xlm_roberta_base_multi",
101
+ num_classes=4,
102
+ preprocessor=None,
103
+ )
104
+ classifier.fit(x=features, y=labels, batch_size=2)
105
+ ```
106
+
107
+ Custom backbone and vocabulary.
108
+ ```python
109
+ features = ["The quick brown fox jumped.", "نسيت الواجب"]
110
+ labels = [0, 3]
111
+
112
+ def train_sentencepiece(ds, vocab_size):
113
+ bytes_io = io.BytesIO()
114
+ sentencepiece.SentencePieceTrainer.train(
115
+ sentence_iterator=ds.as_numpy_iterator(),
116
+ model_writer=bytes_io,
117
+ vocab_size=vocab_size,
118
+ model_type="WORD",
119
+ unk_id=0,
120
+ bos_id=1,
121
+ eos_id=2,
122
+ )
123
+ return bytes_io.getvalue()
124
+ ds = tf.data.Dataset.from_tensor_slices(
125
+ ["the quick brown fox", "the earth is round"]
126
+ )
127
+ proto = train_sentencepiece(ds, vocab_size=10)
128
+ tokenizer = keras_hub.models.XLMRobertaTokenizer(
129
+ proto=proto
130
+ )
131
+ preprocessor = keras_hub.models.XLMRobertaPreprocessor(
132
+ tokenizer,
133
+ sequence_length=128,
134
+ )
135
+ backbone = keras_hub.models.XLMRobertaBackbone(
136
+ vocabulary_size=250002,
137
+ num_layers=4,
138
+ num_heads=4,
139
+ hidden_dim=256,
140
+ intermediate_dim=512,
141
+ max_sequence_length=128,
142
+ )
143
+ classifier = keras_hub.models.XLMRobertaClassifier(
144
+ backbone=backbone,
145
+ preprocessor=preprocessor,
146
+ num_classes=4,
147
+ )
148
+ classifier.fit(x=features, y=labels, batch_size=2)
149
+ ```
150
+ """
151
+
152
+ backbone_cls = XLMRobertaBackbone
153
+ preprocessor_cls = XLMRobertaPreprocessor
154
+
155
+ def __init__(
156
+ self,
157
+ backbone,
158
+ num_classes,
159
+ preprocessor=None,
160
+ activation=None,
161
+ hidden_dim=None,
162
+ dropout=0.0,
163
+ **kwargs,
164
+ ):
165
+ # === Layers ===
166
+ self.backbone = backbone
167
+ self.preprocessor = preprocessor
168
+ self.pooled_dropout = keras.layers.Dropout(
169
+ dropout,
170
+ dtype=backbone.dtype_policy,
171
+ name="pooled_dropout",
172
+ )
173
+ hidden_dim = hidden_dim or backbone.hidden_dim
174
+ self.pooled_dense = keras.layers.Dense(
175
+ hidden_dim,
176
+ activation="tanh",
177
+ dtype=backbone.dtype_policy,
178
+ name="pooled_dense",
179
+ )
180
+ self.output_dropout = keras.layers.Dropout(
181
+ dropout,
182
+ dtype=backbone.dtype_policy,
183
+ name="output_dropout",
184
+ )
185
+ self.output_dense = keras.layers.Dense(
186
+ num_classes,
187
+ kernel_initializer=roberta_kernel_initializer(),
188
+ activation=activation,
189
+ dtype=backbone.dtype_policy,
190
+ name="logits",
191
+ )
192
+
193
+ # === Functional Model ===
194
+ inputs = backbone.input
195
+ x = backbone(inputs)[:, backbone.start_token_index, :]
196
+ x = self.pooled_dropout(x)
197
+ x = self.pooled_dense(x)
198
+ x = self.output_dropout(x)
199
+ outputs = self.output_dense(x)
200
+ # Instantiate using Functional API Model constructor
201
+ super().__init__(
202
+ inputs=inputs,
203
+ outputs=outputs,
204
+ **kwargs,
205
+ )
206
+
207
+ # === Config ===
208
+ self.backbone = backbone
209
+ self.preprocessor = preprocessor
210
+ self.num_classes = num_classes
211
+ self.activation = keras.activations.get(activation)
212
+ self.hidden_dim = hidden_dim
213
+ self.dropout = dropout
214
+
215
+ def get_config(self):
216
+ config = super().get_config()
217
+ config.update(
218
+ {
219
+ "num_classes": self.num_classes,
220
+ "activation": keras.activations.serialize(self.activation),
221
+ "hidden_dim": self.hidden_dim,
222
+ "dropout": self.dropout,
223
+ }
224
+ )
225
+ return config
@@ -0,0 +1,141 @@
1
+ # Copyright 2024 The KerasHub Authors
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+
16
+ import keras
17
+
18
+ from keras_hub.src.api_export import keras_hub_export
19
+ from keras_hub.src.layers.modeling.masked_lm_head import MaskedLMHead
20
+ from keras_hub.src.models.masked_lm import MaskedLM
21
+ from keras_hub.src.models.roberta.roberta_backbone import (
22
+ roberta_kernel_initializer,
23
+ )
24
+ from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
25
+ XLMRobertaBackbone,
26
+ )
27
+ from keras_hub.src.models.xlm_roberta.xlm_roberta_masked_lm_preprocessor import (
28
+ XLMRobertaMaskedLMPreprocessor,
29
+ )
30
+
31
+
32
+ @keras_hub_export("keras_hub.models.XLMRobertaMaskedLM")
33
+ class XLMRobertaMaskedLM(MaskedLM):
34
+ """An end-to-end XLM-RoBERTa model for the masked language modeling task.
35
+
36
+ This model will train XLM-RoBERTa on a masked language modeling task.
37
+ The model will predict labels for a number of masked tokens in the
38
+ input data. For usage of this model with pre-trained weights, see the
39
+ `from_preset()` method.
40
+
41
+ This model can optionally be configured with a `preprocessor` layer, in
42
+ which case inputs can be raw string features during `fit()`, `predict()`,
43
+ and `evaluate()`. Inputs will be tokenized and dynamically masked during
44
+ training and evaluation. This is done by default when creating the model
45
+ with `from_preset()`.
46
+
47
+ Disclaimer: Pre-trained models are provided on an "as is" basis, without
48
+ warranties or conditions of any kind. The underlying model is provided by a
49
+ third party and subject to a separate license, available
50
+ [here](https://github.com/facebookresearch/fairseq).
51
+
52
+ Args:
53
+ backbone: A `keras_hub.models.XLMRobertaBackbone` instance.
54
+ preprocessor: A `keras_hub.models.XLMRobertaMaskedLMPreprocessor` or
55
+ `None`. If `None`, this model will not apply preprocessing, and
56
+ inputs should be preprocessed before calling the model.
57
+
58
+ Examples:
59
+
60
+ Raw string inputs and pretrained backbone.
61
+ ```python
62
+ # Create a dataset with raw string features. Labels are inferred.
63
+ features = ["The quick brown fox jumped.", "I forgot my homework."]
64
+
65
+ # Pretrained language model
66
+ # on an MLM task.
67
+ masked_lm = keras_hub.models.XLMRobertaMaskedLM.from_preset(
68
+ "xlm_roberta_base_multi",
69
+ )
70
+ masked_lm.fit(x=features, batch_size=2)
71
+ ```
72
+
73
+ # Re-compile (e.g., with a new learning rate).
74
+ masked_lm.compile(
75
+ loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
76
+ optimizer=keras.optimizers.Adam(5e-5),
77
+ jit_compile=True,
78
+ )
79
+ # Access backbone programmatically (e.g., to change `trainable`).
80
+ masked_lm.backbone.trainable = False
81
+ # Fit again.
82
+ masked_lm.fit(x=features, batch_size=2)
83
+ ```
84
+
85
+ Preprocessed integer data.
86
+ ```python
87
+ # Create a preprocessed dataset where 0 is the mask token.
88
+ features = {
89
+ "token_ids": np.array([[1, 2, 0, 4, 0, 6, 7, 8]] * 2),
90
+ "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1]] * 2),
91
+ "mask_positions": np.array([[2, 4]] * 2)
92
+ }
93
+ # Labels are the original masked values.
94
+ labels = [[3, 5]] * 2
95
+
96
+ masked_lm = keras_hub.models.XLMRobertaMaskedLM.from_preset(
97
+ "xlm_roberta_base_multi",
98
+ preprocessor=None,
99
+ )
100
+
101
+ masked_lm.fit(x=features, y=labels, batch_size=2)
102
+ ```
103
+ """
104
+
105
+ backbone_cls = XLMRobertaBackbone
106
+ preprocessor_cls = XLMRobertaMaskedLMPreprocessor
107
+
108
+ def __init__(
109
+ self,
110
+ backbone,
111
+ preprocessor=None,
112
+ **kwargs,
113
+ ):
114
+ # === Layers ===
115
+ self.backbone = backbone
116
+ self.preprocessor = preprocessor
117
+ self.masked_lm_head = MaskedLMHead(
118
+ vocabulary_size=backbone.vocabulary_size,
119
+ token_embedding=backbone.token_embedding,
120
+ intermediate_activation="gelu",
121
+ kernel_initializer=roberta_kernel_initializer(),
122
+ dtype=backbone.dtype_policy,
123
+ name="mlm_head",
124
+ )
125
+
126
+ # === Functional Model ===
127
+ inputs = {
128
+ **backbone.input,
129
+ "mask_positions": keras.Input(
130
+ shape=(None,), dtype="int32", name="mask_positions"
131
+ ),
132
+ }
133
+ backbone_outputs = backbone(backbone.input)
134
+ outputs = self.masked_lm_head(
135
+ backbone_outputs, inputs["mask_positions"]
136
+ )
137
+ super().__init__(
138
+ inputs=inputs,
139
+ outputs=outputs,
140
+ **kwargs,
141
+ )
@@ -0,0 +1,195 @@
1
+ # Copyright 2024 The KerasHub Authors
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import keras
16
+ from absl import logging
17
+
18
+ from keras_hub.src.api_export import keras_hub_export
19
+ from keras_hub.src.layers.preprocessing.masked_lm_mask_generator import (
20
+ MaskedLMMaskGenerator,
21
+ )
22
+ from keras_hub.src.models.xlm_roberta.xlm_roberta_preprocessor import (
23
+ XLMRobertaPreprocessor,
24
+ )
25
+
26
+
27
+ @keras_hub_export("keras_hub.models.XLMRobertaMaskedLMPreprocessor")
28
+ class XLMRobertaMaskedLMPreprocessor(XLMRobertaPreprocessor):
29
+ """XLM-RoBERTa preprocessing for the masked language modeling task.
30
+
31
+ This preprocessing layer will prepare inputs for a masked language modeling
32
+ task. It is primarily intended for use with the
33
+ `keras_hub.models.XLMRobertaMaskedLM` task model. Preprocessing will occur in
34
+ multiple steps.
35
+
36
+ 1. Tokenize any number of input segments using the `tokenizer`.
37
+ 2. Pack the inputs together with the appropriate `"<s>"`, `"</s>"` and
38
+ `"<pad>"` tokens, i.e., adding a single `"<s>"` at the start of the
39
+ entire sequence, `"</s></s>"` between each segment,
40
+ and a `"</s>"` at the end of the entire sequence.
41
+ 3. Randomly select non-special tokens to mask, controlled by
42
+ `mask_selection_rate`.
43
+ 4. Construct a `(x, y, sample_weight)` tuple suitable for training with a
44
+ `keras_hub.models.XLMRobertaMaskedLM` task model.
45
+
46
+ Args:
47
+ tokenizer: A `keras_hub.models.XLMRobertaTokenizer` instance.
48
+ sequence_length: int. The length of the packed inputs.
49
+ truncate: string. The algorithm to truncate a list of batched segments
50
+ to fit within `sequence_length`. The value can be either
51
+ `round_robin` or `waterfall`:
52
+ - `"round_robin"`: Available space is assigned one token at a
53
+ time in a round-robin fashion to the inputs that still need
54
+ some, until the limit is reached.
55
+ - `"waterfall"`: The allocation of the budget is done using a
56
+ "waterfall" algorithm that allocates quota in a
57
+ left-to-right manner and fills up the buckets until we run
58
+ out of budget. It supports an arbitrary number of segments.
59
+ mask_selection_rate: float. The probability an input token will be
60
+ dynamically masked.
61
+ mask_selection_length: int. The maximum number of masked tokens
62
+ in a given sample.
63
+ mask_token_rate: float. The probability the a selected token will be
64
+ replaced with the mask token.
65
+ random_token_rate: float. The probability the a selected token will be
66
+ replaced with a random token from the vocabulary. A selected token
67
+ will be left as is with probability
68
+ `1 - mask_token_rate - random_token_rate`.
69
+
70
+ Call arguments:
71
+ x: A tensor of single string sequences, or a tuple of multiple
72
+ tensor sequences to be packed together. Inputs may be batched or
73
+ unbatched. For single sequences, raw python inputs will be converted
74
+ to tensors. For multiple sequences, pass tensors directly.
75
+ y: Label data. Should always be `None` as the layer generates labels.
76
+ sample_weight: Label weights. Should always be `None` as the layer
77
+ generates label weights.
78
+
79
+ Examples:
80
+
81
+ Directly calling the layer on data.
82
+ ```python
83
+ # Load the preprocessor from a preset.
84
+ preprocessor = keras_hub.models.XLMRobertaMaskedLMPreprocessor.from_preset(
85
+ "xlm_roberta_base_multi"
86
+ )
87
+
88
+ # Tokenize and mask a single sentence.
89
+ preprocessor("The quick brown fox jumped.")
90
+ # Tokenize and mask a batch of single sentences.
91
+ preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
92
+ # Tokenize and mask sentence pairs.
93
+ # In this case, always convert input to tensors before calling the layer.
94
+ first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
95
+ second = tf.constant(["The fox tripped.", "Oh look, a whale."])
96
+ preprocessor((first, second))
97
+ ```
98
+
99
+ Mapping with `tf.data.Dataset`.
100
+ ```python
101
+ preprocessor = keras_hub.models.XLMRobertaMaskedLMPreprocessor.from_preset(
102
+ "xlm_roberta_base_multi"
103
+ )
104
+ first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
105
+ second = tf.constant(["The fox tripped.", "Oh look, a whale."])
106
+
107
+ # Map single sentences.
108
+ ds = tf.data.Dataset.from_tensor_slices(first)
109
+ ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
110
+
111
+ # Map sentence pairs.
112
+ ds = tf.data.Dataset.from_tensor_slices((first, second))
113
+ # Watch out for tf.data's default unpacking of tuples here!
114
+ # Best to invoke the `preprocessor` directly in this case.
115
+ ds = ds.map(
116
+ lambda first, second: preprocessor(x=(first, second)),
117
+ num_parallel_calls=tf.data.AUTOTUNE,
118
+ )
119
+ ```
120
+ ```
121
+ """
122
+
123
+ def __init__(
124
+ self,
125
+ tokenizer,
126
+ sequence_length=512,
127
+ truncate="round_robin",
128
+ mask_selection_rate=0.15,
129
+ mask_selection_length=96,
130
+ mask_token_rate=0.8,
131
+ random_token_rate=0.1,
132
+ **kwargs,
133
+ ):
134
+ super().__init__(
135
+ tokenizer,
136
+ sequence_length=sequence_length,
137
+ truncate=truncate,
138
+ **kwargs,
139
+ )
140
+ self.mask_selection_rate = mask_selection_rate
141
+ self.mask_selection_length = mask_selection_length
142
+ self.mask_token_rate = mask_token_rate
143
+ self.random_token_rate = random_token_rate
144
+ self.masker = None
145
+
146
+ def build(self, input_shape):
147
+ super().build(input_shape)
148
+ # Defer masker creation to `build()` so that we can be sure tokenizer
149
+ # assets have loaded when restoring a saved model.
150
+ self.masker = MaskedLMMaskGenerator(
151
+ mask_selection_rate=self.mask_selection_rate,
152
+ mask_selection_length=self.mask_selection_length,
153
+ mask_token_rate=self.mask_token_rate,
154
+ random_token_rate=self.random_token_rate,
155
+ vocabulary_size=self.tokenizer.vocabulary_size(),
156
+ mask_token_id=self.tokenizer.mask_token_id,
157
+ unselectable_token_ids=[
158
+ self.tokenizer.start_token_id,
159
+ self.tokenizer.end_token_id,
160
+ self.tokenizer.pad_token_id,
161
+ ],
162
+ )
163
+
164
+ def get_config(self):
165
+ config = super().get_config()
166
+ config.update(
167
+ {
168
+ "mask_selection_rate": self.mask_selection_rate,
169
+ "mask_selection_length": self.mask_selection_length,
170
+ "mask_token_rate": self.mask_token_rate,
171
+ "random_token_rate": self.random_token_rate,
172
+ }
173
+ )
174
+ return config
175
+
176
+ def call(self, x, y=None, sample_weight=None):
177
+ if y is not None or sample_weight is not None:
178
+ logging.warning(
179
+ f"{self.__class__.__name__} generates `y` and `sample_weight` "
180
+ "based on your input data, but your data already contains `y` "
181
+ "or `sample_weight`. Your `y` and `sample_weight` will be "
182
+ "ignored."
183
+ )
184
+
185
+ x = super().call(x)
186
+ token_ids, padding_mask = x["token_ids"], x["padding_mask"]
187
+ masker_outputs = self.masker(token_ids)
188
+ x = {
189
+ "token_ids": masker_outputs["token_ids"],
190
+ "padding_mask": padding_mask,
191
+ "mask_positions": masker_outputs["mask_positions"],
192
+ }
193
+ y = masker_outputs["mask_ids"]
194
+ sample_weight = masker_outputs["mask_weights"]
195
+ return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)