


"""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__    :							Optional[Any]            =					logging.get_logger(__name__)




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_       ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						       =    os.path.abspath(__lowerCamelCase       )
					logger.info(f"""Converting TensorFlow checkpoint from {tf_path}"""       )
					# Load weights from TF model
					__SCREAMING_SNAKE_CASE						       =    tf.train.list_variables(__lowerCamelCase       )
					__SCREAMING_SNAKE_CASE						       =    []
					__SCREAMING_SNAKE_CASE						       =    []
					__SCREAMING_SNAKE_CASE						       =    []
					for full_name, shape in init_vars:
										# logger.info(f"Loading TF weight {name} with shape {shape}")
										__SCREAMING_SNAKE_CASE						       =    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'
															__SCREAMING_SNAKE_CASE						       =    name[1:]
										# figure out how many levels deep the name is
										__SCREAMING_SNAKE_CASE						       =    0
										for _name in name:
															if _name.startswith("layer_with_weights"       ):
																				depth += 1
															else:
																				break
										layer_depth.append(__lowerCamelCase       )
										# read data
										__SCREAMING_SNAKE_CASE						       =    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       )       )})"""       )
					__SCREAMING_SNAKE_CASE						       =    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       ):
										__SCREAMING_SNAKE_CASE						       =    full_name.split("/"       )
										__SCREAMING_SNAKE_CASE						       =    model
										__SCREAMING_SNAKE_CASE						       =    []
										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"       ):
																				__SCREAMING_SNAKE_CASE						       =    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"]       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"embeddings"       )
																									__SCREAMING_SNAKE_CASE						       =    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       )]       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"encoder"       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"layer"       )
																									__SCREAMING_SNAKE_CASE						       =    pointer[layer_num - 4]
																				elif layer_num == config.num_hidden_layers + 4:
																									# pooler layer
																									trace.extend(["pooler", "dense"]       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"pooler"       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"dense"       )
															elif m_name == "embeddings":
																				trace.append("embeddings"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"embeddings"       )
																				if layer_num == 0:
																									trace.append("word_embeddings"       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"word_embeddings"       )
																				elif layer_num == 1:
																									trace.append("position_embeddings"       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"position_embeddings"       )
																				elif layer_num == 2:
																									trace.append("token_type_embeddings"       )
																									__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"token_type_embeddings"       )
																				else:
																									raise ValueError(f"""Unknown embedding layer with name {full_name}"""       )
																				trace.append("weight"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"weight"       )
															elif m_name == "_attention_layer":
																				# self-attention layer
																				trace.extend(["attention", "self"]       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"attention"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"self"       )
															elif m_name == "_attention_layer_norm":
																				# output attention norm
																				trace.extend(["attention", "output", "LayerNorm"]       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"attention"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"output"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"LayerNorm"       )
															elif m_name == "_attention_output_dense":
																				# output attention dense
																				trace.extend(["attention", "output", "dense"]       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"attention"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"output"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"dense"       )
															elif m_name == "_output_dense":
																				# output dense
																				trace.extend(["output", "dense"]       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"output"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"dense"       )
															elif m_name == "_output_layer_norm":
																				# output dense
																				trace.extend(["output", "LayerNorm"]       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"output"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"LayerNorm"       )
															elif m_name == "_key_dense":
																				# attention key
																				trace.append("key"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"key"       )
															elif m_name == "_query_dense":
																				# attention query
																				trace.append("query"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"query"       )
															elif m_name == "_value_dense":
																				# attention value
																				trace.append("value"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"value"       )
															elif m_name == "_intermediate_dense":
																				# attention intermediate dense
																				trace.extend(["intermediate", "dense"]       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"intermediate"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"dense"       )
															elif m_name == "_output_layer_norm":
																				# output layer norm
																				trace.append("output"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"output"       )
															# weights & biases
															elif m_name in ["bias", "beta"]:
																				trace.append("bias"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"bias"       )
															elif m_name in ["kernel", "gamma"]:
																				trace.append("weight"       )
																				__SCREAMING_SNAKE_CASE						       =    getattr(__lowerCamelCase  ,						"weight"       )
															else:
																				logger.warning(f"""Ignored {m_name}"""       )
        # for certain layers reshape is necessary
										__SCREAMING_SNAKE_CASE						       =    ".".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       ):
															__SCREAMING_SNAKE_CASE						       =    array.reshape(pointer.data.shape       )
										if "kernel" in full_name:
															__SCREAMING_SNAKE_CASE						       =    array.transpose()
										if pointer.shape == array.shape:
															__SCREAMING_SNAKE_CASE						       =    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        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_       ):




					'''simple docstring'''


					logger.info(f"""Loading model based on config from {config_path}..."""       )
					__SCREAMING_SNAKE_CASE						       =    BertConfig.from_json_file(__lowerCamelCase       )
					__SCREAMING_SNAKE_CASE						       =    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__":
		a__    :							Optional[Any]            =					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__    :							Optional[Any]            =					parser.parse_args()
		convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)



