code: 
  """simple docstring"""
  import csv
  import dataclasses
  import json
  from dataclasses import dataclass
  from typing import List, Optional, Union
  
  from ...utils import is_tf_available, is_torch_available, logging
  
  
  __magic_name__		       =				logging.get_logger(__name__)
  
  
  
  
  
  
  
  @dataclass
  class    _lowerCAmelCase						:
  							lowercase_     :     str
  							lowercase_     :     str
  							lowercase_     :     Optional[str]       =							None
  							lowercase_     :     Optional[str]       =							None
  
  
  
  							def       _a   (      self   )							->							List[str]:
  								return json.dumps(dataclasses.asdict(self   ) ,						indent=2   ) + "\n"
  
  
  
  
  
  
  
  @dataclass(frozen=UpperCamelCase__						)
  class    _lowerCAmelCase						:
  							lowercase_     :     List[int]
  							lowercase_     :     Optional[List[int]]       =							None
  							lowercase_     :     Optional[List[int]]       =							None
  							lowercase_     :     Optional[Union[int, float]]       =							None
  
  
  
  							def       _a   (      self   )							->							Optional[int]:
  								return json.dumps(dataclasses.asdict(self   )   ) + "\n"
  
  
  
  
  
  
  
  class    _lowerCAmelCase						:
  
  
  
  							def       _a   (      self ,						a_   )							->							Union[str, Any]:
  								raise NotImplementedError()
  
  
  
  							def       _a   (      self ,						a_   )							->							str:
  								raise NotImplementedError()
  
  
  
  							def       _a   (      self ,						a_   )							->							List[Any]:
  								raise NotImplementedError()
  
  
  
  							def       _a   (      self ,						a_   )							->							Optional[int]:
  								raise NotImplementedError()
  
  
  
  							def       _a   (      self   )							->							Any:
  								raise NotImplementedError()
  
  
  
  							def       _a   (      self ,						a_   )							->							Dict:
  
  
  
  
  								if len(self.get_labels()   ) > 1:
  									_UpperCAmelCase =						self.get_labels()[int(example.label   )]
  								return example
  
  
  
  							@classmethod
  							def       _a   (      cls ,						a_ ,						a_=None   )							->							Union[str, Any]:
  
  
  
  
  								with open(lowerCAmelCase__ ,						"r" ,						encoding="utf-8-sig"   ) as f:
  									return list(csv.reader(lowerCAmelCase__ ,						delimiter="\t" ,						quotechar=lowerCAmelCase__   )   )
  
  
  
  
  
  
  
  class    _lowerCAmelCase						(						UpperCamelCase__						):
  
  
  
  							def __init__(      self ,						a_=None ,						a_=None ,						a_="classification" ,						a_=False   )							->							Any:
  								_UpperCAmelCase =						[] if labels is None else labels
  								_UpperCAmelCase =						[] if examples is None else examples
  								_UpperCAmelCase =						mode
  								_UpperCAmelCase =						verbose
  
  
  
  							def __len__(      self   )							->							Optional[int]:
  								return len(self.examples   )
  
  
  
  							def __getitem__(      self ,						a_   )							->							int:
  
  
  
  
  								if isinstance(lowerCAmelCase__ ,						lowerCAmelCase__   ):
  									return SingleSentenceClassificationProcessor(labels=self.labels ,						examples=self.examples[idx]   )
  								return self.examples[idx]
  
  
  
  							@classmethod
  							def       _a   (      cls ,						a_ ,						a_="" ,						a_=0 ,						a_=1 ,						a_=None ,						a_=False ,						**a_   )							->							List[str]:
  								_UpperCAmelCase =						cls(**lowerCAmelCase__   )
  								processor.add_examples_from_csv(
  								    lowerCAmelCase__ ,						split_name=lowerCAmelCase__ ,						column_label=lowerCAmelCase__ ,						column_text=lowerCAmelCase__ ,						column_id=lowerCAmelCase__ ,						skip_first_row=lowerCAmelCase__ ,						overwrite_labels=lowerCAmelCase__ ,						overwrite_examples=lowerCAmelCase__ ,						)
  								return processor
  
  
  
  							@classmethod
  							def       _a   (      cls ,						a_ ,						a_=None ,						**a_   )							->							int:
  								_UpperCAmelCase =						cls(**lowerCAmelCase__   )
  								processor.add_examples(lowerCAmelCase__ ,						labels=lowerCAmelCase__   )
  								return processor
  
  
  
  							def       _a   (      self ,						a_ ,						a_="" ,						a_=0 ,						a_=1 ,						a_=None ,						a_=False ,						a_=False ,						a_=False ,						)							->							Dict:
  								_UpperCAmelCase =						self._read_tsv(lowerCAmelCase__   )
  								if skip_first_row:
  									_UpperCAmelCase =						lines[1:]
  								_UpperCAmelCase =						[]
  								_UpperCAmelCase =						[]
  								_UpperCAmelCase =						[]
  								for i, line in enumerate(lowerCAmelCase__   ):
  									texts.append(line[column_text]   )
  									labels.append(line[column_label]   )
  									if column_id is not None:
  										ids.append(line[column_id]   )
  									else:
  										_UpperCAmelCase =						f"{split_name}-{i}" if split_name else str(lowerCAmelCase__   )
  										ids.append(lowerCAmelCase__   )
  
  								return self.add_examples(
  								    lowerCAmelCase__ ,						lowerCAmelCase__ ,						lowerCAmelCase__ ,						overwrite_labels=lowerCAmelCase__ ,						overwrite_examples=lowerCAmelCase__   )
  
  
  
  							def       _a   (      self ,						a_ ,						a_=None ,						a_=None ,						a_=False ,						a_=False   )							->							List[Any]:
  
  
  
  
  								if labels is not None and len(lowerCAmelCase__   ) != len(lowerCAmelCase__   ):
  									raise ValueError(
  									    f"Text and labels have mismatched lengths {len(lowerCAmelCase__   )} and {len(lowerCAmelCase__   )}"   )
  								if ids is not None and len(lowerCAmelCase__   ) != len(lowerCAmelCase__   ):
  									raise ValueError(f"Text and ids have mismatched lengths {len(lowerCAmelCase__   )} and {len(lowerCAmelCase__   )}"   )
  								if ids is None:
  									_UpperCAmelCase =						[None] * len(lowerCAmelCase__   )
  								if labels is None:
  									_UpperCAmelCase =						[None] * len(lowerCAmelCase__   )
  								_UpperCAmelCase =						[]
  								_UpperCAmelCase =						set()
  								for text_or_text_and_label, label, guid in zip(lowerCAmelCase__ ,						lowerCAmelCase__ ,						lowerCAmelCase__   ):
  									if isinstance(lowerCAmelCase__ ,						(tuple, list)   ) and label is None:
  										_UpperCAmelCase =						text_or_text_and_label
  									else:
  										_UpperCAmelCase =						text_or_text_and_label
  									added_labels.add(lowerCAmelCase__   )
  									examples.append(InputExample(guid=lowerCAmelCase__ ,						text_a=lowerCAmelCase__ ,						text_b=lowerCAmelCase__ ,						label=lowerCAmelCase__   )   )
  
  								# Update examples
  								if overwrite_examples:
  									_UpperCAmelCase =						examples
  								else:
  									self.examples.extend(lowerCAmelCase__   )
  
  								# Update labels
  								if overwrite_labels:
  									_UpperCAmelCase =						list(lowerCAmelCase__   )
  								else:
  									_UpperCAmelCase =						list(set(self.labels   ).union(lowerCAmelCase__   )   )
  
  								return self.examples
  
  
  
  							def       _a   (      self ,						a_ ,						a_=None ,						a_=False ,						a_=0 ,						a_=True ,						a_=None ,						)							->							int:
  
  
  
  
  								if max_length is None:
  									_UpperCAmelCase =						tokenizer.max_len
  
  								_UpperCAmelCase =						{label: i for i, label in enumerate(self.labels   )}
  
  								_UpperCAmelCase =						[]
  								for ex_index, example in enumerate(self.examples   ):
  									if ex_index % 10000 == 0:
  										logger.info(f"Tokenizing example {ex_index}"   )
  
  									_UpperCAmelCase =						tokenizer.encode(
  									    example.text_a ,						add_special_tokens=lowerCAmelCase__ ,						max_length=min(lowerCAmelCase__ ,						tokenizer.max_len   ) ,						)
  									all_input_ids.append(lowerCAmelCase__   )
  
  								_UpperCAmelCase =						max(len(lowerCAmelCase__   ) for input_ids in all_input_ids   )
  
  								_UpperCAmelCase =						[]
  								for ex_index, (input_ids, example) in enumerate(zip(lowerCAmelCase__ ,						self.examples   )   ):
  									if ex_index % 10000 == 0:
  										logger.info(f"Writing example {ex_index}/{len(self.examples   )}"   )
  									# The mask has 1 for real tokens and 0 for padding tokens. Only real
  									# tokens are attended to.
  									_UpperCAmelCase =						[1 if mask_padding_with_zero else 0] * len(lowerCAmelCase__   )
  
  									# Zero-pad up to the sequence length.
  									_UpperCAmelCase =						batch_length - len(lowerCAmelCase__   )
  									if pad_on_left:
  										_UpperCAmelCase =						([pad_token] * padding_length) + input_ids
  										_UpperCAmelCase =						([0 if mask_padding_with_zero else 1] * padding_length) + attention_mask
  									else:
  										_UpperCAmelCase =						input_ids + ([pad_token] * padding_length)
  										_UpperCAmelCase =						attention_mask + ([0 if mask_padding_with_zero else 1] * padding_length)
  
  									if len(lowerCAmelCase__   ) != batch_length:
  										raise ValueError(f"Error with input length {len(lowerCAmelCase__   )} vs {batch_length}"   )
  									if len(lowerCAmelCase__   ) != batch_length:
  										raise ValueError(f"Error with input length {len(lowerCAmelCase__   )} vs {batch_length}"   )
  
  									if self.mode == "classification":
  										_UpperCAmelCase =						label_map[example.label]
  									elif self.mode == "regression":
  										_UpperCAmelCase =						float(example.label   )
  									else:
  										raise ValueError(self.mode   )
  
  									if ex_index < 5 and self.verbose:
  										logger.info("*** Example ***"   )
  										logger.info(f"guid: {example.guid}"   )
  										logger.info(f"input_ids: {' '.join([str(lowerCAmelCase__   ) for x in input_ids]   )}"   )
  										logger.info(f"attention_mask: {' '.join([str(lowerCAmelCase__   ) for x in attention_mask]   )}"   )
  										logger.info(f"label: {example.label} (id = {label})"   )
  
  									features.append(InputFeatures(input_ids=lowerCAmelCase__ ,						attention_mask=lowerCAmelCase__ ,						label=lowerCAmelCase__   )   )
  
  								if return_tensors is None:
  									return features
  								elif return_tensors == "tf":
  									if not is_tf_available():
  										raise RuntimeError("return_tensors set to 'tf' but TensorFlow 2.0 can't be imported"   )
  									import tensorflow as tf
  
  									def gen():
  										for ex in features:
  											yield ({"input_ids": ex.input_ids, "attention_mask": ex.attention_mask}, ex.label)
  
  									_UpperCAmelCase =						tf.data.Dataset.from_generator(
  									    lowerCAmelCase__ ,						({"input_ids": tf.intaa, "attention_mask": tf.intaa}, tf.intaa) ,						({"input_ids": tf.TensorShape([None]   ), "attention_mask": tf.TensorShape([None]   )}, tf.TensorShape([]   )) ,						)
  									return dataset
  								elif return_tensors == "pt":
  									if not is_torch_available():
  										raise RuntimeError("return_tensors set to 'pt' but PyTorch can't be imported"   )
  									import torch
  									from torch.utils.data import TensorDataset
  
  									_UpperCAmelCase =						torch.tensor([f.input_ids for f in features] ,						dtype=torch.long   )
  									_UpperCAmelCase =						torch.tensor([f.attention_mask for f in features] ,						dtype=torch.long   )
  									if self.mode == "classification":
  										_UpperCAmelCase =						torch.tensor([f.label for f in features] ,						dtype=torch.long   )
  									elif self.mode == "regression":
  										_UpperCAmelCase =						torch.tensor([f.label for f in features] ,						dtype=torch.float   )
  
  									_UpperCAmelCase =						TensorDataset(lowerCAmelCase__ ,						lowerCAmelCase__ ,						lowerCAmelCase__   )
  									return dataset
  								else:
  									raise ValueError("return_tensors should be one of 'tf' or 'pt'"   )
  
  
  
  
  
