




"""simple docstring"""







import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
a      =							logging.get_logger(__name__)

a      =							{
    'post_extract_proj': 'feature_projection.projection',
    'encoder.pos_conv.0': 'encoder.pos_conv_embed.conv',
    'self_attn.k_proj': 'encoder.layers.*.attention.k_proj',
    'self_attn.v_proj': 'encoder.layers.*.attention.v_proj',
    'self_attn.q_proj': 'encoder.layers.*.attention.q_proj',
    'self_attn.out_proj': 'encoder.layers.*.attention.out_proj',
    'self_attn_layer_norm': 'encoder.layers.*.layer_norm',
    'fc1': 'encoder.layers.*.feed_forward.intermediate_dense',
    'fc2': 'encoder.layers.*.feed_forward.output_dense',
    'final_layer_norm': 'encoder.layers.*.final_layer_norm',
    'encoder.layer_norm': 'encoder.layer_norm',
    'w2v_model.layer_norm': 'feature_projection.layer_norm',
    'quantizer.weight_proj': 'quantizer.weight_proj',
    'quantizer.vars': 'quantizer.codevectors',
    'project_q': 'project_q',
    'final_proj': 'project_hid',
    'w2v_encoder.proj': 'lm_head',
    'mask_emb': 'masked_spec_embed',
}
a      =							[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]

def 							lowercase					(snake_case__   : str      ,       snake_case__   : Dict      ,       snake_case__   : List[str]      ,       snake_case__   : Dict      ,       snake_case__   : List[Any] )     ->  Dict:


				'''simple docstring'''






				for attribute in key.split(""".""" ):
								lowerCAmelCase  =				getattr(snake_case__      ,       snake_case__ )

				if weight_type is not None:
								lowerCAmelCase  =				getattr(snake_case__      ,       snake_case__ ).shape
				else:
								lowerCAmelCase  =				hf_pointer.shape

				assert hf_shape == value.shape, (
				    f'''Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be'''
				    f''' {value.shape} for {full_name}'''
				)

				if weight_type == "weight":
								lowerCAmelCase  =				value
				elif weight_type == "weight_g":
								lowerCAmelCase  =				value
				elif weight_type == "weight_v":
								lowerCAmelCase  =				value
				elif weight_type == "bias":
								lowerCAmelCase  =				value
				else:
								lowerCAmelCase  =				value

				logger.info(f'''{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.''' )

def 							lowercase					(snake_case__   : Tuple      ,       snake_case__   : Union[str, Any] )     ->  List[Any]:


				'''simple docstring'''






				lowerCAmelCase  =				[]
				lowerCAmelCase  =				fairseq_model.state_dict()

				lowerCAmelCase  =				hf_model.feature_extractor
				lowerCAmelCase  =				hf_model.adapter

				for name, value in fairseq_dict.items():
								lowerCAmelCase  =				False
								if "conv_layers" in name:
												load_conv_layer(
												    snake_case__      ,       snake_case__      ,       snake_case__      ,       snake_case__      ,       hf_model.config.feat_extract_norm == """group"""      ,       )
												lowerCAmelCase  =				True
								elif any(x in name for x in ["""adaptor""", """w2v_encoder.proj.""", """w2v_proj_ln."""] ):
												load_adapter(snake_case__      ,       snake_case__      ,       snake_case__      ,       snake_case__ )
												lowerCAmelCase  =				True
								else:
												for key, mapped_key in MAPPING.items():
																if key in name or key.split("""w2v_model.""" )[-1] == name.split(""".""" )[0]:
																				lowerCAmelCase  =				True
																				if "*" in mapped_key:
																								lowerCAmelCase  =				name.split(snake_case__ )[0].split(""".""" )[-2]
																								lowerCAmelCase  =				mapped_key.replace("""*"""      ,       snake_case__ )
																				if "weight_g" in name:
																								lowerCAmelCase  =				"""weight_g"""
																				elif "weight_v" in name:
																								lowerCAmelCase  =				"""weight_v"""
																				elif "bias" in name:
																								lowerCAmelCase  =				"""bias"""
																				elif "weight" in name:
																								lowerCAmelCase  =				"""weight"""
																				else:
																								lowerCAmelCase  =				None
																				set_recursively(snake_case__      ,       snake_case__      ,       snake_case__      ,       snake_case__      ,       snake_case__ )
																continue
								if not is_used:
												unused_weights.append(snake_case__ )

				logger.warning(f'''Unused weights: {unused_weights}''' )

