




"""simple docstring"""




import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase=0	):
			# Format the message.
			if name is None:
						__lowercase			:  Optional[int]								=       None
			else:
						__lowercase			:  Optional[int]								=       '''.''' * max(0	,				spaces - 2	) + '''# {:''' + str(50 - spaces	) + '''s}'''
						__lowercase			:  str								=       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      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	):
			# Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
			# for compatibility with later versions of NVIDIA Megatron-LM.
			# The inverse operation is performed inside Megatron-LM to read checkpoints:
			# https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
			# If param is the weight tensor of the self-attention block, the returned tensor
			# will have to be transposed one more time to be read by HuggingFace GPT2.
			__lowercase			:  Optional[int]								=       param.size()
			if checkpoint_version == 1.0:
						# version 1.0 stores [num_heads * hidden_size * num_splits, :]
						__lowercase			:  int								=       (num_heads, hidden_size, num_splits) + input_shape[1:]
						__lowercase			:  str								=       param.view(*UpperCamelCase_	)
						__lowercase			:  Union[str, Any]								=       param.transpose(0	,				2	)
						__lowercase			:  Optional[int]								=       param.transpose(1	,				2	).contiguous()
			elif checkpoint_version >= 2.0:
						# other versions store [num_heads * num_splits * hidden_size, :]
						__lowercase			:  Optional[Any]								=       (num_heads, num_splits, hidden_size) + input_shape[1:]
						__lowercase			:  List[str]								=       param.view(*UpperCamelCase_	)
						__lowercase			:  Union[str, Any]								=       param.transpose(0	,				1	).contiguous()
			__lowercase			:  Tuple								=       param.view(*UpperCamelCase_	)
			return param







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	):
			# The converted output model.
			__lowercase			:  str								=       {}

			# old versions did not store training args
			__lowercase			:  Optional[int]								=       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))

						__lowercase			:  List[Any]								=       ds_args.padded_vocab_size
						__lowercase			:  Optional[Any]								=       ds_args.max_position_embeddings
						__lowercase			:  List[str]								=       ds_args.hidden_size
						__lowercase			:  Optional[Any]								=       ds_args.num_layers
						__lowercase			:  str								=       ds_args.num_attention_heads
						__lowercase			:  Any								=       ds_args.ffn_hidden_size
						# pprint(config)

			# The number of heads.
			__lowercase			:  Optional[Any]								=       config.n_head
			# The hidden_size per head.
			__lowercase			:  Any								=       config.n_embd // config.n_head
			# Megatron-LM checkpoint version
			if "checkpoint_version" in input_state_dict.keys():
						__lowercase			:  Dict								=       input_state_dict['''checkpoint_version''']
			else:
						__lowercase			:  Any								=       0.0

			# The model.
			__lowercase			:  str								=       input_state_dict['''model''']
			# The language model.
			__lowercase			:  Optional[Any]								=       model['''language_model''']
			# The embeddings.
			__lowercase			:  Optional[Any]								=       lm['''embedding''']

			# The word embeddings.
			__lowercase			:  List[Any]								=       embeddings['''word_embeddings''']['''weight''']
			# Truncate the embedding table to vocab_size rows.
			__lowercase			:  Optional[Any]								=       word_embeddings[: config.vocab_size, :]
			__lowercase			:  str								=       word_embeddings

			# The position embeddings.
			__lowercase			:  Any								=       embeddings['''position_embeddings''']['''weight''']
			# Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
			__lowercase			:  Tuple								=       pos_embeddings.size(0	)
			if n_positions != config.n_positions:
						raise ValueError(
						    f"""pos_embeddings.max_sequence_length={n_positions} and config.n_positions={config.n_positions} don't match"""	)
			# Store the position embeddings.
			__lowercase			:  List[Any]								=       pos_embeddings

			# The transformer.
			__lowercase			:  Tuple								=       lm['''transformer'''] if '''transformer''' in lm.keys() else lm['''encoder''']

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

			# The simple map of names for "automated" rules.
			__lowercase			:  Optional[int]								=       {
			    '''attention.dense''': '''.attn.c_proj.''',
			    '''self_attention.dense''': '''.attn.c_proj.''',
			    '''mlp.dense_h_to_4h''': '''.mlp.c_fc.''',
			    '''mlp.dense_4h_to_h''': '''.mlp.c_proj.''',
			}

			# Extract the layers.
			for key, val in transformer.items():
						# Match the name.
						__lowercase			:  List[str]								=       layer_re.match(UpperCamelCase_	)

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

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

						# The name of the layer.
						__lowercase			:  Dict								=       f"""transformer.h.{layer_idx}"""

						# For layernorm(s), simply store the layer norm.
						if op_name.endswith('''layernorm'''	):
									__lowercase			:  int								=       '''ln_1''' if op_name.startswith('''input'''	) else '''ln_2'''
									__lowercase			:  int								=       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":
									# Insert a tensor of 1x1xDxD bias.
									__lowercase			:  Dict								=       torch.tril(torch.ones((n_positions, n_positions)	,				dtype=torch.floataa	)	).view(
									    1	,				1	,				UpperCamelCase_	,				UpperCamelCase_	)
									__lowercase			:  Optional[Any]								=       causal_mask

									# Insert a "dummy" tensor for masked_bias.
									__lowercase			:  str								=       torch.tensor(-1e4	,				dtype=torch.floataa	)
									__lowercase			:  Optional[Any]								=       masked_bias

									__lowercase			:  Optional[int]								=       fix_query_key_value_ordering(UpperCamelCase_	,				UpperCamelCase_	,				3	,				UpperCamelCase_	,				UpperCamelCase_	)
									# Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
									__lowercase			:  Tuple								=       out_val.transpose(0	,				1	).contiguous()
									# Store.
									__lowercase			:  Dict								=       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":
									__lowercase			:  List[str]								=       fix_query_key_value_ordering(UpperCamelCase_	,				UpperCamelCase_	,				3	,				UpperCamelCase_	,				UpperCamelCase_	)
									# Store. No change of shape.
									__lowercase			:  Any								=       out_val

						# Transpose the weights.
						elif weight_or_bias == "weight":
									__lowercase			:  Any								=       megatron_to_transformers[op_name]
									__lowercase			:  Tuple								=       val.transpose(0	,				1	)

						# Copy the bias.
						elif weight_or_bias == "bias":
									__lowercase			:  List[Any]								=       megatron_to_transformers[op_name]
									__lowercase			:  str								=       val

    # DEBUG.
			assert config.n_layer == layer_idx + 1

			# The final layernorm.
			__lowercase			:  Dict								=       transformer['''final_layernorm.weight''']
			__lowercase			:  Dict								=       transformer['''final_layernorm.bias''']

			# For LM head, transformers' wants the matrix to weight embeddings.
			__lowercase			:  Any								=       word_embeddings

			# It should be done!
			return output_state_dict







