


'''simple docstring'''






import argparse
import os
import re

import numpy as np
import PIL
import torch
from timm import create_model
from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator





def      UpperCAmelCase      (		lowerCamelCase_  :List[Any] ):







		'''simple docstring'''
		snake_case_		:							Union[str, Any]			       =						fname.split(os.path.sep )[-1]
		return re.search(R"""^(.*)_\d+\.jpg$"""				,      lowerCAmelCase__ ).groups()[0]
class 						__UpperCamelCase	(						lowercase__							):





							def __init__(  self		:Tuple  ,_UpperCamelCase		:Tuple  ,_UpperCamelCase		:str=None  ,_UpperCamelCase		:int=None			):
									snake_case_		:							int			       =						file_names
									snake_case_		:							Union[str, Any]			       =						image_transform
									snake_case_		:							Union[str, Any]			       =						label_to_id





							def __len__(  self		:Any			):
									return len(self.file_names			)





							def __getitem__(  self		:Optional[int]  ,_UpperCamelCase		:int			):
									snake_case_		:							str			       =						self.file_names[idx]
									snake_case_		:							int			       =						PIL.Image.open(_UpperCAmelCase			)
									snake_case_		:							str			       =						raw_image.convert("""RGB"""			)
									if self.image_transform is not None:
											snake_case_		:							Optional[Any]			       =						self.image_transform(_UpperCAmelCase			)
									snake_case_		:							List[str]			       =						extract_label(_UpperCAmelCase			)
									if self.label_to_id is not None:
											snake_case_		:							str			       =						self.label_to_id[label]
									return {"image": image, "label": label}





