





import argparse
import json
import logging
import os
import sys
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


_lowerCamelCase   				=			[
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        "text-classification",
        "language-modeling",
        "summarization",
        "token-classification",
        "question-answering",
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
	import run_clm_flax
	import run_flax_glue
	import run_flax_ner
	import run_mlm_flax
	import run_qa
	import run_summarization_flax
	import run_ta_mlm_flax


logging.basicConfig(level=logging.DEBUG)

_lowerCamelCase   				=			logging.getLogger()






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      argparse.ArgumentParser()
		parser.add_argument("""-f"""    )
		SCREAMING_SNAKE_CASE__			=      parser.parse_args()
		return args.f






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[Any]					,   UpperCamelCase__:							str="eval"    ):
		SCREAMING_SNAKE_CASE__			=      os.path.join(_SCREAMING_SNAKE_CASE					,   f'''{split}_results.json'''    )
		if os.path.exists(_SCREAMING_SNAKE_CASE    ):
				with open(_SCREAMING_SNAKE_CASE					,   """r"""    ) as f:
						return json.load(_SCREAMING_SNAKE_CASE    )
		raise ValueError(f'''can\'t find {path}'''    )


_lowerCamelCase   				=			logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)




class   UpperCamelCase_       (     UpperCamelCase__    ):


							def      _snake_case  (							self    :List[Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_glue.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --eval_steps=2\n            --warmup_steps=2\n            --seed=42\n            --max_seq_length=128\n            '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_flax_glue.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase		)
											self.assertGreaterEqual(result["""eval_accuracy"""]							,				0.7_5		)


							@slow
							def      _snake_case  (							self    :Union[str, Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_clm_flax.py\n            --model_name_or_path distilgpt2\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --do_train\n            --do_eval\n            --block_size 128\n            --per_device_train_batch_size 4\n            --per_device_eval_batch_size 4\n            --num_train_epochs 2\n            --logging_steps 2 --eval_steps 2\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_clm_flax.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase		)
											self.assertLess(result["""eval_perplexity"""]							,				100		)


							@slow
							def      _snake_case  (							self    :Optional[int]		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_summarization.py\n            --model_name_or_path t5-small\n            --train_file tests/fixtures/tests_samples/xsum/sample.json\n            --validation_file tests/fixtures/tests_samples/xsum/sample.json\n            --test_file tests/fixtures/tests_samples/xsum/sample.json\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --num_train_epochs=3\n            --warmup_steps=8\n            --do_train\n            --do_eval\n            --do_predict\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --predict_with_generate\n        '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_summarization_flax.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase							,				split="""test"""		)
											self.assertGreaterEqual(result["""test_rouge1"""]							,				10		)
											self.assertGreaterEqual(result["""test_rouge2"""]							,				2		)
											self.assertGreaterEqual(result["""test_rougeL"""]							,				7		)
											self.assertGreaterEqual(result["""test_rougeLsum"""]							,				7		)


							@slow
							def      _snake_case  (							self    :List[str]		)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_mlm.py\n            --model_name_or_path distilroberta-base\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --max_seq_length 128\n            --per_device_train_batch_size 4\n            --per_device_eval_batch_size 4\n            --logging_steps 2 --eval_steps 2\n            --do_train\n            --do_eval\n            --num_train_epochs=1\n        '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_mlm_flax.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase		)
											self.assertLess(result["""eval_perplexity"""]							,				42		)


							@slow
							def      _snake_case  (							self    :Any		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_t5_mlm_flax.py\n            --model_name_or_path t5-small\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --do_train\n            --do_eval\n            --max_seq_length 128\n            --per_device_train_batch_size 4\n            --per_device_eval_batch_size 4\n            --num_train_epochs 2\n            --logging_steps 2 --eval_steps 2\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_ta_mlm_flax.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase		)
											self.assertGreaterEqual(result["""eval_accuracy"""]							,				0.4_2		)


							@slow
							def      _snake_case  (							self    :Union[str, Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      7 if get_gpu_count() > 1 else 2

									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_flax_ner.py\n            --model_name_or_path bert-base-uncased\n            --train_file tests/fixtures/tests_samples/conll/sample.json\n            --validation_file tests/fixtures/tests_samples/conll/sample.json\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --do_train\n            --do_eval\n            --warmup_steps=2\n            --learning_rate=2e-4\n            --logging_steps 2 --eval_steps 2\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=2\n            --num_train_epochs={epochs}\n            --seed 7\n        '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_flax_ner.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase		)
											self.assertGreaterEqual(result["""eval_accuracy"""]							,				0.7_5		)
											self.assertGreaterEqual(result["""eval_f1"""]							,				0.3		)


							@slow
							def      _snake_case  (							self    :Union[str, Any]		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__			=      f'''\n            run_qa.py\n            --model_name_or_path bert-base-uncased\n            --version_2_with_negative\n            --train_file tests/fixtures/tests_samples/SQUAD/sample.json\n            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --num_train_epochs=3\n            --warmup_steps=2\n            --do_train\n            --do_eval\n            --logging_steps 2 --eval_steps 2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n        '''.split()

									with patch.object(__UpperCamelCase							,				"""argv"""							,				__UpperCamelCase		):
											run_qa.main()
											SCREAMING_SNAKE_CASE__			=      get_results(__UpperCamelCase		)
											self.assertGreaterEqual(result["""eval_f1"""]							,				30		)
											self.assertGreaterEqual(result["""eval_exact"""]							,				30		)













import argparse

import tensorflow as tf
import torch

from transformers import BertConfig, BertForMaskedLM
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertPooler,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


logging.set_verbosity_info()






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							str					,   UpperCamelCase__:							str					,   UpperCamelCase__:							str    ):
		def get_masked_lm_array(UpperCamelCase__:							str    ):
				SCREAMING_SNAKE_CASE__			=      f'''masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		def get_encoder_array(UpperCamelCase__:							str    ):
				SCREAMING_SNAKE_CASE__			=      f'''encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		def get_encoder_layer_array(UpperCamelCase__:							int					,   UpperCamelCase__:							str    ):
				SCREAMING_SNAKE_CASE__			=      f'''encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		def get_encoder_attention_layer_array(UpperCamelCase__:							int					,   UpperCamelCase__:							str					,   UpperCamelCase__:							Any    ):
				SCREAMING_SNAKE_CASE__			=      f'''encoder/_transformer_layers/{layer_index}/_attention_layer/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
				SCREAMING_SNAKE_CASE__			=      tf.train.load_variable(UpperCamelCase__					,   UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      array.reshape(UpperCamelCase__    )

				if "kernel" in name:
						SCREAMING_SNAKE_CASE__			=      array.transpose()

				return torch.from_numpy(UpperCamelCase__    )

		print(f'''Loading model based on config from {config_path}...'''    )
		SCREAMING_SNAKE_CASE__			=      BertConfig.from_json_file(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      BertForMaskedLM(UpperCamelCase__    )

		# Layers
		for layer_index in range(0					,   config.num_hidden_layers    ):
				SCREAMING_SNAKE_CASE__			=      model.bert.encoder.layer[layer_index]

				# Self-attention
				SCREAMING_SNAKE_CASE__			=      layer.attention.self

				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_query_dense/kernel"""					,   self_attn.query.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_query_dense/bias"""					,   self_attn.query.bias.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_key_dense/kernel"""					,   self_attn.key.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_key_dense/bias"""					,   self_attn.key.bias.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_value_dense/kernel"""					,   self_attn.value.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_value_dense/bias"""					,   self_attn.value.bias.data.shape    )

				# Self-attention Output
				SCREAMING_SNAKE_CASE__			=      layer.attention.output

				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_output_dense/kernel"""					,   self_output.dense.weight.data.shape    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_attention_layer_array(
				    UpperCamelCase__					,   """_output_dense/bias"""					,   self_output.dense.bias.data.shape    )

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_attention_layer_norm/gamma"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_attention_layer_norm/beta"""    )

				# Intermediate
				SCREAMING_SNAKE_CASE__			=      layer.intermediate

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_intermediate_dense/kernel"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_intermediate_dense/bias"""    )

				# Output
				SCREAMING_SNAKE_CASE__			=      layer.output

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_dense/kernel"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_dense/bias"""    )

				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_layer_norm/gamma"""    )
				SCREAMING_SNAKE_CASE__			=      get_encoder_layer_array(UpperCamelCase__					,   """_output_layer_norm/beta"""    )

		# Embeddings
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_position_embedding_layer/embeddings"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_type_embedding_layer/embeddings"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_embedding_norm_layer/gamma"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_embedding_norm_layer/beta"""    )

		# LM Head
		SCREAMING_SNAKE_CASE__			=      model.cls.predictions.transform

		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""dense/kernel"""    )
		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""dense/bias"""    )

		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""layer_norm/gamma"""    )
		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""layer_norm/beta"""    )

		SCREAMING_SNAKE_CASE__			=      get_masked_lm_array("""embedding_table"""    )

		# Pooling
		SCREAMING_SNAKE_CASE__			=      BertPooler(config=UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_pooler_layer/kernel"""    )
		SCREAMING_SNAKE_CASE__			=      get_encoder_array("""_pooler_layer/bias"""    )

		# Export final model
		model.save_pretrained(UpperCamelCase__    )

		# Integration test - should load without any errors ;)
		SCREAMING_SNAKE_CASE__			=      BertForMaskedLM.from_pretrained(UpperCamelCase__    )
		print(new_model.eval()    )

		print("""Model conversion was done sucessfully!"""    )


if __name__ == "__main__":
	_lowerCamelCase   				=			argparse.ArgumentParser()
	parser.add_argument(
	    '--tf_checkpoint_path', type=str, required=True, help='Path to the TensorFlow Token Dropping checkpoint path.'
	)
	parser.add_argument(
	    '--bert_config_file',
	    type=str,
	    required=True,
	    help='The config json file corresponding to the BERT model. This specifies the model architecture.',
	)
	parser.add_argument(
	    '--pytorch_dump_path',
	    type=str,
	    required=True,
	    help='Path to the output PyTorch model.',
	)
	_lowerCamelCase   				=			parser.parse_args()
	convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)