Add template for adding flax models (#12441)

* Add option to add flax

* Add flax template for __init__.py

* Add flax template for .rst

* Copy TF modeling template

* Add a missing line in modeling_tf_... template

* Update first half of modeling_flax_..

* Update encoder flax template

* Copy test_modeling_tf... as test_modeling_flax...

* Replace some TF to Flax in test_modeling_flax_...

* Replace tf to np

some function might not work, like _assert_tensors_equal

* Replace remaining tf to np (might not work)

* Fix cookiecutter

* Add Flax in to_replace_... template

* Update transformers-cli add-new-model

* Save generate_flax in configuration.json

This will be read by transformers-cli

* Fix to_replace_... and cli

* Fix replace cli

* Fix cookiecutter name

* Move docstring earlier to avoid not defined error

* Fix a missing Module

* Add encoder-decoder flax template from bart

* Fix flax test

* Make style

* Fix endif

* Fix replace all "utf-8 -> unp-8"

* Update comment

* Fix flax template (add missing ..._DOCSTRING)

* Use flax_bart imports in template (was t5)

* Fix unp

* Update templates/adding_a_new_model/tests

* Revert "Fix unp"

This reverts commit dc9002a41d902c4f9b07343eab1cb350c8b7fd57.

* Remove one line of copied from to suppress CI error

* Use generate_tensorflow_pytorch_and_flax

* Add a missing part

* fix typo

* fix flax config

* add examples for flax

* small rename

* correct modeling imports

* correct auto loading

* corrects some flax tests

* correct small typo

* correct as type

* finish modif

* correct more templates

* final fixes

* add file testers

* up

* make sure tests match template regex

* correct pytorch

* correct tf

* correct more tf

* correct imports

* minor error

* minor error

* correct init

* more fixes

* correct more flax tests

* correct flax test

* more fixes

* correct docs

* update

* fix

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

templates/adding_a_new_model/tests/encoder-bert-tokenizer.json

@@ -6,6 +6,6 @@
   "authors": "The HuggingFace Team",
   "checkpoint_identifier": "brand-new-bert-base-cased",
   "tokenizer_type": "Based on BERT",
-  "generate_tensorflow_and_pytorch": "PyTorch & TensorFlow",
+  "generate_tensorflow_pytorch_and_flax": "PyTorch, TensorFlow and Flax",
   "is_encoder_decoder_model": "False"
 }

templates/adding_a_new_model/tests/flax-encoder-bert-tokenizer.json

@@ -0,0 +1,11 @@
+{
+  "modelname": "TemplateFLAX",
+  "uppercase_modelname": "TEMPLATE_FLAX",
+  "lowercase_modelname": "template_flax",
+  "camelcase_modelname": "TemplateFlax",
+  "authors": "The HuggingFace Team",
+  "checkpoint_identifier": "brand-new-bert-base-cased",
+  "tokenizer_type": "Based on BERT",
+  "generate_tensorflow_pytorch_and_flax": "Flax",
+  "is_encoder_decoder_model": "False"
+}

templates/adding_a_new_model/tests/flax-seq-2-seq-bart-tokenizer.json

@@ -0,0 +1,11 @@
+{
+  "modelname": "FlaxNewENCDEC",
+  "uppercase_modelname": "FLAX_NEW_ENC_DEC",
+  "lowercase_modelname": "flax_new_enc_dec_template",
+  "camelcase_modelname": "FlaxNewEncDec",
+  "authors": "The HuggingFace Team",
+  "checkpoint_identifier": "new-flax-enc-dec-base",
+  "tokenizer_type": "Based on BART",
+  "generate_tensorflow_pytorch_and_flax": "Flax",
+  "is_encoder_decoder_model": "True"
+}

templates/adding_a_new_model/tests/pt-encoder-bert-tokenizer.json

@@ -6,6 +6,6 @@
   "authors": "The HuggingFace Team",
   "checkpoint_identifier": "brand-new-bert-base-cased",
   "tokenizer_type": "Based on BERT",
-  "generate_tensorflow_and_pytorch": "PyTorch",
+  "generate_tensorflow_pytorch_and_flax": "PyTorch",
   "is_encoder_decoder_model": "False"
 }

templates/adding_a_new_model/tests/pt-seq-2-seq-bart-tokenizer.json

@@ -1,11 +1,11 @@
 {
-  "modelname": "NewENCDEC",
-  "uppercase_modelname": "NEW_ENC_DEC",
-  "lowercase_modelname": "new_enc_dec",
-  "camelcase_modelname": "NewEncDec",
+  "modelname": "PTNewENCDEC",
+  "uppercase_modelname": "PT_NEW_ENC_DEC",
+  "lowercase_modelname": "pt_new_enc_dec_template",
+  "camelcase_modelname": "PtNewEncDec",
   "authors": "The HuggingFace Team",
-  "checkpoint_identifier": "new-enc-dec-base",
+  "checkpoint_identifier": "pt-new-enc-dec-base",
   "tokenizer_type": "Based on BART",
-  "generate_tensorflow_and_pytorch": "PyTorch",
+  "generate_tensorflow_pytorch_and_flax": "PyTorch",
   "is_encoder_decoder_model": "True"
 }

templates/adding_a_new_model/tests/standalone.json

@@ -6,6 +6,6 @@
   "authors": "The HuggingFace Team",
   "checkpoint_identifier": "bi-brand-new-bert-base-cased",
   "tokenizer_type": "Standalone",
-  "generate_tensorflow_and_pytorch": "PyTorch & TensorFlow",
+  "generate_tensorflow_pytorch_and_flax": "PyTorch, TensorFlow and Flax",
   "is_encoder_decoder_model": "False"
 }

templates/adding_a_new_model/tests/tf-encoder-bert-tokenizer.json

@@ -6,6 +6,6 @@
   "authors": "The HuggingFace Team",
   "checkpoint_identifier": "brand-new-bert-base-cased",
   "tokenizer_type": "Based on BERT",
-  "generate_tensorflow_and_pytorch": "TensorFlow",
+  "generate_tensorflow_pytorch_and_flax": "TensorFlow",
   "is_encoder_decoder_model": "False"
 }

templates/adding_a_new_model/tests/tf-seq-2-seq-bart-tokenizer.json

@@ -1,11 +1,11 @@
 {
   "modelname": "NewTFENCDEC",
   "uppercase_modelname": "NEW_TF_ENC_DEC",
-  "lowercase_modelname": "new_tf_enc_dec",
+  "lowercase_modelname": "new_tf_enc_dec_template",
   "camelcase_modelname": "NewTFEncDec",
   "authors": "The HuggingFace Team",
-  "checkpoint_identifier": "new-tf-enc-dec-base",
+  "checkpoint_identifier": "new-tf-enc-dec-base_template",
   "tokenizer_type": "Based on BART",
-  "generate_tensorflow_and_pytorch": "TensorFlow",
+  "generate_tensorflow_pytorch_and_flax": "TensorFlow",
   "is_encoder_decoder_model": "True"
 }