


import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig






def   __lowerCamelCase	(       lowerCamelCase__		:	Optional[Any]     ,       lowerCamelCase__		:	Optional[int]     ,       lowerCamelCase__		:	Optional[Any]=0			):





			'''simple docstring'''




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

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






def   __lowerCamelCase	(       lowerCamelCase__		:	Optional[int]     ,       lowerCamelCase__		:	Dict     ,       lowerCamelCase__		:	Tuple     ,       lowerCamelCase__		:	List[str]			):





			'''simple docstring'''




			lowerCamelCase								=     param.size()
			# other versions store [num_heads * num_splits * hidden_size, :]
			lowerCamelCase								=     (num_heads, num_splits, hidden_size) + input_shape[1:]
			lowerCamelCase								=     param.view(*lowerCamelCase__			)
			lowerCamelCase								=     param.transpose(0     ,       1			).contiguous()
			lowerCamelCase								=     param.view(*lowerCamelCase__			)
			return param






def   __lowerCamelCase	(       lowerCamelCase__		:	Union[str, Any]     ,       lowerCamelCase__		:	List[str]			):





			'''simple docstring'''




			lowerCamelCase								=     config.n_positions
			lowerCamelCase								=     config.n_layer
			lowerCamelCase								=     config.vocab_size
			lowerCamelCase								=     config.n_head
			lowerCamelCase								=     config.n_embd // config.n_head

			lowerCamelCase								=     sd_megatron["""model.language_model.embedding.word_embeddings.weight"""][:vocab_size, :]
			lowerCamelCase								=     {
			    """transformer.wte.weight""": word_embeddings,
			    """transformer.wpe.weight""": sd_megatron["""model.language_model.embedding.position_embeddings.weight"""],
			    """transformer.ln_f.weight""": sd_megatron["""model.language_model.encoder.final_layernorm.weight"""],
			    """transformer.ln_f.bias""": sd_megatron["""model.language_model.encoder.final_layernorm.bias"""],
			}

			lowerCamelCase								=     """model.language_model.encoder.layers."""
			for i in range(lowerCamelCase__			):
						lowerCamelCase								=     torch.tril(torch.ones((n_positions, n_positions)     ,       dtype=torch.bool			)			)
						lowerCamelCase								=     causal_mask.view(1     ,       1     ,       lowerCamelCase__     ,       lowerCamelCase__			)
						lowerCamelCase								=     causal_mask
						lowerCamelCase								=     torch.tensor(-1E4     ,       dtype=torch.bfloataa			)

						lowerCamelCase								=     sd_megatron[f'{pf}{i}.input_layernorm.weight']
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.input_layernorm.bias']

						lowerCamelCase								=     sd_megatron[f'{pf}{i}.self_attention.query_key_value.weight']
						lowerCamelCase								=     fix_query_key_value_ordering(lowerCamelCase__     ,       3     ,       lowerCamelCase__     ,       lowerCamelCase__			)
						lowerCamelCase								=     vala.transpose(0     ,       1			).contiguous()

						lowerCamelCase								=     sd_megatron[f'{pf}{i}.self_attention.query_key_value.bias']
						lowerCamelCase								=     fix_query_key_value_ordering(lowerCamelCase__     ,       3     ,       lowerCamelCase__     ,       lowerCamelCase__			)
						lowerCamelCase								=     vala

						lowerCamelCase								=     sd_megatron[f'{pf}{i}.self_attention.dense.weight'].transpose(
						    0     ,       1			)
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.self_attention.dense.bias']
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.post_attention_layernorm.weight']
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.post_attention_layernorm.bias']
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.mlp.dense_h_to_4h.weight'].transpose(0     ,       1			)
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.mlp.dense_h_to_4h.bias']
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.mlp.dense_4h_to_h.weight'].transpose(
						    0     ,       1			)
						lowerCamelCase								=     sd_megatron[f'{pf}{i}.mlp.dense_4h_to_h.bias']

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

			return sd_hf






