





"""simple docstring"""





# Copyright (c) 2021-, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE=0    ):




		"""simple docstring"""


		if name is None:
				UpperCamelCase   =	None
		else:
				UpperCamelCase   =	"." * max(0  ,	spaces - 2    ) + "# {:" + str(50 - spaces    ) + "s}"
				UpperCamelCase   =	fmt.format(_SCREAMING_SNAKE_CASE    )

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


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	param.size()
		if checkpoint_version == 1.0:
				# version 1.0 stores [num_heads * hidden_size * num_splits, :]
				UpperCamelCase   =	(num_heads, hidden_size, num_splits) + input_shape[1:]
				UpperCamelCase   =	param.view(*_SCREAMING_SNAKE_CASE    )
				UpperCamelCase   =	param.transpose(0  ,	2    )
				UpperCamelCase   =	param.transpose(1  ,	2    ).contiguous()
		elif checkpoint_version >= 2.0:
				# other versions store [num_heads * num_splits * hidden_size, :]
				UpperCamelCase   =	(num_heads, num_splits, hidden_size) + input_shape[1:]
				UpperCamelCase   =	param.view(*_SCREAMING_SNAKE_CASE    )
				UpperCamelCase   =	param.transpose(0  ,	1    ).contiguous()
		UpperCamelCase   =	param.view(*_SCREAMING_SNAKE_CASE    )
		return param


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	{}

		# old versions did not store training args
		UpperCamelCase   =	input_state_dict.get("args"  ,	_SCREAMING_SNAKE_CASE    )
		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))

				UpperCamelCase   =	ds_args.padded_vocab_size
				UpperCamelCase   =	ds_args.max_position_embeddings
				UpperCamelCase   =	ds_args.hidden_size
				UpperCamelCase   =	ds_args.num_layers
				UpperCamelCase   =	ds_args.num_attention_heads
				UpperCamelCase   =	ds_args.ffn_hidden_size
				# pprint(config)

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

		# The model.
		UpperCamelCase   =	input_state_dict["model"]
		# The language model.
		UpperCamelCase   =	model["language_model"]
		# The embeddings.
		UpperCamelCase   =	lm["embedding"]

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

		# The position embeddings.
		UpperCamelCase   =	embeddings["position_embeddings"]["weight"]
		# Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
		UpperCamelCase   =	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.
		UpperCamelCase   =	pos_embeddings

		# The transformer.
		UpperCamelCase   =	lm["transformer"] if "transformer" in lm.keys() else lm["encoder"]

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

		# The simple map of names for "automated" rules.
		UpperCamelCase   =	{
		    "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.
				UpperCamelCase   =	layer_re.match(_SCREAMING_SNAKE_CASE    )

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

				# The index of the layer.
				UpperCamelCase   =	int(m.group(1    )    )
				# The name of the operation.
				UpperCamelCase   =	m.group(2    )
				# Is it a weight or a bias?
				UpperCamelCase   =	m.group(3    )

				# The name of the layer.
				UpperCamelCase   =	F"transformer.h.{layer_idx}"

				# For layernorm(s), simply store the layer norm.
				if op_name.endswith("layernorm"    ):
						UpperCamelCase   =	"ln_1" if op_name.startswith("input"    ) else "ln_2"
						UpperCamelCase   =	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.
						UpperCamelCase   =	torch.tril(torch.ones((n_positions, n_positions)  ,	dtype=torch.floataa    )    ).view(
						    1  ,	1  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
						UpperCamelCase   =	causal_mask

						# Insert a "dummy" tensor for masked_bias.
						UpperCamelCase   =	torch.tensor(-1e4  ,	dtype=torch.floataa    )
						UpperCamelCase   =	masked_bias

						UpperCamelCase   =	fix_query_key_value_ordering(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	3  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
						# Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
						UpperCamelCase   =	out_val.transpose(0  ,	1    ).contiguous()
						# Store.
						UpperCamelCase   =	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":
						UpperCamelCase   =	fix_query_key_value_ordering(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	3  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )
						# Store. No change of shape.
						UpperCamelCase   =	out_val

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

				# Copy the bias.
				elif weight_or_bias == "bias":
						UpperCamelCase   =	megatron_to_transformers[op_name]
						UpperCamelCase   =	val

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

		# The final layernorm.
		UpperCamelCase   =	transformer["final_layernorm.weight"]
		UpperCamelCase   =	transformer["final_layernorm.bias"]

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

		# It should be done!
		return output_state_dict


def       a__					(					):




		"""simple docstring"""


		UpperCamelCase   =	argparse.ArgumentParser()
		parser.add_argument("--print-checkpoint-structure"  ,	action="store_true"    )
		parser.add_argument(
		    "path_to_checkpoint"  ,	type=_SCREAMING_SNAKE_CASE  ,	help="Path to the checkpoint file (.zip archive or direct .pt file)"  ,	)
		parser.add_argument(
		    "--config_file"  ,	default=""  ,	type=_SCREAMING_SNAKE_CASE  ,	help="An optional config json file describing the pre-trained model."  ,	)
		UpperCamelCase   =	parser.parse_args()

		# Extract the basename.
		UpperCamelCase   =	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:
								UpperCamelCase   =	torch.load(_SCREAMING_SNAKE_CASE  ,	map_location="cpu"    )
		else:
				UpperCamelCase   =	torch.load(args.path_to_checkpoint  ,	map_location="cpu"    )

		UpperCamelCase   =	input_state_dict.get("args"  ,	_SCREAMING_SNAKE_CASE    )

		# 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:
								UpperCamelCase   =	"gelu_fast"
						elif ds_args.openai_gelu:
								UpperCamelCase   =	"gelu_new"
						else:
								UpperCamelCase   =	"gelu"
				else:
						# in the very early days this used to be "gelu_new"
						UpperCamelCase   =	"gelu_new"

				# Spell out all parameters in case the defaults change.
				UpperCamelCase   =	GPTaConfig(
				    vocab_size=50_257  ,	n_positions=1_024  ,	n_embd=1_024  ,	n_layer=24  ,	n_head=16  ,	n_inner=4_096  ,	activation_function=_SCREAMING_SNAKE_CASE  ,	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=_SCREAMING_SNAKE_CASE  ,	summary_activation=_SCREAMING_SNAKE_CASE  ,	summary_proj_to_labels=_SCREAMING_SNAKE_CASE  ,	summary_first_dropout=0.1  ,	scale_attn_weights=_SCREAMING_SNAKE_CASE  ,	use_cache=_SCREAMING_SNAKE_CASE  ,	bos_token_id=50_256  ,	eos_token_id=50_256  ,	)
		else:
				UpperCamelCase   =	GPTaConfig.from_json_file(args.config_file    )

		UpperCamelCase   =	["GPT2LMHeadModel"]

		# Convert.
		print("Converting"    )
		UpperCamelCase   =	convert_megatron_checkpoint(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

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

		# Add tokenizer class info to config
		# see https://github.com/huggingface/transformers/issues/13906)
		if ds_args is not None:
				UpperCamelCase   =	ds_args.tokenizer_type
				if tokenizer_type == "GPT2BPETokenizer":
						UpperCamelCase   =	"gpt2"
				elif tokenizer_type == "PretrainedFromHF":
						UpperCamelCase   =	ds_args.tokenizer_name_or_path
				else:
						raise ValueError(F"Unrecognized tokenizer_type {tokenizer_type}"    )
		else:
				UpperCamelCase   =	"gpt2"

		UpperCamelCase   =	AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	type(_SCREAMING_SNAKE_CASE    ).__name__
		UpperCamelCase   =	tokenizer_class

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

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

		# Store the state_dict to file.
		UpperCamelCase   =	os.path.join(_SCREAMING_SNAKE_CASE  ,	"pytorch_model.bin"    )
		print(F"Saving checkpoint to \"{output_checkpoint_file}\""    )
		torch.save(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )


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

if __name__ == "__main__":
						main()

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













"""simple docstring"""





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
						from transformers import AutoModelForSeqaSeqLM, AutoTokenizer





@require_torch
@require_sentencepiece
@require_tokenizers
class        _lowerCamelCase      (       unittest.TestCase ):


						@slow
						def     snake_case_   (self		) ->   List[str]:
								UpperCamelCase   =	AutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"   ,							return_dict=__a		).to(__a		)
								UpperCamelCase   =	AutoTokenizer.from_pretrained("google/mt5-small"		)

								UpperCamelCase   =	tokenizer("Hello there"   ,							return_tensors="pt"		).input_ids
								UpperCamelCase   =	tokenizer("Hi I am"   ,							return_tensors="pt"		).input_ids

								UpperCamelCase   =	model(input_ids.to(__a		)   ,							labels=labels.to(__a		)		).loss
								UpperCamelCase   =	-(labels.shape[-1] * loss.item())

								UpperCamelCase   =	-84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE		) < 1e-4		)






