




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()
__lowerCamelCase   :						List[str]				=			logging.get_logger(__name__)







def 		lowercase__   (		__A:	int							,__A:	Optional[Any]							,__A:	Optional[Any]				):





					'''simple docstring'''






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

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





					'''simple docstring'''






					logger.info(F'''Loading model based on config from {config_path}...'''				)
					__magic_name__							:	Optional[Any]    							=							BertConfig.from_json_file(__A				)
					__magic_name__							:	int    							=							BertModel(__A				)

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

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


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









import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
	from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
	from tensorflow.keras.optimizers import Adam


class    lowerCamelCase    (		tf.keras.optimizers.schedules.LearningRateSchedule						):






			'''simple docstring'''





			def __init__( self			:   Union[str, Any]       ,						lowerCamelCase_			:   float       ,						lowerCamelCase_			:   Callable       ,						lowerCamelCase_			:   int       ,						lowerCamelCase_			:   float = 1.0       ,						lowerCamelCase_			:   str = None       ,						)						->   Optional[int]:
								super().__init__()
								__magic_name__							:	Union[str, Any]    							=							initial_learning_rate
								__magic_name__							:	Optional[int]    							=							warmup_steps
								__magic_name__							:	Union[str, Any]    							=							power
								__magic_name__							:	List[Any]    							=							decay_schedule_fn
								__magic_name__							:	Dict    							=							name

			def __call__( self			:   Any       ,						lowerCamelCase_			:   Dict  )						->   Optional[Any]:
								with tf.name_scope(self.name or '''WarmUp'''  ) as name:
													# Implements polynomial warmup. i.e., if global_step < warmup_steps, the
													# learning rate will be `global_step/num_warmup_steps * init_lr`.
													__magic_name__							:	Union[str, Any]    							=							tf.cast(lowerCamelCase_       ,						tf.floataa  )
													__magic_name__							:	Dict    							=							tf.cast(self.warmup_steps       ,						tf.floataa  )
													__magic_name__							:	Dict    							=							global_step_float / warmup_steps_float
													__magic_name__							:	List[str]    							=							self.initial_learning_rate * tf.math.pow(lowerCamelCase_       ,						self.power  )
													return tf.cond(
													    global_step_float < warmup_steps_float       ,						lambda: warmup_learning_rate       ,						lambda: self.decay_schedule_fn(step - self.warmup_steps  )       ,						name=lowerCamelCase_       ,						)


			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   str:
								return {
								    "initial_learning_rate": self.initial_learning_rate,
								    "decay_schedule_fn": self.decay_schedule_fn,
								    "warmup_steps": self.warmup_steps,
								    "power": self.power,
								    "name": self.name,
								}