def 							lowercase					(snake_case__   : str      ,       snake_case__   : Optional[int]      ,       snake_case__   : Any      ,       snake_case__   : Dict      ,       snake_case__   : Dict )     ->  str:


				'''simple docstring'''






				lowerCAmelCase  =				full_name.split("""conv_layers.""" )[-1]
				lowerCAmelCase  =				name.split(""".""" )
				lowerCAmelCase  =				int(items[0] )
				lowerCAmelCase  =				int(items[1] )

				if type_id == 0:
								if "bias" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
												    f'''{full_name} has size {value.shape}, but'''
												    f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
												)
												lowerCAmelCase  =				value
												logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.''' )
								elif "weight" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
												    f'''{full_name} has size {value.shape}, but'''
												    f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
												)
												lowerCAmelCase  =				value
												logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.''' )
				elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
								if "bias" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
												    f'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
												    " found."
												)
												lowerCAmelCase  =				value
												logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
								elif "weight" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
												    f'''{full_name} has size {value.shape}, but'''
												    f''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
												)
												lowerCAmelCase  =				value
												logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
				else:
								unused_weights.append(snake_case__ )

def 							lowercase					(snake_case__   : int      ,       snake_case__   : Optional[int]      ,       snake_case__   : Optional[Any]      ,       snake_case__   : int )     ->  Union[str, Any]:


				'''simple docstring'''






				lowerCAmelCase  =				full_name.split("""adaptor.""" )[-1]
				lowerCAmelCase  =				name.split(""".""" )

				if items[1].isdigit():
								lowerCAmelCase  =				int(items[1] )
				else:
								lowerCAmelCase  =				None

				if "adaptor" not in full_name:
								if "proj_ln" in full_name:
												# has to be layer norm
												if "bias" in name:
																assert (
																    value.shape == adapter.proj_layer_norm.bias.data.shape
																), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
																lowerCAmelCase  =				value
																logger.info(f'''Adapter proj layer norm bias was initialized from {full_name}.''' )
												if "weight" in name:
																assert (
																    value.shape == adapter.proj_layer_norm.weight.data.shape
																), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
																lowerCAmelCase  =				value
								else:
												# has to be projection layer
												if "bias" in name:
																assert (
																    value.shape == adapter.proj.bias.data.shape
																), f'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
																lowerCAmelCase  =				value
																logger.info(f'''Adapter proj layer bias was initialized from {full_name}.''' )
												if "weight" in name:
																assert (
																    value.shape == adapter.proj.weight.data.shape
																), f'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
																lowerCAmelCase  =				value
																logger.info(f'''Adapter proj layer weight was initialized from {full_name}.''' )
				elif isinstance(snake_case__      ,       snake_case__ ):
								if "bias" in name:
												assert (
												    value.shape == adapter.layers[layer_id].conv.bias.data.shape
												), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
												lowerCAmelCase  =				value
												logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
								elif "weight" in name:
												assert (
												    value.shape == adapter.layers[layer_id].conv.weight.data.shape
												), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
												lowerCAmelCase  =				value
												logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
				else:
								unused_weights.append(snake_case__ )

def 							lowercase					(snake_case__   : Dict )     ->  List[str]:


				'''simple docstring'''






				lowerCAmelCase ,						lowerCAmelCase  =				emb.weight.shape
				lowerCAmelCase  =				nn.Linear(snake_case__      ,       snake_case__      ,       bias=snake_case__ )
				lowerCAmelCase  =				emb.weight.data
				return lin_layer

