






"""simple docstring"""




import argparse
import os
import re
import zipfile

import torch

from transformers import MegatronBertConfig

def     lowerCAmelCase__  (  UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__=0       ):



		'''simple docstring'''






		if name is None:
				_a						: Union[str, Any]    =			None
		else:
				_a						: int    =			""".""" * max(0  ,     spaces - 2       ) + """# {:""" + str(5_0 - spaces       ) + """s}"""
				_a						: Optional[int]    =			fmt.format(UpperCamelCase__       )

		# Print and recurse (if needed).
		if isinstance(UpperCamelCase__  ,     UpperCamelCase__       ):
				if msg is not None:
						print(UpperCamelCase__       )
				for k in val.keys():
						recursive_print(UpperCamelCase__  ,     val[k]  ,     spaces + 2       )
		elif isinstance(UpperCamelCase__  ,     torch.Tensor       ):
				print(UpperCamelCase__  ,     """:"""  ,     val.size()       )
		else:
				print(UpperCamelCase__  ,     """:"""  ,     UpperCamelCase__       )

def     lowerCAmelCase__  (  UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__       ):



		'''simple docstring'''
		_a						: Union[str, Any]    =			param.size()
		if checkpoint_version == 1.0:
				# version 1.0 stores [num_heads * hidden_size * num_splits, :]
				_a						: Any    =			(num_heads, hidden_size, num_splits) + input_shape[1:]
				_a						: Any    =			param.view(*UpperCamelCase__       )
				_a						: Optional[Any]    =			param.transpose(0  ,     2       )
				_a						: Dict    =			param.transpose(1  ,     2       ).contiguous()
		elif checkpoint_version >= 2.0:
				# other versions store [num_heads * num_splits * hidden_size, :]
				_a						: Union[str, Any]    =			(num_heads, num_splits, hidden_size) + input_shape[1:]
				_a						: Dict    =			param.view(*UpperCamelCase__       )
				_a						: str    =			param.transpose(0  ,     1       ).contiguous()
		_a						: Optional[Any]    =			param.view(*UpperCamelCase__       )
		return param

