



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

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

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

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


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)










"""simple docstring"""


import math

def 			_lowerCAmelCase				(       lowercase_						):
						assert isinstance(lowercase_		,     lowercase_						) and (
						    number >= 0
						), "'number' must been an int and positive"

						if 1 < number < 4:
												# 2 and 3 are primes
												return True
						elif number < 2 or not number % 2:
												# Negatives, 0, 1 and all even numbers are not primes
												return False

						UpperCAmelCase					=			range(3		,     int(math.sqrt(lowercase_						) + 1						)		,     2						)
						return not any(not number % i for i in odd_numbers						)






def 			_lowerCAmelCase				(       lowercase_		,     lowercase_=1		,     **lowercase_						):
						UpperCAmelCase					=			factor * value
						UpperCAmelCase					=			value

						while not is_prime(lowercase_						):
												value += 1 if not ("desc" in kwargs and kwargs["desc"] is True) else -1

						if value == first_value_val:
												return next_prime(value + 1		,     **lowercase_						)
						return value