@torch.no_grad()
def 							lowercase					(snake_case__   : Dict      ,       snake_case__   : Optional[int]      ,       snake_case__   : Optional[int]      ,       snake_case__   : Tuple      ,       snake_case__   : List[str]      ,       snake_case__   : Union[str, Any]      ,       snake_case__   : Tuple      ,       snake_case__   : Tuple      ,       snake_case__   : str      ,       snake_case__   : Union[str, Any]      ,       snake_case__   : Optional[int]      ,       )     ->  Optional[int]:


				'''simple docstring'''






				lowerCAmelCase  =				WavaVecaConfig.from_pretrained(
				    snake_case__      ,       add_adapter=snake_case__      ,       adapter_stride=snake_case__      ,       adapter_kernel_size=snake_case__      ,       use_auth_token=snake_case__      ,       output_hidden_size=snake_case__      ,       )
				lowerCAmelCase  =				MBartConfig.from_pretrained(snake_case__ )

				# load model
				lowerCAmelCase ,						lowerCAmelCase ,						lowerCAmelCase  =				fairseq.checkpoint_utils.load_model_ensemble_and_task(
				    [checkpoint_path]      ,       arg_overrides={
				        """config_yaml""": config_yaml_path,
				        """data""": """/""".join(dict_path.split("""/""" )[:-1] ),
				        """w2v_path""": checkpoint_path,
				        """load_pretrained_decoder_from""": None,
				    }      ,       )
				lowerCAmelCase  =				model[0].eval()

				# load feature extractor
				lowerCAmelCase  =				WavaVecaFeatureExtractor.from_pretrained(snake_case__      ,       use_auth_token=snake_case__ )

				# set weights for wav2vec2 encoder
				lowerCAmelCase  =				WavaVecaModel(snake_case__ )

				recursively_load_weights_wavaveca(model.encoder      ,       snake_case__ )

				# load decoder weights
				lowerCAmelCase  =				MBartForCausalLM(snake_case__ )
				lowerCAmelCase ,						lowerCAmelCase  =				hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()      ,       strict=snake_case__ )
				logger.warning(f'''The following keys are missing when loading the decoder weights: {missing_keys}''' )
				logger.warning(f'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}''' )

				lowerCAmelCase  =				SpeechEncoderDecoderModel(encoder=snake_case__      ,       decoder=snake_case__ )
				lowerCAmelCase  =				False

				lowerCAmelCase  =				MBartaaTokenizer(snake_case__ )
				tokenizer.save_pretrained(snake_case__ )

				lowerCAmelCase  =				hf_wavavec.config.to_dict()
				lowerCAmelCase  =				tokenizer.pad_token_id
				lowerCAmelCase  =				tokenizer.bos_token_id
				lowerCAmelCase  =				tokenizer.eos_token_id
				lowerCAmelCase  =				"""mbart50"""
				lowerCAmelCase  =				"""wav2vec2"""

				lowerCAmelCase  =				tokenizer.eos_token_id
				lowerCAmelCase  =				250_004
				lowerCAmelCase  =				tokenizer.eos_token_id

				lowerCAmelCase  =				SpeechEncoderDecoderConfig.from_dict(snake_case__ )

				hf_wavavec.save_pretrained(snake_case__ )
				feature_extractor.save_pretrained(snake_case__ )


if __name__ == "__main__":
		a      =							argparse.ArgumentParser()
		parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
		parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
		parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
		parser.add_argument('--config_yaml_path', default=None, type=str, help='Path to yaml file of fine-tuned model')
		parser.add_argument(
		    '--encoder_config_path',
		    default='facebook/wav2vec2-xls-r-1b',
		    type=str,
		    help='Path to hf encoder wav2vec2 checkpoint config',
		)
		parser.add_argument(
		    '--decoder_config_path',
		    default='facebook/mbart-large-50-one-to-many-mmt',
		    type=str,
		    help='Path to hf decoder checkpoint config',
		)
		parser.add_argument('--add_adapter', default=True, type=bool, help='whethere to add model adapter layers')
		parser.add_argument('--adapter_stride', default=2, type=int, help='stride of adapter layers')
		parser.add_argument('--adapter_kernel_size', default=3, type=int, help='kernel size of adapter layers')
		parser.add_argument('--encoder_output_dim', default=1_0_2_4, type=int, help='encoder output dim')
		parser.add_argument('--start_token_id', default=2_5_0_0_0_4, type=int, help='`decoder_start_token_id` of model config')

		a      =							parser.parse_args()
		convert_wavaveca_checkpoint(
		    args.checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.dict_path,
		    args.config_yaml_path,
		    encoder_config_path=args.encoder_config_path,
		    decoder_config_path=args.decoder_config_path,
		    add_adapter=args.add_adapter,
		    adapter_kernel_size=args.adapter_kernel_size,
		    adapter_stride=args.adapter_stride,
		    decoder_start_token_id=args.start_token_id,
		    encoder_output_dim=args.encoder_output_dim,
		)











