




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







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










from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__					    =    logging.get_logger(__name__)

UpperCamelCase__					    =    {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class     a__					(      snake_case__   ):
						_a							:					Tuple      =					"""nllb-moe"""
						_a							:					Dict      =					["""past_key_values"""]
						_a							:					Optional[int]      =					{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






						def __init__(     self	,    _A=1_2_8_1_1_2	,    _A=1_0_2_4	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=0.05	,    _A=0.05	,    _A=True	,    _A=True	,    _A="relu"	,    _A=1_0_2_4	,    _A=0.1	,    _A=0.1	,    _A=0.0	,    _A=0.02	,    _A=2	,    _A=True	,    _A=False	,    _A="float32"	,    _A=False	,    _A=1_2_8	,    _A=6_4	,    _A=4	,    _A=4	,    _A=0.0_01	,    _A=0.0_01	,    _A="all"	,    _A=False	,    _A=False	,    _A=1.0	,    _A=0.2	,    _A=1	,    _A=0	,    _A=2	,    _A=False	,    **_A	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = d_model
											__lowerCAmelCase            = encoder_ffn_dim
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = encoder_attention_heads
											__lowerCAmelCase            = decoder_ffn_dim
											__lowerCAmelCase            = decoder_layers
											__lowerCAmelCase            = decoder_attention_heads
											__lowerCAmelCase            = dropout
											__lowerCAmelCase            = attention_dropout
											__lowerCAmelCase            = activation_dropout
											__lowerCAmelCase            = activation_function
											__lowerCAmelCase            = init_std
											__lowerCAmelCase            = encoder_layerdrop
											__lowerCAmelCase            = decoder_layerdrop
											__lowerCAmelCase            = use_cache
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = scale_embedding  # scale factor will be sqrt(d_model) if True
											__lowerCAmelCase            = router_z_loss_coef
											__lowerCAmelCase            = router_aux_loss_coef
											__lowerCAmelCase            = decoder_sparse_step
											__lowerCAmelCase            = encoder_sparse_step
											__lowerCAmelCase            = num_experts
											__lowerCAmelCase            = expert_capacity
											__lowerCAmelCase            = router_bias
											if router_dtype not in ["float32", "float16", "bfloat16"]:
																raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""    )
											__lowerCAmelCase            = router_dtype

											__lowerCAmelCase            = router_ignore_padding_tokens
											__lowerCAmelCase            = batch_prioritized_routing
											__lowerCAmelCase            = second_expert_policy
											__lowerCAmelCase            = normalize_router_prob_before_dropping
											__lowerCAmelCase            = moe_eval_capacity_token_fraction
											__lowerCAmelCase            = moe_token_dropout
											__lowerCAmelCase            = output_router_logits
											super().__init__(
											    pad_token_id=_A	,    bos_token_id=_A	,    eos_token_id=_A	,    is_encoder_decoder=_A	,    decoder_start_token_id=_A	,    **_A	,    )