def   __lowerCamelCase	(       lowerCamelCase__		:	Optional[Any]     ,       lowerCamelCase__		:	Tuple			):





			'''simple docstring'''




			lowerCamelCase								=     64000
			lowerCamelCase								=     config_megatron["""encoder_seq_length"""]
			lowerCamelCase								=     config_megatron["""hidden_size"""]
			lowerCamelCase								=     config_megatron["""num_layers"""]
			lowerCamelCase								=     config_megatron["""num_attention_heads"""]
			lowerCamelCase								=     config_megatron["""ffn_hidden_size"""]
			lowerCamelCase								=     """gelu"""
			lowerCamelCase								=     0.1
			lowerCamelCase								=     0.1
			lowerCamelCase								=     0.1
			lowerCamelCase								=     config_megatron["""layernorm_epsilon"""]  # 1e-5
			lowerCamelCase								=     config_megatron["""init_method_std"""]  # 0.02
			lowerCamelCase								=     config_megatron["""apply_query_key_layer_scaling"""]  # True
			lowerCamelCase								=     True
			lowerCamelCase								=     True

			# This identifies the 6.7B (7B) model which uses a different tokenizer
			if config_megatron["hidden_size"] == 4096:
						lowerCamelCase								=     1  # <|endoftext|>
						lowerCamelCase								=     1  # <|endoftext|>
						lowerCamelCase								=     0  # <unk>
			else:
						lowerCamelCase								=     2  # <s>
						lowerCamelCase								=     3  # <|endoftext|>
						lowerCamelCase								=     0  # <pad>

			return config_hf






def   __lowerCamelCase	(       lowerCamelCase__		:	int			):





			'''simple docstring'''




			print(lowerCamelCase__			)

			lowerCamelCase								=     args.checkpoint_path
			lowerCamelCase								=     args.save_path
			if isfile(lowerCamelCase__			):
						raise FileNotFoundError(f'ERROR! could not find file {checkpoint_path}'			)

			# Load the model.
			lowerCamelCase								=     torch.load(lowerCamelCase__     ,       map_location="""cpu"""			)

			# Load the config.
			lowerCamelCase								=     checkpoint["""hyper_parameters"""]["""cfg"""]
			lowerCamelCase								=     GPTaConfig()
			lowerCamelCase								=     copy_config(config_hf=lowerCamelCase__     ,       config_megatron=lowerCamelCase__			)
			lowerCamelCase								=     ["""GPT2LMHeadModel"""]

			lowerCamelCase								=     checkpoint["""state_dict"""]

			# Convert.
			print("""Converting"""			)
			lowerCamelCase								=     convert_megatron_checkpoint(lowerCamelCase__     ,       lowerCamelCase__			)

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

			lowerCamelCase								=     """GPTSw3Tokenizer"""

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

			# Store the state_dict to file.
			lowerCamelCase								=     os.path.join(lowerCamelCase__     ,       """pytorch_model.bin"""			)
			print(f'Saving checkpoint to "{output_checkpoint_file}"'			)
			torch.save(lowerCamelCase__     ,       lowerCamelCase__			)


if __name__ == "__main__":
		UpperCAmelCase      :						str    =		argparse.ArgumentParser()
		parser.add_argument(
		    "--checkpoint_path",
		    type=str,
		    required=True,
		    help="e.g. megatron_gpt--val_loss=2.42-step=38000-consumed_samples=54720000",
		)
		parser.add_argument("--save_path", type=str, required=True, help="e.g. /home/user/gpt-sw3/hf")
		parser.add_argument("--print-checkpoint-structure", action="store_true")
		UpperCAmelCase      :						Optional[Any]    =		parser.parse_args()
		main(_args)





import inspect
import unittest

from transformers import EfficientNetConfig
from transformers.testing_utils import is_pipeline_test, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import EfficientNetForImageClassification, EfficientNetModel
		from transformers.models.efficientnet.modeling_efficientnet import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor







class 					__lowercase					:



		"""simple docstring"""

		def __init__(   self		,			A		,			A=13		,			A=32		,			A=3		,			A=[3, 3, 5]		,			A=[32, 16, 24]		,			A=[16, 24, 40]		,			A=[1, 1, 2]		,			A=[1, 1, 2]		,			A=[1, 6, 6]		,			A=True		,			A=True		,			A=37		,			A="gelu"		,			A=10		,			)     ->   Dict:







					'''simple docstring'''







					lowerCamelCase								=     parent
					lowerCamelCase								=     batch_size
					lowerCamelCase								=     image_size
					lowerCamelCase								=     num_channels
					lowerCamelCase								=     kernel_sizes
					lowerCamelCase								=     in_channels
					lowerCamelCase								=     out_channels
					lowerCamelCase								=     strides
					lowerCamelCase								=     num_block_repeats
					lowerCamelCase								=     expand_ratios
					lowerCamelCase								=     is_training
					lowerCamelCase								=     hidden_act
					lowerCamelCase								=     num_labels
					lowerCamelCase								=     use_labels







		def 				__A      (   self  )     ->   Optional[Any]:







					'''simple docstring'''







					lowerCamelCase								=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

					lowerCamelCase								=     None
					if self.use_labels:
								lowerCamelCase								=     ids_tensor([self.batch_size]		,			self.num_labels  )

					lowerCamelCase								=     self.get_config()
					return config, pixel_values, labels







		def 				__A      (   self  )     ->   Optional[int]:







					'''simple docstring'''







					return EfficientNetConfig(
					    num_channels=self.num_channels		,			kernel_sizes=self.kernel_sizes		,			in_channels=self.in_channels		,			out_channels=self.out_channels		,			strides=self.strides		,			num_block_repeats=self.num_block_repeats		,			expand_ratios=self.expand_ratios		,			hidden_act=self.hidden_act		,			num_labels=self.num_labels		,			)







		def 				__A      (   self		,			A		,			A		,			A  )     ->   List[Any]:







					'''simple docstring'''







					lowerCamelCase								=     EfficientNetModel(config=A  )
					model.to(A  )
					model.eval()
					lowerCamelCase								=     model(A  )
					# expected last hidden states: B, C, H // 4, W // 4
					self.parent.assertEqual(
					    result.last_hidden_state.shape		,			(self.batch_size, config.hidden_dim, self.image_size // 4, self.image_size // 4)		,			)







		def 				__A      (   self		,			A		,			A		,			A  )     ->   Optional[int]:







					'''simple docstring'''







					lowerCamelCase								=     EfficientNetForImageClassification(A  )
					model.to(A  )
					model.eval()
					lowerCamelCase								=     model(A		,			labels=A  )
					self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_labels)  )







		def 				__A      (   self  )     ->   Any:







					'''simple docstring'''







					lowerCamelCase								=     self.prepare_config_and_inputs()
					lowerCamelCase		,		lowerCamelCase		,		lowerCamelCase								=     config_and_inputs
					lowerCamelCase								=     {"""pixel_values""": pixel_values}
					return config, inputs_dict