"""simple docstring"""







import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
a      =							logging.get_logger(__name__)





def 							lowercase					(snake_case__   : str      ,       snake_case__   : Tuple      ,       snake_case__   : Dict )     ->  int:


				'''simple docstring'''






				lowerCAmelCase  =				os.path.abspath(snake_case__ )
				logger.info(f'''Converting TensorFlow checkpoint from {tf_path}''' )
				# Load weights from TF model
				lowerCAmelCase  =				tf.train.list_variables(snake_case__ )
				lowerCAmelCase  =				[]
				lowerCAmelCase  =				[]
				lowerCAmelCase  =				[]
				for full_name, shape in init_vars:
								# logger.info(f"Loading TF weight {name} with shape {shape}")
								lowerCAmelCase  =				full_name.split("""/""" )
								if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
												logger.info(f'''Skipping non-model layer {full_name}''' )
												continue
								if "optimizer" in full_name:
												logger.info(f'''Skipping optimization layer {full_name}''' )
												continue
								if name[0] == "model":
												# ignore initial 'model'
												lowerCAmelCase  =				name[1:]
								# figure out how many levels deep the name is
								lowerCAmelCase  =				0
								for _name in name:
												if _name.startswith("""layer_with_weights""" ):
																depth += 1
												else:
																break
								layer_depth.append(snake_case__ )
								# read data
								lowerCAmelCase  =				tf.train.load_variable(snake_case__      ,       snake_case__ )
								names.append("""/""".join(snake_case__ ) )
								arrays.append(snake_case__ )
				logger.info(f'''Read a total of {len(snake_case__ ):,} layers''' )

				# Sanity check
				if len(set(snake_case__ ) ) != 1:
								raise ValueError(f'''Found layer names with different depths (layer depth {list(set(snake_case__ ) )})''' )
				lowerCAmelCase  =				list(set(snake_case__ ) )[0]
				if layer_depth != 1:
								raise ValueError(
								    """The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"""
								    """ heads.""" )

				# convert layers
				logger.info("""Converting weights...""" )
				for full_name, array in zip(snake_case__      ,       snake_case__ ):
								lowerCAmelCase  =				full_name.split("""/""" )
								lowerCAmelCase  =				model
								lowerCAmelCase  =				[]
								for i, m_name in enumerate(snake_case__ ):
												if m_name == ".ATTRIBUTES":
																# variable names end with .ATTRIBUTES/VARIABLE_VALUE
																break
												if m_name.startswith("""layer_with_weights""" ):
																lowerCAmelCase  =				int(m_name.split("""-""" )[-1] )
																if layer_num <= 2:
																				# embedding layers
																				# layer_num 0: word_embeddings
																				# layer_num 1: position_embeddings
																				# layer_num 2: token_type_embeddings
																				continue
																elif layer_num == 3:
																				# embedding LayerNorm
																				trace.extend(["""embeddings""", """LayerNorm"""] )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """embeddings""" )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """LayerNorm""" )
																elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
																				# encoder layers
																				trace.extend(["""encoder""", """layer""", str(layer_num - 4 )] )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """encoder""" )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """layer""" )
																				lowerCAmelCase  =				pointer[layer_num - 4]
																elif layer_num == config.num_hidden_layers + 4:
																				# pooler layer
																				trace.extend(["""pooler""", """dense"""] )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """pooler""" )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """dense""" )
												elif m_name == "embeddings":
																trace.append("""embeddings""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """embeddings""" )
																if layer_num == 0:
																				trace.append("""word_embeddings""" )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """word_embeddings""" )
																elif layer_num == 1:
																				trace.append("""position_embeddings""" )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """position_embeddings""" )
																elif layer_num == 2:
																				trace.append("""token_type_embeddings""" )
																				lowerCAmelCase  =				getattr(snake_case__      ,       """token_type_embeddings""" )
																else:
																				raise ValueError(f'''Unknown embedding layer with name {full_name}''' )
																trace.append("""weight""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """weight""" )
												elif m_name == "_attention_layer":
																# self-attention layer
																trace.extend(["""attention""", """self"""] )
																lowerCAmelCase  =				getattr(snake_case__      ,       """attention""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """self""" )
												elif m_name == "_attention_layer_norm":
																# output attention norm
																trace.extend(["""attention""", """output""", """LayerNorm"""] )
																lowerCAmelCase  =				getattr(snake_case__      ,       """attention""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """output""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """LayerNorm""" )
												elif m_name == "_attention_output_dense":
																# output attention dense
																trace.extend(["""attention""", """output""", """dense"""] )
																lowerCAmelCase  =				getattr(snake_case__      ,       """attention""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """output""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """dense""" )
												elif m_name == "_output_dense":
																# output dense
																trace.extend(["""output""", """dense"""] )
																lowerCAmelCase  =				getattr(snake_case__      ,       """output""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """dense""" )
												elif m_name == "_output_layer_norm":
																# output dense
																trace.extend(["""output""", """LayerNorm"""] )
																lowerCAmelCase  =				getattr(snake_case__      ,       """output""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """LayerNorm""" )
												elif m_name == "_key_dense":
																# attention key
																trace.append("""key""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """key""" )
												elif m_name == "_query_dense":
																# attention query
																trace.append("""query""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """query""" )
												elif m_name == "_value_dense":
																# attention value
																trace.append("""value""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """value""" )
												elif m_name == "_intermediate_dense":
																# attention intermediate dense
																trace.extend(["""intermediate""", """dense"""] )
																lowerCAmelCase  =				getattr(snake_case__      ,       """intermediate""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """dense""" )
												elif m_name == "_output_layer_norm":
																# output layer norm
																trace.append("""output""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """output""" )
												# weights & biases
												elif m_name in ["bias", "beta"]:
																trace.append("""bias""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """bias""" )
												elif m_name in ["kernel", "gamma"]:
																trace.append("""weight""" )
																lowerCAmelCase  =				getattr(snake_case__      ,       """weight""" )
												else:
																logger.warning(f'''Ignored {m_name}''' )
        # for certain layers reshape is necessary
								lowerCAmelCase  =				""".""".join(snake_case__ )
								if re.match(R"""(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)"""      ,       snake_case__ ) or re.match(
								    R"""(\S+)\.attention\.output\.dense\.weight"""      ,       snake_case__ ):
												lowerCAmelCase  =				array.reshape(pointer.data.shape )
								if "kernel" in full_name:
												lowerCAmelCase  =				array.transpose()
								if pointer.shape == array.shape:
												lowerCAmelCase  =				torch.from_numpy(snake_case__ )
								else:
												raise ValueError(
												    f'''Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'''
												    f''' {array.shape}''' )
								logger.info(f'''Successfully set variable {full_name} to PyTorch layer {trace}''' )
				return model





def 							lowercase					(snake_case__   : Tuple      ,       snake_case__   : Any      ,       snake_case__   : Any )     ->  str:


				'''simple docstring'''






				logger.info(f'''Loading model based on config from {config_path}...''' )
				lowerCAmelCase  =				BertConfig.from_json_file(snake_case__ )
				lowerCAmelCase  =				BertModel(snake_case__ )

				# Load weights from checkpoint
				logger.info(f'''Loading weights from checkpoint {tf_checkpoint_path}...''' )
				load_tfa_weights_in_bert(snake_case__      ,       snake_case__      ,       snake_case__ )

				# Save pytorch-model
				logger.info(f'''Saving PyTorch model to {pytorch_dump_path}...''' )
				torch.save(model.state_dict()      ,       snake_case__ )


if __name__ == "__main__":
		a      =							argparse.ArgumentParser()
		parser.add_argument(
		    '--tf_checkpoint_path', type=str, required=True, help='Path to the TensorFlow 2.x 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 (must include filename).',
		)
		a      =							parser.parse_args()
		convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)