"""simple docstring"""







import argparse
import re

from flax.traverse_util import flatten_dict, unflatten_dict
from tax import checkpoints

from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
from transformers.utils import logging


logging.set_verbosity_info()


# should not include what is already done by the `from_pt` argument
a__    :							Optional[Any]            =					{
    '''/attention/''': '''/0/SelfAttention/''',
    '''/self_attention/''': '''/0/SelfAttention/''',
    '''/encoder_decoder_attention/''': '''/1/EncDecAttention/''',
    '''value''': '''v''',
    '''query''': '''q''',
    '''key''': '''k''',
    '''out''': '''o''',
    '''pre_self_attention_layer_norm''': '''0/layer_norm''',
    '''pre_cross_attention_layer_norm''': '''1/layer_norm''',
    '''pre_attention_layer_norm''': '''0/layer_norm''',  # previously 1, but seems wrong
    '''token_embedder''': '''shared''',
    '''encoder_norm''': '''final_layer_norm''',
    '''decoder_norm''': '''final_layer_norm''',
    '''relpos_bias/rel_embedding''': '''block/0/layer/0/SelfAttention/relative_attention_bias/weight''',
    '''router/router_weights/w/''': '''router/classifier/''',
    '''roer/roer_weights/w/''': '''router/classifier/''',
    '''logits_dense''': '''lm_head''',
}




