





import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_lowerCamelCase   				=			get_tests_dir('fixtures')
_lowerCamelCase   				=			get_tests_dir('fixtures/dummy_feature_extractor_config.json')
_lowerCamelCase   				=			get_tests_dir('fixtures/dummy-config.json')




class   UpperCamelCase_       (     unittest.TestCase    ):


							def      _snake_case  (							self    :Optional[Any]		)   ->   Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      0


							def      _snake_case  (							self    :str		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained("""facebook/wav2vec2-base-960h"""		)
									self.assertIsInstance(__A							,				__A		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A		)
									self.assertIsInstance(__A							,				__A		)


							def      _snake_case  (							self    :Optional[int]		)   ->   List[Any]:



									"""simple docstring"""



									with tempfile.TemporaryDirectory() as tmpdirname:
											SCREAMING_SNAKE_CASE__			=      WavaVecaConfig()

											# remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A		).to_dict()

											config_dict.pop("""feature_extractor_type"""		)
											SCREAMING_SNAKE_CASE__			=      WavaVecaFeatureExtractor(**__A		)

											# save in new folder
											model_config.save_pretrained(__A		)
											config.save_pretrained(__A		)

											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A		)

											# make sure private variable is not incorrectly saved
											SCREAMING_SNAKE_CASE__			=      json.loads(config.to_json_string()		)
											self.assertTrue("""_processor_class""" not in dict_as_saved		)

									self.assertIsInstance(__A							,				__A		)


							def      _snake_case  (							self    :List[str]		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A		)
									self.assertIsInstance(__A							,				__A		)


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



									"""simple docstring"""



									with self.assertRaisesRegex(
									    __A							,				"""bert-base is not a local folder and is not a valid model identifier"""		):
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained("""bert-base"""		)


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



									"""simple docstring"""



									with self.assertRaisesRegex(
									    __A							,				r"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""		):
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A							,				revision="""aaaaaa"""		)


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



									"""simple docstring"""



									with self.assertRaisesRegex(
									    __A							,				"""hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."""							,				):
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained("""hf-internal-testing/config-no-model"""		)


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



									"""simple docstring"""



									with self.assertRaises(__A		):
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(
											    """hf-internal-testing/test_dynamic_feature_extractor"""		)
									# If remote code is disabled, we can't load this config.
									with self.assertRaises(__A		):
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(
											    """hf-internal-testing/test_dynamic_feature_extractor"""							,				trust_remote_code=__A		)

									SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(
									    """hf-internal-testing/test_dynamic_feature_extractor"""							,				trust_remote_code=__A		)
									self.assertEqual(feature_extractor.__class__.__name__							,				"""NewFeatureExtractor"""		)

									# Test feature extractor can be reloaded.
									with tempfile.TemporaryDirectory() as tmp_dir:
											feature_extractor.save_pretrained(__A		)
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A							,				trust_remote_code=__A		)
									self.assertEqual(reloaded_feature_extractor.__class__.__name__							,				"""NewFeatureExtractor"""		)


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



									"""simple docstring"""



									try:
											AutoConfig.register("""custom"""							,				__A		)
											AutoFeatureExtractor.register(__A							,				__A		)
											# Trying to register something existing in the Transformers library will raise an error
											with self.assertRaises(__A		):
													AutoFeatureExtractor.register(__A							,				__A		)

											# Now that the config is registered, it can be used as any other config with the auto-API
											SCREAMING_SNAKE_CASE__			=      CustomFeatureExtractor.from_pretrained(__A		)
											with tempfile.TemporaryDirectory() as tmp_dir:
													feature_extractor.save_pretrained(__A		)
													SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(__A		)
													self.assertIsInstance(__A							,				__A		)

									finally:
											if "custom" in CONFIG_MAPPING._extra_content:
													del CONFIG_MAPPING._extra_content["custom"]
											if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
													del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]


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



									"""simple docstring"""



									class   UpperCamelCase_       (     UpperCamelCase__    ):
																lowerCamelCase_           =					True

									try:
											AutoConfig.register("""custom"""							,				__A		)
											AutoFeatureExtractor.register(__A							,				__A		)
											# If remote code is not set, the default is to use local
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(
											    """hf-internal-testing/test_dynamic_feature_extractor"""		)
											self.assertEqual(feature_extractor.__class__.__name__							,				"""NewFeatureExtractor"""		)
											self.assertTrue(feature_extractor.is_local		)

											# If remote code is disabled, we load the local one.
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(
											    """hf-internal-testing/test_dynamic_feature_extractor"""							,				trust_remote_code=__A		)
											self.assertEqual(feature_extractor.__class__.__name__							,				"""NewFeatureExtractor"""		)
											self.assertTrue(feature_extractor.is_local		)

											# If remote is enabled, we load from the Hub
											SCREAMING_SNAKE_CASE__			=      AutoFeatureExtractor.from_pretrained(
											    """hf-internal-testing/test_dynamic_feature_extractor"""							,				trust_remote_code=__A		)
											self.assertEqual(feature_extractor.__class__.__name__							,				"""NewFeatureExtractor"""		)
											self.assertTrue(not hasattr(__A							,				"""is_local"""		)		)

									finally:
											if "custom" in CONFIG_MAPPING._extra_content:
													del CONFIG_MAPPING._extra_content["custom"]
											if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
													del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]






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)