def      __UpperCAmelCase     ( ):
			# Create the argument parser.
			__lowercase			:  List[str]								=       argparse.ArgumentParser()
			parser.add_argument('''--print-checkpoint-structure'''	,				action='''store_true'''	)
			parser.add_argument(
			    '''path_to_checkpoint'''	,				type=UpperCamelCase_	,				help='''Path to the checkpoint file (.zip archive or direct .pt file)'''	,				)
			parser.add_argument(
			    '''--config_file'''	,				default=''''''	,				type=UpperCamelCase_	,				help='''An optional config json file describing the pre-trained model.'''	,				)
			__lowercase			:  Dict								=       parser.parse_args()

			# Extract the basename.
			__lowercase			:  Union[str, 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:
												__lowercase			:  Tuple								=       torch.load(UpperCamelCase_	,				map_location='''cpu'''	)
			else:
						__lowercase			:  Dict								=       torch.load(args.path_to_checkpoint	,				map_location='''cpu'''	)

			__lowercase			:  Dict								=       input_state_dict.get('''args'''	,				UpperCamelCase_	)

			# Read the config, or default to the model released by NVIDIA.
			if args.config_file == "":
						if ds_args is not None:
									if ds_args.bias_gelu_fusion:
												__lowercase			:  str								=       '''gelu_fast'''
									elif ds_args.openai_gelu:
												__lowercase			:  Dict								=       '''gelu_new'''
									else:
												__lowercase			:  List[str]								=       '''gelu'''
						else:
									# in the very early days this used to be "gelu_new"
									__lowercase			:  List[str]								=       '''gelu_new'''

						# Spell out all parameters in case the defaults change.
						__lowercase			:  Any								=       GPTaConfig(
						    vocab_size=5_02_57	,				n_positions=10_24	,				n_embd=10_24	,				n_layer=24	,				n_head=16	,				n_inner=40_96	,				activation_function=UpperCamelCase_	,				resid_pdrop=0.1	,				embd_pdrop=0.1	,				attn_pdrop=0.1	,				layer_norm_epsilon=1e-5	,				initializer_range=0.02	,				summary_type='''cls_index'''	,				summary_use_proj=UpperCamelCase_	,				summary_activation=UpperCamelCase_	,				summary_proj_to_labels=UpperCamelCase_	,				summary_first_dropout=0.1	,				scale_attn_weights=UpperCamelCase_	,				use_cache=UpperCamelCase_	,				bos_token_id=5_02_56	,				eos_token_id=5_02_56	,				)
			else:
						__lowercase			:  List[Any]								=       GPTaConfig.from_json_file(args.config_file	)

			__lowercase			:  Dict								=       ['''GPT2LMHeadModel''']

			# Convert.
			print('''Converting'''	)
			__lowercase			:  Tuple								=       convert_megatron_checkpoint(UpperCamelCase_	,				UpperCamelCase_	,				UpperCamelCase_	)

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

			# Add tokenizer class info to config
			# see https://github.com/huggingface/transformers/issues/13906)
			if ds_args is not None:
						__lowercase			:  Optional[int]								=       ds_args.tokenizer_type
						if tokenizer_type == "GPT2BPETokenizer":
									__lowercase			:  Dict								=       '''gpt2'''
						elif tokenizer_type == "PretrainedFromHF":
									__lowercase			:  Dict								=       ds_args.tokenizer_name_or_path
						else:
									raise ValueError(f"""Unrecognized tokenizer_type {tokenizer_type}"""	)
			else:
						__lowercase			:  int								=       '''gpt2'''

			__lowercase			:  List[str]								=       AutoTokenizer.from_pretrained(UpperCamelCase_	)
			__lowercase			:  List[str]								=       type(UpperCamelCase_	).__name__
			__lowercase			:  Optional[int]								=       tokenizer_class

			# Store the config to file.
			print('''Saving config'''	)
			config.save_pretrained(UpperCamelCase_	)

			# Save tokenizer based on args
			print(f"""Adding {tokenizer_class} tokenizer files"""	)
			tokenizer.save_pretrained(UpperCamelCase_	)

			# Store the state_dict to file.
			__lowercase			:  Optional[Any]								=       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'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
							from ..ta.tokenization_ta import TaTokenizer
else:
							from ...utils.dummy_sentencepiece_objects import TaTokenizer

__SCREAMING_SNAKE_CASE					:				str									=		TaTokenizer

if is_tokenizers_available():
							from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
							from ...utils.dummy_tokenizers_objects import TaTokenizerFast

__SCREAMING_SNAKE_CASE					:				Any									=		TaTokenizerFast

__SCREAMING_SNAKE_CASE					:				Tuple									=		{'''configuration_mt5''': ['''MT5Config''', '''MT5OnnxConfig''']}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__SCREAMING_SNAKE_CASE					:				List[Any]									=		[
							    '''MT5EncoderModel''',
							    '''MT5ForConditionalGeneration''',
							    '''MT5ForQuestionAnswering''',
							    '''MT5Model''',
							    '''MT5PreTrainedModel''',
							    '''MT5Stack''',
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__SCREAMING_SNAKE_CASE					:				str									=		['''TFMT5EncoderModel''', '''TFMT5ForConditionalGeneration''', '''TFMT5Model''']

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__SCREAMING_SNAKE_CASE					:				List[str]									=		['''FlaxMT5EncoderModel''', '''FlaxMT5ForConditionalGeneration''', '''FlaxMT5Model''']


if TYPE_CHECKING:
							from .configuration_mta import MTaConfig, MTaOnnxConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_mta import (
														    MTaEncoderModel,
														    MTaForConditionalGeneration,
														    MTaForQuestionAnswering,
														    MTaModel,
														    MTaPreTrainedModel,
														    MTaStack,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
							import sys

							__SCREAMING_SNAKE_CASE					:				Union[str, Any]									=		_LazyModule(
							    __name__,
							    globals()['''__file__'''],
							    _import_structure,
							    extra_objects={'''MT5Tokenizer''': MTaTokenizer, '''MT5TokenizerFast''': MTaTokenizerFast},
							    module_spec=__spec__,
							)





