


'''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()
__snake_case							      =     logging.get_logger(__name__)


def 			a      (			__a ,					__a ,					__a   )  ->	Tuple:



					'''simple docstring'''






					UpperCamelCase__   :Union[str, Any]												= os.path.abspath(lowerCamelCase_   )
					logger.info(f'''Converting TensorFlow checkpoint from {tf_path}'''   )
					# Load weights from TF model
					UpperCamelCase__   :Optional[Any]												= tf.train.list_variables(lowerCamelCase_   )
					UpperCamelCase__   :Optional[Any]												= []
					UpperCamelCase__   :List[str]												= []
					UpperCamelCase__   :str												= []
					for full_name, shape in init_vars:
										# logger.info(f"Loading TF weight {name} with shape {shape}")
										UpperCamelCase__   :List[Any]												= 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'
															UpperCamelCase__   :Union[str, Any]												= name[1:]
										# figure out how many levels deep the name is
										UpperCamelCase__   :str												= 0
										for _name in name:
															if _name.startswith('''layer_with_weights'''   ):
																				depth += 1
															else:
																				break
										layer_depth.append(lowerCamelCase_   )
										# read data
										UpperCamelCase__   :Union[str, Any]												= tf.train.load_variable(lowerCamelCase_ ,					lowerCamelCase_   )
										names.append('''/'''.join(lowerCamelCase_   )   )
										arrays.append(lowerCamelCase_   )
					logger.info(f'''Read a total of {len(lowerCamelCase_   ):,} layers'''   )

					# Sanity check
					if len(set(lowerCamelCase_   )   ) != 1:
										raise ValueError(f'''Found layer names with different depths (layer depth {list(set(lowerCamelCase_   )   )})'''   )
					UpperCamelCase__   :Any												= list(set(lowerCamelCase_   )   )[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(lowerCamelCase_ ,					lowerCamelCase_   ):
										UpperCamelCase__   :Tuple												= full_name.split('''/'''   )
										UpperCamelCase__   :Optional[int]												= model
										UpperCamelCase__   :str												= []
										for i, m_name in enumerate(lowerCamelCase_   ):
															if m_name == ".ATTRIBUTES":
																				# variable names end with .ATTRIBUTES/VARIABLE_VALUE
																				break
															if m_name.startswith('''layer_with_weights'''   ):
																				UpperCamelCase__   :Optional[Any]												= 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''']   )
																									UpperCamelCase__   :Union[str, Any]												= getattr(lowerCamelCase_ ,					'''embeddings'''   )
																									UpperCamelCase__   :Optional[int]												= getattr(lowerCamelCase_ ,					'''LayerNorm'''   )
																				elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
																									# encoder layers
																									trace.extend(['''encoder''', '''layer''', str(layer_num - 4   )]   )
																									UpperCamelCase__   :List[str]												= getattr(lowerCamelCase_ ,					'''encoder'''   )
																									UpperCamelCase__   :Optional[int]												= getattr(lowerCamelCase_ ,					'''layer'''   )
																									UpperCamelCase__   :Dict												= pointer[layer_num - 4]
																				elif layer_num == config.num_hidden_layers + 4:
																									# pooler layer
																									trace.extend(['''pooler''', '''dense''']   )
																									UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''pooler'''   )
																									UpperCamelCase__   :Tuple												= getattr(lowerCamelCase_ ,					'''dense'''   )
															elif m_name == "embeddings":
																				trace.append('''embeddings'''   )
																				UpperCamelCase__   :Dict												= getattr(lowerCamelCase_ ,					'''embeddings'''   )
																				if layer_num == 0:
																									trace.append('''word_embeddings'''   )
																									UpperCamelCase__   :str												= getattr(lowerCamelCase_ ,					'''word_embeddings'''   )
																				elif layer_num == 1:
																									trace.append('''position_embeddings'''   )
																									UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''position_embeddings'''   )
																				elif layer_num == 2:
																									trace.append('''token_type_embeddings'''   )
																									UpperCamelCase__   :List[Any]												= getattr(lowerCamelCase_ ,					'''token_type_embeddings'''   )
																				else:
																									raise ValueError(f'''Unknown embedding layer with name {full_name}'''   )
																				trace.append('''weight'''   )
																				UpperCamelCase__   :List[Any]												= getattr(lowerCamelCase_ ,					'''weight'''   )
															elif m_name == "_attention_layer":
																				# self-attention layer
																				trace.extend(['''attention''', '''self''']   )
																				UpperCamelCase__   :Tuple												= getattr(lowerCamelCase_ ,					'''attention'''   )
																				UpperCamelCase__   :Tuple												= getattr(lowerCamelCase_ ,					'''self'''   )
															elif m_name == "_attention_layer_norm":
																				# output attention norm
																				trace.extend(['''attention''', '''output''', '''LayerNorm''']   )
																				UpperCamelCase__   :Dict												= getattr(lowerCamelCase_ ,					'''attention'''   )
																				UpperCamelCase__   :int												= getattr(lowerCamelCase_ ,					'''output'''   )
																				UpperCamelCase__   :Dict												= getattr(lowerCamelCase_ ,					'''LayerNorm'''   )
															elif m_name == "_attention_output_dense":
																				# output attention dense
																				trace.extend(['''attention''', '''output''', '''dense''']   )
																				UpperCamelCase__   :Union[str, Any]												= getattr(lowerCamelCase_ ,					'''attention'''   )
																				UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''output'''   )
																				UpperCamelCase__   :Dict												= getattr(lowerCamelCase_ ,					'''dense'''   )
															elif m_name == "_output_dense":
																				# output dense
																				trace.extend(['''output''', '''dense''']   )
																				UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''output'''   )
																				UpperCamelCase__   :Any												= getattr(lowerCamelCase_ ,					'''dense'''   )
															elif m_name == "_output_layer_norm":
																				# output dense
																				trace.extend(['''output''', '''LayerNorm''']   )
																				UpperCamelCase__   :List[Any]												= getattr(lowerCamelCase_ ,					'''output'''   )
																				UpperCamelCase__   :Dict												= getattr(lowerCamelCase_ ,					'''LayerNorm'''   )
															elif m_name == "_key_dense":
																				# attention key
																				trace.append('''key'''   )
																				UpperCamelCase__   :List[Any]												= getattr(lowerCamelCase_ ,					'''key'''   )
															elif m_name == "_query_dense":
																				# attention query
																				trace.append('''query'''   )
																				UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''query'''   )
															elif m_name == "_value_dense":
																				# attention value
																				trace.append('''value'''   )
																				UpperCamelCase__   :Any												= getattr(lowerCamelCase_ ,					'''value'''   )
															elif m_name == "_intermediate_dense":
																				# attention intermediate dense
																				trace.extend(['''intermediate''', '''dense''']   )
																				UpperCamelCase__   :int												= getattr(lowerCamelCase_ ,					'''intermediate'''   )
																				UpperCamelCase__   :Optional[int]												= getattr(lowerCamelCase_ ,					'''dense'''   )
															elif m_name == "_output_layer_norm":
																				# output layer norm
																				trace.append('''output'''   )
																				UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''output'''   )
															# weights & biases
															elif m_name in ["bias", "beta"]:
																				trace.append('''bias'''   )
																				UpperCamelCase__   :Optional[Any]												= getattr(lowerCamelCase_ ,					'''bias'''   )
															elif m_name in ["kernel", "gamma"]:
																				trace.append('''weight'''   )
																				UpperCamelCase__   :Optional[int]												= getattr(lowerCamelCase_ ,					'''weight'''   )
															else:
																				logger.warning(f'''Ignored {m_name}'''   )
        # for certain layers reshape is necessary
										UpperCamelCase__   :Any												= '.'.join(lowerCamelCase_   )
										if re.match(R'''(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)''' ,					lowerCamelCase_   ) or re.match(
										    R'''(\S+)\.attention\.output\.dense\.weight''' ,					lowerCamelCase_   ):
															UpperCamelCase__   :Any												= array.reshape(pointer.data.shape   )
										if "kernel" in full_name:
															UpperCamelCase__   :Union[str, Any]												= array.transpose()
										if pointer.shape == array.shape:
															UpperCamelCase__   :Optional[int]												= torch.from_numpy(lowerCamelCase_   )
										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 			a      (			__a ,					__a ,					__a   )  ->	Optional[int]:



					'''simple docstring'''






					logger.info(f'''Loading model based on config from {config_path}...'''   )
					UpperCamelCase__   :Union[str, Any]												= BertConfig.from_json_file(lowerCamelCase_   )
					UpperCamelCase__   :Union[str, Any]												= BertModel(lowerCamelCase_   )

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

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


