




'''simple docstring'''





import tensorflow as tf

from ...tf_utils import shape_list
class 		SCREAMING_SNAKE_CASE			(     tf.keras.layers.Layer							):




				'''simple docstring'''





				def __init__(   self  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__=1  ,						SCREAMING_SNAKE_CASE__=False  ,						**SCREAMING_SNAKE_CASE__				):
							'''simple docstring'''







							super().__init__(**_lowerCamelCase				)

							snake_case:       Dict        =       vocab_size
							snake_case:       Optional[Any]        =       d_embed
							snake_case:       Tuple        =       d_proj

							snake_case:       Optional[int]        =       cutoffs + [vocab_size]
							snake_case:       Optional[Any]        =       [0] + self.cutoffs
							snake_case:       List[Any]        =       div_val

							snake_case:       List[Any]        =       self.cutoffs[0]
							snake_case:       Any        =       len(self.cutoffs				) - 1
							snake_case:       Optional[Any]        =       self.shortlist_size + self.n_clusters
							snake_case:       Dict        =       keep_order

							snake_case:       str        =       []
							snake_case:       Tuple        =       []







				def        _UpperCamelCase							(   self  ,						SCREAMING_SNAKE_CASE__				):
							'''simple docstring'''







							if self.n_clusters > 0:
										snake_case:       Optional[Any]        =       self.add_weight(
										    shape=(self.n_clusters, self.d_embed)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name='cluster_weight'				)
										snake_case:       Any        =       self.add_weight(
										    shape=(self.n_clusters,)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name='cluster_bias'				)

							if self.div_val == 1:
										for i in range(len(self.cutoffs				)				):
													if self.d_proj != self.d_embed:
																snake_case:       str        =       self.add_weight(
																    shape=(self.d_embed, self.d_proj)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name=F"""out_projs_._{i}"""  ,						)
																self.out_projs.append(_lowerCamelCase				)
													else:
																self.out_projs.append(_lowerCamelCase				)
													snake_case:       Dict        =       self.add_weight(
													    shape=(self.vocab_size, self.d_embed)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name=F"""out_layers_._{i}_._weight"""  ,						)
													snake_case:       Union[str, Any]        =       self.add_weight(
													    shape=(self.vocab_size,)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name=F"""out_layers_._{i}_._bias"""  ,						)
													self.out_layers.append((weight, bias)				)
							else:
										for i in range(len(self.cutoffs				)				):
													snake_case    ,						snake_case:       str        =       self.cutoff_ends[i], self.cutoff_ends[i + 1]
													snake_case:       Tuple        =       self.d_embed // (self.div_val**i)

													snake_case:       Union[str, Any]        =       self.add_weight(
													    shape=(d_emb_i, self.d_proj)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name=F"""out_projs_._{i}"""				)
													self.out_projs.append(_lowerCamelCase				)
													snake_case:       List[Any]        =       self.add_weight(
													    shape=(r_idx - l_idx, d_emb_i)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name=F"""out_layers_._{i}_._weight"""  ,						)
													snake_case:       List[str]        =       self.add_weight(
													    shape=(r_idx - l_idx,)  ,						initializer='zeros'  ,						trainable=_lowerCamelCase  ,						name=F"""out_layers_._{i}_._bias"""  ,						)
													self.out_layers.append((weight, bias)				)
							super().build(_lowerCamelCase				)







				@staticmethod
				def        _UpperCamelCase							(   SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__=None				):
							'''simple docstring'''







							snake_case:       Optional[int]        =       x
							if proj is not None:
										snake_case:       List[Any]        =       tf.einsum('ibd,ed->ibe'  ,						_lowerCamelCase  ,						_lowerCamelCase				)
							return tf.einsum('ibd,nd->ibn'  ,						_lowerCamelCase  ,						_lowerCamelCase				) + b







				@staticmethod
				def        _UpperCamelCase							(   SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__				):
							'''simple docstring'''







							snake_case:       Dict        =       shape_list(_lowerCamelCase				)
							snake_case:       List[str]        =       tf.range(lp_size[0]  ,						dtype=target.dtype				)
							snake_case:       List[Any]        =       tf.stack([r, target]  ,						1				)
							return tf.gather_nd(_lowerCamelCase  ,						_lowerCamelCase				)







				def        _UpperCamelCase							(   self  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__=True  ,						SCREAMING_SNAKE_CASE__=False				):
							'''simple docstring'''







							snake_case:       Tuple        =       0
							if self.n_clusters == 0:
										snake_case:       Dict        =       self._logit(_lowerCamelCase  ,						self.out_layers[0][0]  ,						self.out_layers[0][1]  ,						self.out_projs[0]				)
										if target is not None:
													snake_case:       Any        =       tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_lowerCamelCase  ,						logits=_lowerCamelCase				)
										snake_case:       List[str]        =       tf.nn.log_softmax(_lowerCamelCase  ,						axis=-1				)
							else:
										snake_case:       Tuple        =       shape_list(_lowerCamelCase				)
										snake_case:       Tuple        =       []
										snake_case:       Any        =       tf.zeros(hidden_sizes[:2]				)
										for i in range(len(self.cutoffs				)				):
													snake_case    ,						snake_case:       Dict        =       self.cutoff_ends[i], self.cutoff_ends[i + 1]
													if target is not None:
																snake_case:       Union[str, Any]        =       (target >= l_idx) & (target < r_idx)
																snake_case:       List[Any]        =       tf.where(_lowerCamelCase				)
																snake_case:       List[str]        =       tf.boolean_mask(_lowerCamelCase  ,						_lowerCamelCase				) - l_idx

													if self.div_val == 1:
																snake_case:       Dict        =       self.out_layers[0][0][l_idx:r_idx]
																snake_case:       int        =       self.out_layers[0][1][l_idx:r_idx]
													else:
																snake_case:       List[str]        =       self.out_layers[i][0]
																snake_case:       Optional[Any]        =       self.out_layers[i][1]

													if i == 0:
																snake_case:       Optional[Any]        =       tf.concat([cur_W, self.cluster_weight]  ,						0				)
																snake_case:       Any        =       tf.concat([cur_b, self.cluster_bias]  ,						0				)

																snake_case:       List[str]        =       self._logit(_lowerCamelCase  ,						_lowerCamelCase  ,						_lowerCamelCase  ,						self.out_projs[0]				)
																snake_case:       List[Any]        =       tf.nn.log_softmax(_lowerCamelCase				)
																out.append(head_logprob[..., : self.cutoffs[0]]				)
																if target is not None:
																			snake_case:       int        =       tf.boolean_mask(_lowerCamelCase  ,						_lowerCamelCase				)
																			snake_case:       Optional[int]        =       self._gather_logprob(_lowerCamelCase  ,						_lowerCamelCase				)
													else:
																snake_case:       Any        =       self._logit(_lowerCamelCase  ,						_lowerCamelCase  ,						_lowerCamelCase  ,						self.out_projs[i]				)
																snake_case:       Dict        =       tf.nn.log_softmax(_lowerCamelCase				)
																snake_case:       List[Any]        =       self.cutoffs[0] + i - 1  # No probability for the head cluster
																snake_case:       Dict        =       head_logprob[..., cluster_prob_idx, None] + tail_logprob
																out.append(_lowerCamelCase				)
																if target is not None:
																			snake_case:       int        =       tf.boolean_mask(_lowerCamelCase  ,						_lowerCamelCase				)
																			snake_case:       Dict        =       tf.boolean_mask(_lowerCamelCase  ,						_lowerCamelCase				)
																			snake_case:       Optional[Any]        =       self._gather_logprob(_lowerCamelCase  ,						_lowerCamelCase				)
																			cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
													if target is not None:
																loss += tf.scatter_nd(_lowerCamelCase  ,						-cur_logprob  ,						shape_list(_lowerCamelCase				)				)
										snake_case:       Tuple        =       tf.concat(_lowerCamelCase  ,						axis=-1				)

							if target is not None:
										if return_mean:
													snake_case:       Optional[Any]        =       tf.reduce_mean(_lowerCamelCase				)
										# Add the training-time loss value to the layer using `self.add_loss()`.
										self.add_loss(_lowerCamelCase				)

										# Log the loss as a metric (we could log arbitrary metrics,
										# including different metrics for training and inference.
										self.add_metric(_lowerCamelCase  ,						name=self.name  ,						aggregation='mean' if return_mean else ''				)

							return out

"""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()
_lowercase							=				logging.get_logger(__name__)



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

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

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

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


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