style_context: 
  """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_bert/convert_megatron_bert_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 MegatronBertConfig
  
  
  
  
  def     __lowerCamelCase  (		UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__=0      ):
  
  
  	"""simple docstring"""
  
  
  
  
  
  
  	if name is None:
  		_UpperCAmelCase =						None
  	else:
  		_UpperCAmelCase =						"." * max(0				,     spaces - 2      ) + "# {:" + str(50 - spaces      ) + "s}"
  		_UpperCAmelCase =						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"""
  
  
  
  
  
  
  	_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(*UpperCamelCase__      )
  		_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(*UpperCamelCase__      )
  		_UpperCAmelCase =						param.transpose(0				,     1      ).contiguous()
  	_UpperCAmelCase =						param.view(*UpperCamelCase__      )
  	return param
  
  
  
  
  def     __lowerCamelCase  (		UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__      ):
  
  
  	"""simple docstring"""
  
  
  
  
  
  
  	_UpperCAmelCase =						{}
  
  	# old versions did not store training args
  	_UpperCAmelCase =						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))
  
  		_UpperCAmelCase =						ds_args.tokenizer_type
  		_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 if "ffn_hidden_size" in ds_args else 4 * ds_args.hidden_size
  		# pprint(config)
  
  	# The number of heads.
  	_UpperCAmelCase =						config.num_attention_heads
  	# The hidden_size per head.
  	_UpperCAmelCase =						config.hidden_size // heads
  	# 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, :]
  	# Store the word embeddings.
  	_UpperCAmelCase =						word_embeddings
  
  	# The position embeddings.
  	_UpperCAmelCase =						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.
  	_UpperCAmelCase =						pos_embeddings
  
  	# The token-type embeddings.
  	_UpperCAmelCase =						embeddings["tokentype_embeddings"]["weight"]
  	# Store the position embeddings.
  	_UpperCAmelCase =						tokentype_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": ".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.
  	_UpperCAmelCase =						None
  
  	# Extract the layers.
  	for key, val in transformer.items():
  		# Match the name.
  		_UpperCAmelCase =						layer_re.match(UpperCamelCase__      )
  
  		# 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"bert.encoder.layer.{layer_idx}"
  
  		# For layernorm(s), simply store the layer norm.
  		if op_name.endswith("layernorm"      ):
  			_UpperCAmelCase =						"attention.ln" if op_name.startswith("input"      ) else "ln"
  			_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":
  			# Make sure the QKV pointer is nil.
  			assert attention_qkv_weight is None, ""
  
  			_UpperCAmelCase =						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.
  			_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":
  			# 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.
  			_UpperCAmelCase =						attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
  			_UpperCAmelCase =						attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
  			_UpperCAmelCase =						attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]
  
  			_UpperCAmelCase =						fix_query_key_value_ordering(UpperCamelCase__				,     UpperCamelCase__				,     3				,     UpperCamelCase__				,     UpperCamelCase__      )
  			# Split the bias.
  			_UpperCAmelCase =						out_val[0 * config.hidden_size : 1 * config.hidden_size]
  			_UpperCAmelCase =						out_val[1 * config.hidden_size : 2 * config.hidden_size]
  			_UpperCAmelCase =						out_val[2 * config.hidden_size : 3 * config.hidden_size]
  
  			# Store.
  			_UpperCAmelCase =						q
  			_UpperCAmelCase =						q_bias
  			_UpperCAmelCase =						k
  			_UpperCAmelCase =						k_bias
  			_UpperCAmelCase =						v
  			_UpperCAmelCase =						v_bias
  
  			# Clear the stored tensor.
  			_UpperCAmelCase =						None
  
  		# Copy weights and biases as is.
  		elif weight_or_bias in ["weight", "bias"]:
  			_UpperCAmelCase =						megatron_to_transformers[op_name]
  			_UpperCAmelCase =						val
  
      # The final layernorm.
  	_UpperCAmelCase =						transformer["final_layernorm.weight"]
  	_UpperCAmelCase =						transformer["final_layernorm.bias"]
  
  	# The pooler.
  	_UpperCAmelCase =						lm["pooler"]
  
  	# Store the matrix and the bias.
  	_UpperCAmelCase =						pooler["dense.weight"]
  	_UpperCAmelCase =						pooler["dense.bias"]
  
  	# The LM head from Megatron (for RACE).
  	_UpperCAmelCase =						model["lm_head"]
  
  	# The transform matrix.
  	_UpperCAmelCase =						lm_head["dense.weight"]
  	_UpperCAmelCase =						lm_head["dense.bias"]
  
  	# The transform LN.
  	_UpperCAmelCase =						lm_head["layernorm.weight"]
  	_UpperCAmelCase =						lm_head["layernorm.bias"]
  
  	# For the decoder, we replicate the weights.
  	_UpperCAmelCase =						word_embeddings
  	_UpperCAmelCase =						lm_head["bias"]
  
  	# The classifier from Megatron (for MLNI).
  	_UpperCAmelCase =						model["binary_head"]
  
  	# Store the classifier.
  	_UpperCAmelCase =						binary_head["weight"]
  	_UpperCAmelCase =						binary_head["bias"]
  
  	# It should be done!
  	return output_state_dict
  
  
  
  
  def     __lowerCamelCase  (		):
  
  
  	"""simple docstring"""
  
  
  
  
  
  
  	_UpperCAmelCase =						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."				,     )
  	_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(UpperCamelCase__				,     map_location="cpu"      )
  	else:
  		_UpperCAmelCase =						torch.load(args.path_to_checkpoint				,     map_location="cpu"      )
  
  	if args.config_file == "":
  		# Default config of megatron-bert 345m
  		_UpperCAmelCase =						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
  		_UpperCAmelCase =						input_state_dict["model"]["lm_head"]["bias"].numel()
  	else:
  		_UpperCAmelCase =						MegatronBertConfig.from_json_file(args.config_file      )
  
  	# Convert.
  	print("Converting"      )
  	_UpperCAmelCase =						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.
  	_UpperCAmelCase =						os.path.join(UpperCamelCase__				,     "pytorch_model.bin"      )
  	print(f"Saving checkpoint to \"{output_checkpoint_file}\""      )
  	torch.save(UpperCamelCase__				,     UpperCamelCase__      )
  
  
  ####################################################################################################
  
  if __name__ == "__main__":
  			main()
  
  ####################################################################################################
  
  
  
  
  