if __name__ == "__main__":
		__snake_case							      =     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).''',
		)
		__snake_case							      =     parser.parse_args()
		convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_tf_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE       :      List[str]												= {
    "configuration_speech_to_text": ["SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP", "Speech2TextConfig"],
    "processing_speech_to_text": ["Speech2TextProcessor"],
}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE       :      Union[str, Any]												= ["Speech2TextTokenizer"]

try:
 if not is_speech_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE       :      int												= ["Speech2TextFeatureExtractor"]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE       :      Optional[Any]												= [
     "TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST",
     "TFSpeech2TextForConditionalGeneration",
     "TFSpeech2TextModel",
     "TFSpeech2TextPreTrainedModel",
 ]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE       :      str												= [
     "SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST",
     "Speech2TextForConditionalGeneration",
     "Speech2TextModel",
     "Speech2TextPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
 from .processing_speech_to_text import SpeechaTextProcessor

 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_speech_to_text import SpeechaTextTokenizer

 try:
  if not is_speech_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_speech_to_text import (
      TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFSpeechaTextForConditionalGeneration,
      TFSpeechaTextModel,
      TFSpeechaTextPreTrainedModel,
  )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_speech_to_text import (
      SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
      SpeechaTextForConditionalGeneration,
      SpeechaTextModel,
      SpeechaTextPreTrainedModel,
  )

else:
 import sys

 SCREAMING_SNAKE_CASE       :      Optional[Any]												= _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