def     lowerCAmelCase__  (  UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__       ):



		'''simple docstring'''
		_a						: Union[str, Any]    =			{}

		# old versions did not store training args
		_a						: Any    =			input_state_dict.get("""args"""  ,     UpperCamelCase__       )
		if ds_args is not None:
				# do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
				# from pprint import pprint
				# pprint(vars(ds_args))

				_a						: int    =			ds_args.tokenizer_type
				_a						: Union[str, Any]    =			ds_args.padded_vocab_size
				_a						: Tuple    =			ds_args.max_position_embeddings
				_a						: str    =			ds_args.hidden_size
				_a						: Union[str, Any]    =			ds_args.num_layers
				_a						: Optional[Any]    =			ds_args.num_attention_heads
				_a						: Any    =			ds_args.ffn_hidden_size if """ffn_hidden_size""" in ds_args else 4 * ds_args.hidden_size
				# pprint(config)

		# The number of heads.
		_a						: Tuple    =			config.num_attention_heads
		# The hidden_size per head.
		_a						: Any    =			config.hidden_size // heads
		# Megatron-LM checkpoint version
		if "checkpoint_version" in input_state_dict.keys():
				_a						: Optional[int]    =			input_state_dict["""checkpoint_version"""]
		else:
				_a						: Any    =			0.0

		# The model.
		_a						: List[Any]    =			input_state_dict["""model"""]
		# The language model.
		_a						: int    =			model["""language_model"""]
		# The embeddings.
		_a						: int    =			lm["""embedding"""]

		# The word embeddings.
		_a						: str    =			embeddings["""word_embeddings"""]["""weight"""]
		# Truncate the embedding table to vocab_size rows.
		_a						: str    =			word_embeddings[: config.vocab_size, :]
		# Store the word embeddings.
		_a						: List[str]    =			word_embeddings

		# The position embeddings.
		_a						: Optional[int]    =			embeddings["""position_embeddings"""]["""weight"""]
		assert pos_embeddings.size(0       ) == config.max_position_embeddings and pos_embeddings.size(1       ) == config.hidden_size
		# Store the position embeddings.
		_a						: Tuple    =			pos_embeddings

		# The token-type embeddings.
		_a						: Optional[Any]    =			embeddings["""tokentype_embeddings"""]["""weight"""]
		# Store the position embeddings.
		_a						: Any    =			tokentype_embeddings

		# The transformer.
		_a						: List[str]    =			lm["""transformer"""] if """transformer""" in lm.keys() else lm["""encoder"""]

		# The regex to extract layer names.
		_a						: str    =			re.compile(R"""layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"""       )

		# The simple map of names for "automated" rules.
		_a						: Dict    =			{
		    """attention.dense""": """.attention.output.dense.""",
		    """self_attention.dense""": """.attention.output.dense.""",
		    """mlp.dense_h_to_4h""": """.intermediate.dense.""",
		    """mlp.dense_4h_to_h""": """.output.dense.""",
		}

		# Keep track of the attention/query/value tensor.
		_a						: str    =			None

		# Extract the layers.
		for key, val in transformer.items():
				# Match the name.
				_a						: Dict    =			layer_re.match(UpperCamelCase__       )

				# Stop if that's not a layer
				if m is None:
						break

				# The index of the layer.
				_a						: Tuple    =			int(m.group(1       )       )
				# The name of the operation.
				_a						: Tuple    =			m.group(2       )
				# Is it a weight or a bias?
				_a						: Optional[int]    =			m.group(3       )

				# The name of the layer.
				_a						: int    =			F"""bert.encoder.layer.{layer_idx}"""

				# For layernorm(s), simply store the layer norm.
				if op_name.endswith("""layernorm"""       ):
						_a						: Optional[Any]    =			"""attention.ln""" if op_name.startswith("""input"""       ) else """ln"""
						_a						: Union[str, Any]    =			val

				# Transpose the QKV matrix.
				elif (
				    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
				) and weight_or_bias == "weight":
						# Make sure the QKV pointer is nil.
						assert attention_qkv_weight is None, ""

						_a						: Union[str, Any]    =			fix_query_key_value_ordering(UpperCamelCase__  ,     UpperCamelCase__  ,     3  ,     UpperCamelCase__  ,     UpperCamelCase__       )
						# Store the tensor as we need the bias as well to interleave QKV and biases.
						_a						: List[Any]    =			out_val

				# Transpose the bias.
				elif (
				    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
				) and weight_or_bias == "bias":
						# Make sure we read the weight tensor.
						assert attention_qkv_weight is not None, ""

						# Split the QKV matrix into Q, K and V. Megatron stores Q,K,V interleaved.
						_a						: List[Any]    =			attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
						_a						: Any    =			attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
						_a						: Union[str, Any]    =			attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

						_a						: Optional[Any]    =			fix_query_key_value_ordering(UpperCamelCase__  ,     UpperCamelCase__  ,     3  ,     UpperCamelCase__  ,     UpperCamelCase__       )
						# Split the bias.
						_a						: int    =			out_val[0 * config.hidden_size : 1 * config.hidden_size]
						_a						: Optional[int]    =			out_val[1 * config.hidden_size : 2 * config.hidden_size]
						_a						: List[str]    =			out_val[2 * config.hidden_size : 3 * config.hidden_size]

						# Store.
						_a						: str    =			q
						_a						: Tuple    =			q_bias
						_a						: List[Any]    =			k
						_a						: Tuple    =			k_bias
						_a						: str    =			v
						_a						: Dict    =			v_bias

						# Clear the stored tensor.
						_a						: Dict    =			None

				# Copy weights and biases as is.
				elif weight_or_bias in ["weight", "bias"]:
						_a						: Optional[int]    =			megatron_to_transformers[op_name]
						_a						: Any    =			val

    # The final layernorm.
		_a						: Dict    =			transformer["""final_layernorm.weight"""]
		_a						: Optional[int]    =			transformer["""final_layernorm.bias"""]

		# The pooler.
		_a						: Dict    =			lm["""pooler"""]

		# Store the matrix and the bias.
		_a						: Optional[int]    =			pooler["""dense.weight"""]
		_a						: List[str]    =			pooler["""dense.bias"""]

		# The LM head from Megatron (for RACE).
		_a						: Any    =			model["""lm_head"""]

		# The transform matrix.
		_a						: Optional[int]    =			lm_head["""dense.weight"""]
		_a						: Optional[int]    =			lm_head["""dense.bias"""]

		# The transform LN.
		_a						: Optional[int]    =			lm_head["""layernorm.weight"""]
		_a						: Tuple    =			lm_head["""layernorm.bias"""]

		# For the decoder, we replicate the weights.
		_a						: Optional[int]    =			word_embeddings
		_a						: Optional[Any]    =			lm_head["""bias"""]

		# The classifier from Megatron (for MLNI).
		_a						: str    =			model["""binary_head"""]

		# Store the classifier.
		_a						: List[Any]    =			binary_head["""weight"""]
		_a						: str    =			binary_head["""bias"""]

		# It should be done!
		return output_state_dict