def      UpperCAmelCase      (		lowerCamelCase_  :Any				,      lowerCamelCase_  :Optional[Any] ):







		'''simple docstring'''
		# Initialize accelerator
		if args.with_tracking:
				snake_case_		:							List[Any]			       =						Accelerator(
				    cpu=args.cpu				,      mixed_precision=args.mixed_precision				,      log_with="""all"""				,      project_dir=args.project_dir )
		else:
				snake_case_		:							Optional[int]			       =						Accelerator(cpu=args.cpu				,      mixed_precision=args.mixed_precision )

		# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
		snake_case_		:							Optional[int]			       =						config["""lr"""]
		snake_case_		:							Optional[Any]			       =						int(config["""num_epochs"""] )
		snake_case_		:							Optional[int]			       =						int(config["""seed"""] )
		snake_case_		:							List[str]			       =						int(config["""batch_size"""] )
		snake_case_		:							int			       =						config["""image_size"""]
		if not isinstance(lowerCAmelCase__				,      (list, tuple) ):
				snake_case_		:							Optional[int]			       =						(image_size, image_size)

		# Parse out whether we are saving every epoch or after a certain number of batches
		if hasattr(args.checkpointing_steps				,      """isdigit""" ):
				if args.checkpointing_steps == "epoch":
						snake_case_		:							int			       =						args.checkpointing_steps
				elif args.checkpointing_steps.isdigit():
						snake_case_		:							Union[str, Any]			       =						int(args.checkpointing_steps )
				else:
						raise ValueError(
						    F'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.''' )
		else:
				snake_case_		:							int			       =						None

		# We need to initialize the trackers we use, and also store our configuration
		if args.with_tracking:
				snake_case_		:							str			       =						os.path.split(lowerCAmelCase__ )[-1].split(""".""" )[0]
				accelerator.init_trackers(lowerCAmelCase__				,      lowerCAmelCase__ )

		# Grab all the image filenames
		snake_case_		:							Optional[Any]			       =						[os.path.join(args.data_dir				,      lowerCAmelCase__ ) for fname in os.listdir(args.data_dir ) if fname.endswith(""".jpg""" )]

		# Build the label correspondences
		snake_case_		:							Any			       =						[extract_label(lowerCAmelCase__ ) for fname in file_names]
		snake_case_		:							Optional[int]			       =						list(set(lowerCAmelCase__ ) )
		id_to_label.sort()
		snake_case_		:							Union[str, Any]			       =						{lbl: i for i, lbl in enumerate(lowerCAmelCase__ )}

		# Set the seed before splitting the data.
		np.random.seed(lowerCAmelCase__ )
		torch.manual_seed(lowerCAmelCase__ )
		torch.cuda.manual_seed_all(lowerCAmelCase__ )

		# Split our filenames between train and validation
		snake_case_		:							Dict			       =						np.random.permutation(len(lowerCAmelCase__ ) )
		snake_case_		:							Dict			       =						int(0.8 * len(lowerCAmelCase__ ) )
		snake_case_		:							List[Any]			       =						random_perm[:cut]
		snake_case_		:							Dict			       =						random_perm[cut:]

		# For training we use a simple RandomResizedCrop
		snake_case_		:							Optional[int]			       =						Compose([RandomResizedCrop(lowerCAmelCase__				,      scale=(0.5, 1.0) ), ToTensor()] )
		snake_case_		:							Optional[Any]			       =						PetsDataset(
		    [file_names[i] for i in train_split]				,      image_transform=lowerCAmelCase__				,      label_to_id=lowerCAmelCase__ )

		# For evaluation, we use a deterministic Resize
		snake_case_		:							Optional[Any]			       =						Compose([Resize(lowerCAmelCase__ ), ToTensor()] )
		snake_case_		:							List[str]			       =						PetsDataset([file_names[i] for i in eval_split]				,      image_transform=lowerCAmelCase__				,      label_to_id=lowerCAmelCase__ )

		# Instantiate dataloaders.
		snake_case_		:							int			       =						DataLoader(lowerCAmelCase__				,      shuffle=lowerCAmelCase__				,      batch_size=lowerCAmelCase__				,      num_workers=4 )
		snake_case_		:							int			       =						DataLoader(lowerCAmelCase__				,      shuffle=lowerCAmelCase__				,      batch_size=lowerCAmelCase__				,      num_workers=4 )

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		snake_case_		:							Optional[int]			       =						create_model("""resnet50d"""				,      pretrained=lowerCAmelCase__				,      num_classes=len(lowerCAmelCase__ ) )

		# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
		# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
		# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
		snake_case_		:							Tuple			       =						model.to(accelerator.device )

		# Freezing the base model
		for param in model.parameters():
				snake_case_		:							Tuple			       =						False
		for param in model.get_classifier().parameters():
				snake_case_		:							List[str]			       =						True

		# We normalize the batches of images to be a bit faster.
		snake_case_		:							Any			       =						torch.tensor(model.default_cfg["""mean"""] )[None, :, None, None].to(accelerator.device )
		snake_case_		:							Optional[Any]			       =						torch.tensor(model.default_cfg["""std"""] )[None, :, None, None].to(accelerator.device )

		# Instantiate optimizer
		snake_case_		:							Optional[int]			       =						torch.optim.Adam(params=model.parameters()				,      lr=lr / 25 )

		# Instantiate learning rate scheduler
		snake_case_		:							Optional[Any]			       =						OneCycleLR(optimizer=lowerCAmelCase__				,      max_lr=lowerCAmelCase__				,      epochs=lowerCAmelCase__				,      steps_per_epoch=len(lowerCAmelCase__ ) )

		# Prepare everything
		# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
		# prepare method.
		snake_case_       ,      snake_case_       ,      snake_case_       ,      snake_case_       ,      snake_case_		:							Dict			       =						accelerator.prepare(
		    lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__ )
		# We need to keep track of how many total steps we have iterated over
		snake_case_		:							Any			       =						0
		# We also need to keep track of the starting epoch so files are named properly
		snake_case_		:							int			       =						0

		# Potentially load in the weights and states from a previous save
		if args.resume_from_checkpoint:
				if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
						accelerator.print(F'''Resumed from checkpoint: {args.resume_from_checkpoint}''' )
						accelerator.load_state(args.resume_from_checkpoint )
						snake_case_		:							str			       =						os.path.basename(args.resume_from_checkpoint )
				else:
						# Get the most recent checkpoint
						snake_case_		:							str			       =						[f.name for f in os.scandir(os.getcwd() ) if f.is_dir()]
						dirs.sort(key=os.path.getctime )
						snake_case_		:							Union[str, Any]			       =						dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
				# Extract `epoch_{i}` or `step_{i}`
				snake_case_		:							Union[str, Any]			       =						os.path.splitext(lowerCAmelCase__ )[0]

				if "epoch" in training_difference:
						snake_case_		:							Optional[Any]			       =						int(training_difference.replace("""epoch_"""				,      """""" ) ) + 1
						snake_case_		:							List[str]			       =						None
				else:
						snake_case_		:							List[str]			       =						int(training_difference.replace("""step_"""				,      """""" ) )
						snake_case_		:							List[Any]			       =						resume_step // len(lowerCAmelCase__ )
						resume_step -= starting_epoch * len(lowerCAmelCase__ )

    # Now we train the model
		for epoch in range(lowerCAmelCase__				,      lowerCAmelCase__ ):
				model.train()
				if args.with_tracking:
						snake_case_		:							str			       =						0
				if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
						# We need to skip steps until we reach the resumed step
						snake_case_		:							Tuple			       =						accelerator.skip_first_batches(lowerCAmelCase__				,      lowerCAmelCase__ )
						overall_step += resume_step
				else:
						# After the first iteration though, we need to go back to the original dataloader
						snake_case_		:							Optional[Any]			       =						train_dataloader
				for batch in active_dataloader:
						# We could avoid this line since we set the accelerator with `device_placement=True`.
						snake_case_		:							Optional[int]			       =						{k: v.to(accelerator.device ) for k, v in batch.items()}
						snake_case_		:							str			       =						(batch["""image"""] - mean) / std
						snake_case_		:							Optional[Any]			       =						model(lowerCAmelCase__ )
						snake_case_		:							Any			       =						torch.nn.functional.cross_entropy(lowerCAmelCase__				,      batch["""label"""] )
						# We keep track of the loss at each epoch
						if args.with_tracking:
								total_loss += loss.detach().float()
						accelerator.backward(lowerCAmelCase__ )
						optimizer.step()
						lr_scheduler.step()
						optimizer.zero_grad()
						overall_step += 1
						if isinstance(lowerCAmelCase__				,      lowerCAmelCase__ ):
								snake_case_		:							List[str]			       =						F'''step_{overall_step}'''
								if overall_step % checkpointing_steps == 0:
										if args.output_dir is not None:
												snake_case_		:							List[Any]			       =						os.path.join(args.output_dir				,      lowerCAmelCase__ )
										accelerator.save_state(lowerCAmelCase__ )
				model.eval()
				snake_case_		:							Optional[Any]			       =						0
				snake_case_		:							Tuple			       =						0
				for step, batch in enumerate(lowerCAmelCase__ ):
						# We could avoid this line since we set the accelerator with `device_placement=True`.
						snake_case_		:							Optional[int]			       =						{k: v.to(accelerator.device ) for k, v in batch.items()}
						snake_case_		:							str			       =						(batch["""image"""] - mean) / std
						with torch.no_grad():
								snake_case_		:							List[Any]			       =						model(lowerCAmelCase__ )
						snake_case_		:							Any			       =						outputs.argmax(dim=-1 )
						snake_case_       ,      snake_case_		:							Optional[Any]			       =						accelerator.gather_for_metrics((predictions, batch["""label"""]) )
						snake_case_		:							Any			       =						predictions == references
						num_elems += accurate_preds.shape[0]
						accurate += accurate_preds.long().sum()

				snake_case_		:							Dict			       =						accurate.item() / num_elems
				# Use accelerator.print to print only on the main process.
				accelerator.print(F'''epoch {epoch}: {1_00 * eval_metric:.2f}''' )
				if args.with_tracking:
						accelerator.log(
						    {
						        """accuracy""": 1_00 * eval_metric,
						        """train_loss""": total_loss.item() / len(lowerCAmelCase__ ),
						        """epoch""": epoch,
						    }				,      step=lowerCAmelCase__				,      )
				if checkpointing_steps == "epoch":
						snake_case_		:							Optional[int]			       =						F'''epoch_{epoch}'''
						if args.output_dir is not None:
								snake_case_		:							str			       =						os.path.join(args.output_dir				,      lowerCAmelCase__ )
						accelerator.save_state(lowerCAmelCase__ )

		if args.with_tracking:
				accelerator.end_training()