@require_torch
class 					__lowercase					(       a_ ,       a_ ,       unittest.TestCase     ):



		"""simple docstring"""

		UpperCamelCase    :				str 					=		(EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
		UpperCamelCase    :				List[Any] 					=		(
		    {"feature-extraction": EfficientNetModel, "image-classification": EfficientNetForImageClassification}
		    if is_torch_available()
		    else {}
		)

		UpperCamelCase    :				Any 					=		False
		UpperCamelCase    :				Dict 					=		False
		UpperCamelCase    :				Optional[int] 					=		False
		UpperCamelCase    :				str 					=		False
		UpperCamelCase    :				Optional[Any] 					=		False







		def 				__A      (   self  )     ->   str:







					'''simple docstring'''







					lowerCamelCase								=     EfficientNetModelTester(self  )
					lowerCamelCase								=     ConfigTester(
					    self		,			config_class=A		,			has_text_modality=A		,			hidden_size=37  )







		def 				__A      (   self  )     ->   Optional[Any]:







					'''simple docstring'''







					self.create_and_test_config_common_properties()
					self.config_tester.create_and_test_config_to_json_string()
					self.config_tester.create_and_test_config_to_json_file()
					self.config_tester.create_and_test_config_from_and_save_pretrained()
					self.config_tester.create_and_test_config_with_num_labels()
					self.config_tester.check_config_can_be_init_without_params()
					self.config_tester.check_config_arguments_init()







		def 				__A      (   self  )     ->   Union[str, Any]:







					'''simple docstring'''







					return







		@unittest.skip(reason="""EfficientNet does not use inputs_embeds"""  )
		def 				__A      (   self  )     ->   Optional[Any]:







					'''simple docstring'''







					pass







		@unittest.skip(reason="""EfficientNet does not support input and output embeddings"""  )
		def 				__A      (   self  )     ->   Any:







					'''simple docstring'''







					pass







		@unittest.skip(reason="""EfficientNet does not use feedforward chunking"""  )
		def 				__A      (   self  )     ->   List[Any]:







					'''simple docstring'''







					pass







		def 				__A      (   self  )     ->   Optional[int]:







					'''simple docstring'''







					lowerCamelCase		,		lowerCamelCase								=     self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
								lowerCamelCase								=     model_class(A  )
								lowerCamelCase								=     inspect.signature(model.forward  )
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								lowerCamelCase								=     [*signature.parameters.keys()]

								lowerCamelCase								=     ["""pixel_values"""]
								self.assertListEqual(arg_names[:1]		,			A  )







		def 				__A      (   self  )     ->   Dict:







					'''simple docstring'''







					lowerCamelCase								=     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*A  )







		def 				__A      (   self  )     ->   Tuple:







					'''simple docstring'''







					def check_hidden_states_output(A		,			A		,			A  ):
								lowerCamelCase								=     model_class(A  )
								model.to(A  )
								model.eval()

								with torch.no_grad():
											lowerCamelCase								=     model(**self._prepare_for_class(A		,			A  )  )

								lowerCamelCase								=     outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
								lowerCamelCase								=     sum(config.num_block_repeats  ) * 4
								self.assertEqual(len(A  )		,			A  )

								# EfficientNet's feature maps are of shape (batch_size, num_channels, height, width)
								self.assertListEqual(
								    list(hidden_states[0].shape[-2:]  )		,			[self.model_tester.image_size // 2, self.model_tester.image_size // 2]		,			)

					lowerCamelCase		,		lowerCamelCase								=     self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
								lowerCamelCase								=     True
								check_hidden_states_output(A		,			A		,			A  )

								# check that output_hidden_states also work using config
								del inputs_dict["output_hidden_states"]
								lowerCamelCase								=     True

								check_hidden_states_output(A		,			A		,			A  )







		def 				__A      (   self  )     ->   int:







					'''simple docstring'''







					lowerCamelCase								=     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*A  )







		@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
		def 				__A      (   self  )     ->   Union[str, Any]:







					'''simple docstring'''







					pass







		@slow
		def 				__A      (   self  )     ->   Any:







					'''simple docstring'''







					for model_name in EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								lowerCamelCase								=     EfficientNetModel.from_pretrained(A  )
								self.assertIsNotNone(A  )







		@is_pipeline_test
		@require_vision
		@slow
		def 				__A      (   self  )     ->   Optional[Any]:







					'''simple docstring'''







					super().test_pipeline_image_classification()






def   __lowerCamelCase	(       ):





			'''simple docstring'''




			lowerCamelCase								=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
			return image







@require_torch
@require_vision
class 					__lowercase					(       unittest.TestCase     ):



		"""simple docstring"""

		@cached_property
		def 				__A      (   self  )     ->   Dict:







					'''simple docstring'''







					return AutoImageProcessor.from_pretrained("""google/efficientnet-b7"""  ) if is_vision_available() else None







		@slow
		def 				__A      (   self  )     ->   List[str]:







					'''simple docstring'''







					lowerCamelCase								=     EfficientNetForImageClassification.from_pretrained("""google/efficientnet-b7"""  ).to(A  )

					lowerCamelCase								=     self.default_image_processor
					lowerCamelCase								=     prepare_img()
					lowerCamelCase								=     image_processor(images=A		,			return_tensors="""pt"""  ).to(A  )

					# forward pass
					with torch.no_grad():
								lowerCamelCase								=     model(**A  )

					# verify the logits
					lowerCamelCase								=     torch.Size((1, 10_00)  )
					self.assertEqual(outputs.logits.shape		,			A  )

					lowerCamelCase								=     torch.tensor([0.0001, 0.0002, 0.0002]  ).to(A  )
					self.assertTrue(torch.allclose(outputs.logits[0, :3]		,			A		,			atol=1e-4  )  )