def        UpperCAmelCase__	(lowerCAmelCase_       ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						       =    list(s_dict.keys()       )
					for key in keys:
										__SCREAMING_SNAKE_CASE						       =    R".*/layers_(\d+)"
										__SCREAMING_SNAKE_CASE						       =    key
										if re.match(lowerCAmelCase_  ,						lowerCAmelCase_       ):
															__SCREAMING_SNAKE_CASE						       =    re.sub(R"layers_(\d+)"  ,						R"block/\1/layer"  ,						lowerCAmelCase_       )

										__SCREAMING_SNAKE_CASE						       =    R"(encoder|decoder)\/"

										if re.match(lowerCAmelCase_  ,						lowerCAmelCase_       ):
															__SCREAMING_SNAKE_CASE						       =    re.match(lowerCAmelCase_  ,						lowerCAmelCase_       ).groups()
															if groups[0] == "encoder":
																				__SCREAMING_SNAKE_CASE						       =    re.sub(R"/mlp/"  ,						R"/1/mlp/"  ,						lowerCAmelCase_       )
																				__SCREAMING_SNAKE_CASE						       =    re.sub(R"/pre_mlp_layer_norm/"  ,						R"/1/layer_norm/"  ,						lowerCAmelCase_       )

															elif groups[0] == "decoder":
																				__SCREAMING_SNAKE_CASE						       =    re.sub(R"/mlp/"  ,						R"/2/mlp/"  ,						lowerCAmelCase_       )
																				__SCREAMING_SNAKE_CASE						       =    re.sub(R"/pre_mlp_layer_norm/"  ,						R"/2/layer_norm/"  ,						lowerCAmelCase_       )

        # 2. Convert other classic mappings
										for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
															if old_key in new_key:
																				__SCREAMING_SNAKE_CASE						       =    new_key.replace(lowerCAmelCase_  ,						lowerCAmelCase_       )

										print(f"""{key} -> {new_key}"""       )
										__SCREAMING_SNAKE_CASE						       =    s_dict.pop(lowerCAmelCase_       )

					if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
										__SCREAMING_SNAKE_CASE						       =    s_dict[
										    "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
										].T
					if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
										__SCREAMING_SNAKE_CASE						       =    s_dict[
										    "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
										].T

					# 3. Take extra care of the EXPERTS layer
					for key in list(s_dict.keys()       ):
										if "expert" in key:
															__SCREAMING_SNAKE_CASE						       =    s_dict[key].shape[0]
															__SCREAMING_SNAKE_CASE						       =    s_dict[key]
															for idx in range(lowerCAmelCase_       ):
																				__SCREAMING_SNAKE_CASE						       =    expert_weihts[idx]
																				print(f"""{key} -> {key.replace('expert/'  ,						'nested fstring'       )}"""       )

															s_dict.pop(lowerCAmelCase_       )

					return s_dict


a__    :							List[Any]            =					{
    '''NUM_ENCODER_LAYERS''': '''num_layers''',
    '''NUM_DECODER_LAYERS''': '''num_decoder_layers''',
    '''NUM_HEADS''': '''num_heads''',
    '''HEAD_DIM''': '''d_kv''',
    '''EMBED_DIM''': '''d_model''',
    '''MLP_DIM''': '''d_ff''',
    '''NUM_SELECTED_EXPERTS''': '''num_selected_experts''',
    '''NUM_ENCODER_SPARSE_LAYERS''': '''num_sparse_encoder_layers''',
    '''NUM_DECODER_SPARSE_LAYERS''': '''num_sparse_decoder_layers''',
    '''dense.MlpBlock.activations''': '''feed_forward_proj''',
}




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_       ):




					'''simple docstring'''


					import regex as re

					with open(lowerCAmelCase_  ,						"r"       ) as f:
										__SCREAMING_SNAKE_CASE						       =    f.read()

					__SCREAMING_SNAKE_CASE						       =    re.findall(R"(.*) = ([0-9.]*)"  ,						lowerCAmelCase_       )
					__SCREAMING_SNAKE_CASE						       =    {}
					for param, value in regex_match:
										if param in GIN_TO_CONFIG_MAPPING and value != "":
															__SCREAMING_SNAKE_CASE						       =    float(lowerCAmelCase_       ) if "." in value else int(lowerCAmelCase_       )

					__SCREAMING_SNAKE_CASE						       =    re.findall(R"(.*activations) = \(\'(.*)\',\)"  ,						lowerCAmelCase_       )[0]
					__SCREAMING_SNAKE_CASE						       =    str(activation[1]       )

					__SCREAMING_SNAKE_CASE						       =    num_experts
					__SCREAMING_SNAKE_CASE						       =    SwitchTransformersConfig(**lowerCAmelCase_       )
					return config




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_="./"  ,						lowerCAmelCase_=8       ):




					'''simple docstring'''


					print(f"""Loading flax weights from : {flax_checkpoint_path}"""       )
					__SCREAMING_SNAKE_CASE						       =    checkpoints.load_tax_checkpoint(lowerCAmelCase_       )

					if gin_file is not None:
										__SCREAMING_SNAKE_CASE						       =    convert_gin_to_config(lowerCAmelCase_  ,						lowerCAmelCase_       )
					else:
										__SCREAMING_SNAKE_CASE						       =    SwitchTransformersConfig.from_pretrained(lowerCAmelCase_       )

					__SCREAMING_SNAKE_CASE						       =    SwitchTransformersForConditionalGeneration(lowerCAmelCase_       )

					__SCREAMING_SNAKE_CASE						       =    flax_params["target"]
					__SCREAMING_SNAKE_CASE						       =    flatten_dict(lowerCAmelCase_  ,						sep="/"       )
					__SCREAMING_SNAKE_CASE						       =    rename_keys(lowerCAmelCase_       )
					__SCREAMING_SNAKE_CASE						       =    unflatten_dict(lowerCAmelCase_  ,						sep="/"       )

					# Load the flax params in the PT model
					load_flax_weights_in_pytorch_model(lowerCAmelCase_  ,						lowerCAmelCase_       )

					print(f"""Save PyTorch model to {pytorch_dump_path}"""       )
					pt_model.save_pretrained(lowerCAmelCase_       )


if __name__ == "__main__":
		a__    :							List[Any]            =					argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--switch_t5x_checkpoint_path''',
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        '''The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the'''
		        ''' model architecture. If not provided, a `gin_file` has to be provided.'''
		    ),
		)
		parser.add_argument(
		    '''--gin_file''',
		    default=None,
		    type=str,
		    required=False,
		    help='''Path to the gin config file. If not provided, a `config_file` has to be passed   ''',
		)
		parser.add_argument(
		    '''--config_name''', default=None, type=str, required=False, help='''Config name of SwitchTransformers model.'''
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output pytorch model.'''
		)
		parser.add_argument('''--num_experts''', default=8, type=int, required=False, help='''Number of experts''')
		a__    :							Tuple            =					parser.parse_args()
		convert_flax_checkpoint_to_pytorch(
		    args.switch_tax_checkpoint_path,
		    args.config_name,
		    args.gin_file,
		    args.pytorch_dump_folder_path,
		    args.num_experts,
		)