def     lowerCAmelCase__  (  ):



		'''simple docstring'''
		_a						: List[Any]    =			argparse.ArgumentParser()
		parser.add_argument("""--print-checkpoint-structure"""  ,     action="""store_true"""       )
		parser.add_argument("""path_to_checkpoint"""  ,     type=UpperCamelCase__  ,     help="""Path to the ZIP file containing the checkpoint"""       )
		parser.add_argument(
		    """--config_file"""  ,     default=""""""  ,     type=UpperCamelCase__  ,     help="""An optional config json file describing the pre-trained model."""  ,     )
		_a						: Union[str, Any]    =			parser.parse_args()

		# Extract the basename.
		_a						: Optional[Any]    =			os.path.dirname(args.path_to_checkpoint       )

		# Load the model.
		# the .zip is very optional, let's keep it for backward compatibility
		print(F"""Extracting PyTorch state dictionary from \"{args.path_to_checkpoint}\""""       )
		if args.path_to_checkpoint.endswith(""".zip"""       ):
				with zipfile.ZipFile(args.path_to_checkpoint  ,     """r"""       ) as checkpoint:
						with checkpoint.open("""release/mp_rank_00/model_optim_rng.pt"""       ) as pytorch_dict:
								_a						: str    =			torch.load(UpperCamelCase__  ,     map_location="""cpu"""       )
		else:
				_a						: Optional[Any]    =			torch.load(args.path_to_checkpoint  ,     map_location="""cpu"""       )

		if args.config_file == "":
				# Default config of megatron-bert 345m
				_a						: Tuple    =			MegatronBertConfig()

				# different megatron-bert-*-345m models have different vocab sizes, so override the default
				# config (which is for megatron-bert-cased-345m) with the actual vocab dimension
				_a						: int    =			input_state_dict["""model"""]["""lm_head"""]["""bias"""].numel()
		else:
				_a						: List[Any]    =			MegatronBertConfig.from_json_file(args.config_file       )

		# Convert.
		print("""Converting"""       )
		_a						: str    =			convert_megatron_checkpoint(UpperCamelCase__  ,     UpperCamelCase__  ,     UpperCamelCase__       )

		# Print the structure of converted state dict.
		if args.print_checkpoint_structure:
				recursive_print(UpperCamelCase__  ,     UpperCamelCase__       )

		# Store the config to file.
		print("""Saving config"""       )
		config.save_pretrained(UpperCamelCase__       )

		# Store the state_dict to file.
		_a						: List[str]    =			os.path.join(UpperCamelCase__  ,     """pytorch_model.bin"""       )
		print(F"""Saving checkpoint to \"{output_checkpoint_file}\""""       )
		torch.save(UpperCamelCase__  ,     UpperCamelCase__       )


####################################################################################################

if __name__ == "__main__":
	main()

####################################################################################################









"""simple docstring"""




import os
import unittest

from transformers.models.bertweet.tokenization_bertweet import VOCAB_FILES_NAMES, BertweetTokenizer

from ...test_tokenization_common import TokenizerTesterMixin






class  UpperCamelCase    (			snake_case_	,  unittest.TestCase   ):
 UpperCamelCase      : Union[str, Any]    =      BertweetTokenizer
 UpperCamelCase      : Optional[Any]    =      False


 def _lowercase			(  self	:		Dict							)      -> List[Any]:
   super().setUp()

   # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
   _a						: Tuple    =			["""I""", """m""", """V@@""", """R@@""", """r""", """e@@"""]
   _a						: str    =			dict(zip(UpperCAmelCase__							,		range(len(UpperCAmelCase__							)							)							)							)
   _a						: Dict    =			["""#version: 0.2""", """a m</w>"""]
   _a						: int    =			{"""unk_token""": """<unk>"""}

   _a						: Union[str, Any]    =			os.path.join(self.tmpdirname							,		VOCAB_FILES_NAMES["""vocab_file"""]							)
   _a						: List[Any]    =			os.path.join(self.tmpdirname							,		VOCAB_FILES_NAMES["""merges_file"""]							)
   with open(self.vocab_file							,		"""w"""							,		encoding="""utf-8"""							) as fp:
     for token in vocab_tokens:
       fp.write(f"""{token} {vocab_tokens[token]}\n"""							)
   with open(self.merges_file							,		"""w"""							,		encoding="""utf-8"""							) as fp:
     fp.write("""\n""".join(UpperCAmelCase__							)							)


 def _lowercase			(  self	:		Any							,		**UpperCAmelCase__	:		Optional[Any]							)      -> Dict:
   kwargs.update(self.special_tokens_map							)
   return BertweetTokenizer.from_pretrained(self.tmpdirname							,		**UpperCAmelCase__							)


 def _lowercase			(  self	:		Dict							,		UpperCAmelCase__	:		Optional[int]							)      -> Dict:
   _a						: List[str]    =			"""I am VinAI Research"""
   _a						: Optional[int]    =			"""I <unk> m V<unk> <unk> <unk> I Re<unk> e<unk> <unk> <unk> <unk>"""
   return input_text, output_text







 def _lowercase			(  self	:		Tuple							)      -> Optional[int]:
   _a						: Optional[int]    =			BertweetTokenizer(self.vocab_file							,		self.merges_file							,		**self.special_tokens_map							)
   _a						: Tuple    =			"""I am VinAI Research"""
   _a						: Dict    =			"""I a@@ m V@@ i@@ n@@ A@@ I R@@ e@@ s@@ e@@ a@@ r@@ c@@ h""".split()
   _a						: str    =			tokenizer.tokenize(UpperCAmelCase__							)
   self.assertListEqual(UpperCAmelCase__							,		UpperCAmelCase__							)

   _a						: Tuple    =			tokens + [tokenizer.unk_token]

   _a						: Tuple    =			[4, 3, 5, 6, 3, 3, 3, 4, 7, 9, 3, 9, 3, 3, 3, 3, 3]
   self.assertListEqual(tokenizer.convert_tokens_to_ids(UpperCAmelCase__							)							,		UpperCAmelCase__							)