def 		lowercase__   (		__A:	float							,__A:	int							,__A:	int							,__A:	float = 0.0							,__A:	float = 0.9							,__A:	float = 0.999							,__A:	float = 1e-8							,__A:	Optional[float] = None							,__A:	Optional[float] = None							,__A:	float = 0.0							,__A:	float = 1.0							,__A:	Optional[List[str]] = None							,):





					'''simple docstring'''






					__magic_name__							:	Any    							=							tf.keras.optimizers.schedules.PolynomialDecay(
					    initial_learning_rate=__A							,decay_steps=num_train_steps - num_warmup_steps							,end_learning_rate=init_lr * min_lr_ratio							,power=__A							,)
					if num_warmup_steps:
										__magic_name__							:	Tuple    							=							WarmUp(
										    initial_learning_rate=__A							,decay_schedule_fn=__A							,warmup_steps=__A							,)
					if weight_decay_rate > 0.0:
										__magic_name__							:	Union[str, Any]    							=							AdamWeightDecay(
										    learning_rate=__A							,weight_decay_rate=__A							,beta_a=__A							,beta_a=__A							,epsilon=__A							,clipnorm=__A							,global_clipnorm=__A							,exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias''']							,include_in_weight_decay=__A							,)
					else:
										__magic_name__							:	str    							=							tf.keras.optimizers.Adam(
										    learning_rate=__A							,beta_a=__A							,beta_a=__A							,epsilon=__A							,clipnorm=__A							,global_clipnorm=__A							,)
					# We return the optimizer and the LR scheduler in order to better track the
					# evolution of the LR independently of the optimizer.
					return optimizer, lr_schedule


class    lowerCamelCase    (		_lowerCamelCase						):






			'''simple docstring'''





			def __init__( self			:   Optional[int]       ,						lowerCamelCase_			:   Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.0_0_1       ,						lowerCamelCase_			:   float = 0.9       ,						lowerCamelCase_			:   float = 0.9_9_9       ,						lowerCamelCase_			:   float = 1E-7       ,						lowerCamelCase_			:   bool = False       ,						lowerCamelCase_			:   float = 0.0       ,						lowerCamelCase_			:   Optional[List[str]] = None       ,						lowerCamelCase_			:   Optional[List[str]] = None       ,						lowerCamelCase_			:   str = "AdamWeightDecay"       ,						**lowerCamelCase_			:   int       ,						)						->   List[str]:
								super().__init__(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						**lowerCamelCase_  )
								__magic_name__							:	int    							=							weight_decay_rate
								__magic_name__							:	Tuple    							=							include_in_weight_decay
								__magic_name__							:	Union[str, Any]    							=							exclude_from_weight_decay

			@classmethod
			def 							UpperCAmelCase__       ( cls			:   int       ,						lowerCamelCase_			:   int  )						->   Optional[Any]:
								__magic_name__							:	Tuple    							=							{'''WarmUp''': WarmUp}
								return super(lowerCamelCase_       ,						cls  ).from_config(lowerCamelCase_       ,						custom_objects=lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   Union[str, Any]       ,						lowerCamelCase_			:   int       ,						lowerCamelCase_			:   Optional[int]       ,						lowerCamelCase_			:   Optional[int]  )						->   int:
								super(lowerCamelCase_       ,						self  )._prepare_local(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
								__magic_name__							:	Any    							=							tf.constant(
								    self.weight_decay_rate       ,						name='''adam_weight_decay_rate'''  )

			def 							UpperCAmelCase__       ( self			:   Optional[int]       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   Any  )						->   Optional[Any]:
								__magic_name__							:	Optional[int]    							=							self._do_use_weight_decay(var.name  )
								if do_decay:
													return var.assign_sub(
													    learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['''weight_decay_rate''']       ,						use_locking=self._use_locking       ,						)
								return tf.no_op()

			def 							UpperCAmelCase__       ( self			:   str       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   int=None       ,						**lowerCamelCase_			:   Union[str, Any]  )						->   Optional[int]:
								__magic_name__							,    __magic_name__							:	Dict    							=							list(zip(*lowerCamelCase_  )  )
								return super(lowerCamelCase_       ,						self  ).apply_gradients(zip(lowerCamelCase_       ,						lowerCamelCase_  )       ,						name=lowerCamelCase_       ,						**lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   List[Any]       ,						lowerCamelCase_			:   Optional[int]       ,						lowerCamelCase_			:   int       ,						lowerCamelCase_			:   Union[str, Any]  )						->   Union[str, Any]:
								if apply_state is None:
													return self._decayed_lr_t[var_dtype], {}

								__magic_name__							:	str    							=							apply_state or {}
								__magic_name__							:	Optional[Any]    							=							apply_state.get((var_device, var_dtype)  )
								if coefficients is None:
													__magic_name__							:	List[str]    							=							self._fallback_apply_state(lowerCamelCase_       ,						lowerCamelCase_  )
													__magic_name__							:	List[Any]    							=							coefficients

								return coefficients["lr_t"], {"apply_state": apply_state}

			def 							UpperCAmelCase__       ( self			:   List[str]       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Dict=None  )						->   Optional[Any]:
								__magic_name__							,    __magic_name__							:	Union[str, Any]    							=							self._get_lr(var.device       ,						var.dtype.base_dtype       ,						lowerCamelCase_  )
								__magic_name__							:	Optional[Any]    							=							self._decay_weights_op(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
								with tf.control_dependencies([decay]  ):
													return super(lowerCamelCase_       ,						self  )._resource_apply_dense(lowerCamelCase_       ,						lowerCamelCase_       ,						**lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   Dict       ,						lowerCamelCase_			:   List[Any]       ,						lowerCamelCase_			:   List[Any]       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   Union[str, Any]=None  )						->   Any:
								__magic_name__							,    __magic_name__							:	Tuple    							=							self._get_lr(var.device       ,						var.dtype.base_dtype       ,						lowerCamelCase_  )
								__magic_name__							:	Optional[int]    							=							self._decay_weights_op(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
								with tf.control_dependencies([decay]  ):
													return super(lowerCamelCase_       ,						self  )._resource_apply_sparse(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						**lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   int  )						->   Dict:
								__magic_name__							:	List[Any]    							=							super().get_config()
								config.update({'''weight_decay_rate''': self.weight_decay_rate}  )
								return config


			def 							UpperCAmelCase__       ( self			:   Dict       ,						lowerCamelCase_			:   Union[str, Any]  )						->   Union[str, Any]:
								if self.weight_decay_rate == 0:
													return False

								if self._include_in_weight_decay:
													for r in self._include_in_weight_decay:
																		if re.search(lowerCamelCase_       ,						lowerCamelCase_  ) is not None:
																							return True

								if self._exclude_from_weight_decay:
													for r in self._exclude_from_weight_decay:
																		if re.search(lowerCamelCase_       ,						lowerCamelCase_  ) is not None:
																							return False
								return True






class    lowerCamelCase    (		_lowerCamelCase						):






			'''simple docstring'''





			def __init__( self			:   Optional[Any]  )						->   Any:
								__magic_name__							:	Optional[int]    							=							[]
								__magic_name__							:	Dict    							=							None

			@property
			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Any:
								if self._accum_steps is None:
													__magic_name__							:	Optional[Any]    							=							tf.Variable(
													    tf.constant(0       ,						dtype=tf.intaa  )       ,						trainable=lowerCamelCase_       ,						synchronization=tf.VariableSynchronization.ON_READ       ,						aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA       ,						)

								return self._accum_steps.value()

			@property
			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Optional[Any]:
								if not self._gradients:
													raise ValueError('''The accumulator should be called first to initialize the gradients'''  )
								return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]

			def __call__( self			:   Optional[int]       ,						lowerCamelCase_			:   Any  )						->   List[str]:
								if not self._gradients:
													__magic_name__							:	int    							=							self.step  # Create the step variable.
													self._gradients.extend(
													    [
													        tf.Variable(
													            tf.zeros_like(lowerCamelCase_  )       ,						trainable=lowerCamelCase_       ,						synchronization=tf.VariableSynchronization.ON_READ       ,						aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA       ,						)
													        if gradient is not None
													        else gradient
													        for gradient in gradients
													    ]  )
								if len(lowerCamelCase_  ) != len(self._gradients  ):
													raise ValueError(F'''Expected {len(self._gradients  )} gradients, but got {len(lowerCamelCase_  )}'''  )

								for accum_gradient, gradient in zip(self._gradients       ,						lowerCamelCase_  ):
													if accum_gradient is not None and gradient is not None:
																		accum_gradient.assign_add(lowerCamelCase_  )

								self._accum_steps.assign_add(1  )


			def 							UpperCAmelCase__       ( self			:   str  )						->   Optional[int]:
								if not self._gradients:
													return
								self._accum_steps.assign(0  )
								for gradient in self._gradients:
													if gradient is not None:
																		gradient.assign(tf.zeros_like(lowerCamelCase_  )  )