def      UpperCAmelCase      (		):







		'''simple docstring'''
		snake_case_		:							Tuple			       =						argparse.ArgumentParser(description="""Simple example of training script.""" )
		parser.add_argument("""--data_dir"""				,      required=lowerCAmelCase__				,      help="""The data folder on disk.""" )
		parser.add_argument("""--fp16"""				,      action="""store_true"""				,      help="""If passed, will use FP16 training.""" )
		parser.add_argument(
		    """--mixed_precision"""				,      type=lowerCAmelCase__				,      default=lowerCAmelCase__				,      choices=["""no""", """fp16""", """bf16""", """fp8"""]				,      help="""Whether to use mixed precision. Choose"""
		    """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
		    """and an Nvidia Ampere GPU."""				,      )
		parser.add_argument("""--cpu"""				,      action="""store_true"""				,      help="""If passed, will train on the CPU.""" )
		parser.add_argument(
		    """--checkpointing_steps"""				,      type=lowerCAmelCase__				,      default=lowerCAmelCase__				,      help="""Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."""				,      )
		parser.add_argument(
		    """--output_dir"""				,      type=lowerCAmelCase__				,      default="""."""				,      help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""				,      )
		parser.add_argument(
		    """--resume_from_checkpoint"""				,      type=lowerCAmelCase__				,      default=lowerCAmelCase__				,      help="""If the training should continue from a checkpoint folder."""				,      )
		parser.add_argument(
		    """--with_tracking"""				,      action="""store_true"""				,      help="""Whether to load in all available experiment trackers from the environment and use them for logging."""				,      )
		parser.add_argument(
		    """--project_dir"""				,      type=lowerCAmelCase__				,      default="""logs"""				,      help="""Location on where to store experiment tracking logs` and relevent project information"""				,      )
		snake_case_		:							Tuple			       =						parser.parse_args()
		snake_case_		:							List[str]			       =						{"""lr""": 3E-2, """num_epochs""": 3, """seed""": 42, """batch_size""": 64, """image_size""": 2_24}
		training_function(lowerCAmelCase__				,      lowerCAmelCase__ )


if __name__ == "__main__":
							main()






'''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      UpperCAmelCase      (		lowerCamelCase_  :Tuple				,      lowerCamelCase_  :Union[str, Any]				,      lowerCamelCase_  :List[Any]=0 ):







		'''simple docstring'''
		# Format the message.
		if name is None:
				snake_case_		:							Tuple			       =						None
		else:
				snake_case_		:							Optional[Any]			       =						""".""" * max(0				,      spaces - 2 ) + """# {:""" + str(50 - spaces ) + """s}"""
				snake_case_		:							Optional[Any]			       =						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      UpperCAmelCase      (		lowerCamelCase_  :Optional[int]				,      lowerCamelCase_  :List[str]				,      lowerCamelCase_  :List[Any]				,      lowerCamelCase_  :List[str]				,      lowerCamelCase_  :int ):







		'''simple docstring'''
		# 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.
		snake_case_		:							Any			       =						param.size()
		if checkpoint_version == 1.0:
				# version 1.0 stores [num_heads * hidden_size * num_splits, :]
				snake_case_		:							List[str]			       =						(num_heads, hidden_size, num_splits) + input_shape[1:]
				snake_case_		:							List[Any]			       =						param.view(*lowerCamelCase_ )
				snake_case_		:							Tuple			       =						param.transpose(0				,      2 )
				snake_case_		:							List[str]			       =						param.transpose(1				,      2 ).contiguous()
		elif checkpoint_version >= 2.0:
				# other versions store [num_heads * num_splits * hidden_size, :]
				snake_case_		:							Tuple			       =						(num_heads, num_splits, hidden_size) + input_shape[1:]
				snake_case_		:							int			       =						param.view(*lowerCamelCase_ )
				snake_case_		:							str			       =						param.transpose(0				,      1 ).contiguous()
		snake_case_		:							int			       =						param.view(*lowerCamelCase_ )
		return param





def      UpperCAmelCase      (		lowerCamelCase_  :List[str]				,      lowerCamelCase_  :Optional[Any]				,      lowerCamelCase_  :Tuple ):







		'''simple docstring'''
		# The converted output model.
		snake_case_		:							Tuple			       =						{}

		# old versions did not store training args
		snake_case_		:							Optional[Any]			       =						input_state_dict.get("""args"""				,      lowerCamelCase_ )
		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))

				snake_case_		:							Optional[int]			       =						ds_args.padded_vocab_size
				snake_case_		:							str			       =						ds_args.max_position_embeddings
				snake_case_		:							Tuple			       =						ds_args.hidden_size
				snake_case_		:							List[str]			       =						ds_args.num_layers
				snake_case_		:							Union[str, Any]			       =						ds_args.num_attention_heads
				snake_case_		:							Tuple			       =						ds_args.ffn_hidden_size
				# pprint(config)

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

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

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

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

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

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

		# The simple map of names for "automated" rules.
		snake_case_		:							Tuple			       =						{
		    """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.
				snake_case_		:							List[Any]			       =						layer_re.match(lowerCamelCase_ )

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

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

				# The name of the layer.
				snake_case_		:							List[Any]			       =						F'''transformer.h.{layer_idx}'''

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

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

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

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

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

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

		# The final layernorm.
		snake_case_		:							str			       =						transformer["""final_layernorm.weight"""]
		snake_case_		:							int			       =						transformer["""final_layernorm.bias"""]

		# For LM head, transformers' wants the matrix to weight embeddings.
		snake_case_		:							Optional[int]			       =						word_embeddings

		# It should be done!
		return output_state_dict





def      UpperCAmelCase      (		):







		'''simple docstring'''
		# Create the argument parser.
		snake_case_		:							str			       =						argparse.ArgumentParser()
		parser.add_argument("""--print-checkpoint-structure"""				,      action="""store_true""" )
		parser.add_argument(
		    """path_to_checkpoint"""				,      type=lowerCamelCase_				,      help="""Path to the checkpoint file (.zip archive or direct .pt file)"""				,      )
		parser.add_argument(
		    """--config_file"""				,      default=""""""				,      type=lowerCamelCase_				,      help="""An optional config json file describing the pre-trained model."""				,      )
		snake_case_		:							Tuple			       =						parser.parse_args()

		# Extract the basename.
		snake_case_		:							List[str]			       =						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:
								snake_case_		:							Dict			       =						torch.load(lowerCamelCase_				,      map_location="""cpu""" )
		else:
				snake_case_		:							Optional[Any]			       =						torch.load(args.path_to_checkpoint				,      map_location="""cpu""" )

		snake_case_		:							Optional[int]			       =						input_state_dict.get("""args"""				,      lowerCamelCase_ )

		# 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:
								snake_case_		:							Tuple			       =						"""gelu_fast"""
						elif ds_args.openai_gelu:
								snake_case_		:							int			       =						"""gelu_new"""
						else:
								snake_case_		:							Optional[Any]			       =						"""gelu"""
				else:
						# in the very early days this used to be "gelu_new"
						snake_case_		:							int			       =						"""gelu_new"""

				# Spell out all parameters in case the defaults change.
				snake_case_		:							int			       =						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=lowerCamelCase_				,      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=lowerCamelCase_				,      summary_activation=lowerCamelCase_				,      summary_proj_to_labels=lowerCamelCase_				,      summary_first_dropout=0.1				,      scale_attn_weights=lowerCamelCase_				,      use_cache=lowerCamelCase_				,      bos_token_id=5_02_56				,      eos_token_id=5_02_56				,      )
		else:
				snake_case_		:							int			       =						GPTaConfig.from_json_file(args.config_file )

		snake_case_		:							Optional[int]			       =						["""GPT2LMHeadModel"""]

		# Convert.
		print("""Converting""" )
		snake_case_		:							Optional[Any]			       =						convert_megatron_checkpoint(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ )

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

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

		snake_case_		:							Union[str, Any]			       =						AutoTokenizer.from_pretrained(lowerCamelCase_ )
		snake_case_		:							int			       =						type(lowerCamelCase_ ).__name__
		snake_case_		:							Union[str, Any]			       =						tokenizer_class

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

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

		# Store the state_dict to file.
		snake_case_		:							Optional[Any]			       =						os.path.join(lowerCamelCase_				,      """pytorch_model.bin""" )
		print(F'''Saving checkpoint to "{output_checkpoint_file}"''' )
		torch.save(lowerCamelCase_				,      lowerCamelCase_ )


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

if __name__ == "__main__":
							main()

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