import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


if TYPE_CHECKING:
				from transformers.pipelines.conversational import Conversation

__SCREAMING_SNAKE_CASE						:    Dict     = logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE						:    List[Any]     = {
    'vocab_file': 'vocab.json',
    'merges_file': 'merges.txt',
}

__SCREAMING_SNAKE_CASE						:    Tuple     = {
    'vocab_file': {
        'gpt2': 'https://huggingface.co/gpt2/resolve/main/vocab.json',
        'gpt2-medium': 'https://huggingface.co/gpt2-medium/resolve/main/vocab.json',
        'gpt2-large': 'https://huggingface.co/gpt2-large/resolve/main/vocab.json',
        'gpt2-xl': 'https://huggingface.co/gpt2-xl/resolve/main/vocab.json',
        'distilgpt2': 'https://huggingface.co/distilgpt2/resolve/main/vocab.json',
    },
    'merges_file': {
        'gpt2': 'https://huggingface.co/gpt2/resolve/main/merges.txt',
        'gpt2-medium': 'https://huggingface.co/gpt2-medium/resolve/main/merges.txt',
        'gpt2-large': 'https://huggingface.co/gpt2-large/resolve/main/merges.txt',
        'gpt2-xl': 'https://huggingface.co/gpt2-xl/resolve/main/merges.txt',
        'distilgpt2': 'https://huggingface.co/distilgpt2/resolve/main/merges.txt',
    },
}

__SCREAMING_SNAKE_CASE						:    Optional[Any]     = {
    'gpt2': 1_0_2_4,
    'gpt2-medium': 1_0_2_4,
    'gpt2-large': 1_0_2_4,
    'gpt2-xl': 1_0_2_4,
    'distilgpt2': 1_0_2_4,
}
@lru_cache()
def 					snake_case						()       ->       int:

						'''simple docstring'''



						_snake_case						:     List[str]				=    (
						    list(range(ord("!"						)					,			ord("~"						) + 1						)						) + list(range(ord("¡"						)					,			ord("¬"						) + 1						)						) + list(range(ord("®"						)					,			ord("ÿ"						) + 1						)						)
						)
						_snake_case						:     Union[str, Any]				=    bs[:]
						_snake_case						:     List[str]				=    0
						for b in range(2**8						):
												if b not in bs:
																		bs.append(__lowercase						)
																		cs.append(2**8 + n						)
																		n += 1
						_snake_case						:     Tuple				=    [chr(__lowercase						) for n in cs]
						return dict(zip(__lowercase					,			__lowercase						)						)
def 					snake_case						(__lowercase						)       ->       Optional[int]:

						'''simple docstring'''



						_snake_case						:     List[Any]				=    set()
						_snake_case						:     Optional[int]				=    word[0]
						for char in word[1:]:
												pairs.add((prev_char, char)						)
												_snake_case						:     Optional[Any]				=    char
						return pairs



class    lowercase_    (     __snake_case		):
							_lowerCamelCase						=		VOCAB_FILES_NAMES
							_lowerCamelCase						=		PRETRAINED_VOCAB_FILES_MAP
							_lowerCamelCase						=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							_lowerCamelCase						=		['input_ids', 'attention_mask']



							def __init__(						self      ,				lowercase_      ,				lowercase_      ,				lowercase_="replace"      ,				lowercase_="<|endoftext|>"      ,				lowercase_="<|endoftext|>"      ,				lowercase_="<|endoftext|>"      ,				lowercase_=None      ,				lowercase_=False      ,				lowercase_=False      ,				**lowercase_      ,				):
													_snake_case						:     int				=    AddedToken(lowercase_      ,				lstrip=lowercase_      ,				rstrip=lowercase_				) if isinstance(lowercase_      ,				lowercase_				) else bos_token
													_snake_case						:     Dict				=    AddedToken(lowercase_      ,				lstrip=lowercase_      ,				rstrip=lowercase_				) if isinstance(lowercase_      ,				lowercase_				) else eos_token
													_snake_case						:     Optional[Any]				=    AddedToken(lowercase_      ,				lstrip=lowercase_      ,				rstrip=lowercase_				) if isinstance(lowercase_      ,				lowercase_				) else unk_token
													_snake_case						:     Union[str, Any]				=    AddedToken(lowercase_      ,				lstrip=lowercase_      ,				rstrip=lowercase_				) if isinstance(lowercase_      ,				lowercase_				) else pad_token
													super().__init__(
													    errors=lowercase_      ,				unk_token=lowercase_      ,				bos_token=lowercase_      ,				eos_token=lowercase_      ,				pad_token=lowercase_      ,				add_prefix_space=lowercase_      ,				add_bos_token=lowercase_      ,				**lowercase_      ,				)
													_snake_case						:     str				=    add_bos_token

													with open(lowercase_      ,				encoding="utf-8"				) as vocab_handle:
																			_snake_case						:     str				=    json.load(lowercase_				)
													_snake_case						:     Optional[int]				=    {v: k for k, v in self.encoder.items()}
													_snake_case						:     List[str]				=    errors  # how to handle errors in decoding
													_snake_case						:     Tuple				=    bytes_to_unicode()
													_snake_case						:     Optional[Any]				=    {v: k for k, v in self.byte_encoder.items()}
													with open(lowercase_      ,				encoding="utf-8"				) as merges_handle:
																			_snake_case						:     Union[str, Any]				=    merges_handle.read().split("\n"				)[1:-1]
													_snake_case						:     Optional[Any]				=    [tuple(merge.split()				) for merge in bpe_merges]
													_snake_case						:     Optional[int]				=    dict(zip(lowercase_      ,				range(len(lowercase_				)				)				)				)
													_snake_case						:     str				=    {}
													_snake_case						:     str				=    add_prefix_space

													# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
													_snake_case						:     str				=    re.compile(R"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"				)



							@property
							def        UpperCamelCase       (						self				):
													return len(self.encoder				)



							def        UpperCamelCase       (						self				):
													return dict(self.encoder      ,				**self.added_tokens_encoder				)



							def        UpperCamelCase       (						self      ,				lowercase_				):
													if token in self.cache:
																			return self.cache[token]
													_snake_case						:     List[Any]				=    tuple(lowercase_				)
													_snake_case						:     List[str]				=    get_pairs(lowercase_				)

													if not pairs:
																			return token

													while True:
																			_snake_case						:     int				=    min(lowercase_      ,				key=lambda lowercase_				: self.bpe_ranks.get(lowercase_      ,				float("inf"				)				)				)
																			if bigram not in self.bpe_ranks:
																									break
																			_snake_case				,_snake_case						:     Tuple				=    bigram
																			_snake_case						:     Dict				=    []
																			_snake_case						:     List[Any]				=    0
																			while i < len(lowercase_				):
																									try:
																															_snake_case						:     Any				=    word.index(lowercase_      ,				lowercase_				)
																									except ValueError:
																															new_word.extend(word[i:]				)
																															break
																									else:
																															new_word.extend(word[i:j]				)
																															_snake_case						:     int				=    j

																									if word[i] == first and i < len(lowercase_				) - 1 and word[i + 1] == second:
																															new_word.append(first + second				)
																															i += 2
																									else:
																															new_word.append(word[i]				)
																															i += 1
																			_snake_case						:     Optional[Any]				=    tuple(lowercase_				)
																			_snake_case						:     str				=    new_word
																			if len(lowercase_				) == 1:
																									break
																			else:
																									_snake_case						:     List[Any]				=    get_pairs(lowercase_				)
													_snake_case						:     List[Any]				=    " ".join(lowercase_				)
													_snake_case						:     Dict				=    word
													return word



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_=None				):
													if self.add_bos_token:
																			_snake_case						:     Union[str, Any]				=    [self.bos_token_id]
													else:
																			_snake_case						:     Union[str, Any]				=    []

													_snake_case						:     Optional[Any]				=    bos_token_ids + token_ids_a

													if token_ids_a is None:
																			return output

													return output + bos_token_ids + token_ids_a



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_ = None      ,				lowercase_ = False				):
													if already_has_special_tokens:
																			return super().get_special_tokens_mask(
																			    token_ids_a=lowercase_      ,				token_ids_a=lowercase_      ,				already_has_special_tokens=lowercase_				)

													if not self.add_bos_token:
																			return super().get_special_tokens_mask(
																			    token_ids_a=lowercase_      ,				token_ids_a=lowercase_      ,				already_has_special_tokens=lowercase_				)

													if token_ids_a is None:
																			return [1] + ([0] * len(lowercase_				))
													return [1] + ([0] * len(lowercase_				)) + [1] + ([0] * len(lowercase_				))



							def        UpperCamelCase       (						self      ,				lowercase_				):
													_snake_case						:     Dict				=    []
													for token in re.findall(self.pat      ,				lowercase_				):
																			_snake_case						:     Union[str, Any]				=    "".join(
																			    self.byte_encoder[b] for b in token.encode("utf-8"				)				)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
																			bpe_tokens.extend(bpe_token for bpe_token in self.bpe(lowercase_				).split(" "				)				)
													return bpe_tokens



							def        UpperCamelCase       (						self      ,				lowercase_				):
													return self.encoder.get(lowercase_      ,				self.encoder.get(self.unk_token				)				)



							def        UpperCamelCase       (						self      ,				lowercase_				):
													return self.decoder.get(lowercase_				)



							def        UpperCamelCase       (						self      ,				lowercase_				):
													_snake_case						:     int				=    "".join(lowercase_				)
													_snake_case						:     str				=    bytearray([self.byte_decoder[c] for c in text]				).decode("utf-8"      ,				errors=self.errors				)
													return text



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_ = None				):
													if not os.path.isdir(lowercase_				):
																			logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""				)
																			return
													_snake_case						:     Optional[int]				=    os.path.join(
													    lowercase_      ,				(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]				)
													_snake_case						:     Optional[Any]				=    os.path.join(
													    lowercase_      ,				(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]				)

													with open(lowercase_      ,				"w"      ,				encoding="utf-8"				) as f:
																			f.write(json.dumps(self.encoder      ,				indent=2      ,				sort_keys=lowercase_      ,				ensure_ascii=lowercase_				) + "\n"				)

													_snake_case						:     str				=    0
													with open(lowercase_      ,				"w"      ,				encoding="utf-8"				) as writer:
																			writer.write("#version: 0.2\n"				)
																			for bpe_tokens, token_index in sorted(self.bpe_ranks.items()      ,				key=lambda lowercase_				: kv[1]				):
																									if index != token_index:
																															logger.warning(
																															    f"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
																															    " Please check that the tokenizer is not corrupted!"				)
																															_snake_case						:     Dict				=    token_index
																									writer.write(" ".join(lowercase_				) + "\n"				)
																									index += 1

													return vocab_file, merge_file



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_=False      ,				**lowercase_				):
													_snake_case						:     Optional[Any]				=    kwargs.pop("add_prefix_space"      ,				self.add_prefix_space				)
													if is_split_into_words or add_prefix_space:
																			_snake_case						:     Optional[int]				=    " " + text
													return (text, kwargs)



							def        UpperCamelCase       (						self      ,				lowercase_				):
													_snake_case						:     str				=    []
													for is_user, text in conversation.iter_texts():
																			input_ids.extend(self.encode(lowercase_      ,				add_special_tokens=lowercase_				) + [self.eos_token_id]				)
													if len(lowercase_				) > self.model_max_length:
																			_snake_case						:     str				=    input_ids[-self.model_max_length :]
													return input_ids






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

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





def SCREAMING_SNAKE_CASE_   (       snake_case   :					Optional[int]						,  snake_case   :					Optional[int]						,  snake_case   :					List[str]						,  snake_case   :					Union[str, Any]						,  snake_case   :					int		)->		List[str]:
       # 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 BERT.
       _lowerCamelCase    	=					param.size()
       if checkpoint_version == 1.0:
              # version 1.0 stores [num_heads * hidden_size * num_splits, :]
              _lowerCamelCase    	=					(num_heads, hidden_size, num_splits) + input_shape[1:]
              _lowerCamelCase    	=					param.view(*snake_case		)
              _lowerCamelCase    	=					param.transpose(0						,  2		)
              _lowerCamelCase    	=					param.transpose(1						,  2		).contiguous()
       elif checkpoint_version >= 2.0:
              # other versions store [num_heads * num_splits * hidden_size, :]
              _lowerCamelCase    	=					(num_heads, num_splits, hidden_size) + input_shape[1:]
              _lowerCamelCase    	=					param.view(*snake_case		)
              _lowerCamelCase    	=					param.transpose(0						,  1		).contiguous()
       _lowerCamelCase    	=					param.view(*snake_case		)
       return param





def SCREAMING_SNAKE_CASE_   (       snake_case   :					Union[str, Any]						,  snake_case   :					List[Any]						,  snake_case   :					List[Any]		)->		Optional[Any]:
       # The converted output model.
       _lowerCamelCase    	=					{}

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

              _lowerCamelCase    	=					ds_args.tokenizer_type
              _lowerCamelCase    	=					ds_args.padded_vocab_size
              _lowerCamelCase    	=					ds_args.max_position_embeddings
              _lowerCamelCase    	=					ds_args.hidden_size
              _lowerCamelCase    	=					ds_args.num_layers
              _lowerCamelCase    	=					ds_args.num_attention_heads
              _lowerCamelCase    	=					ds_args.ffn_hidden_size if 'ffn_hidden_size' in ds_args else 4 * ds_args.hidden_size
              # pprint(config)

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

       # The model.
       _lowerCamelCase    	=					input_state_dict['model']
       # The language model.
       _lowerCamelCase    	=					model['language_model']
       # The embeddings.
       _lowerCamelCase    	=					lm['embedding']

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

       # The position embeddings.
       _lowerCamelCase    	=					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.
       _lowerCamelCase    	=					pos_embeddings

       # The token-type embeddings.
       _lowerCamelCase    	=					embeddings['tokentype_embeddings']['weight']
       # Store the position embeddings.
       _lowerCamelCase    	=					tokentype_embeddings

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

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

       # The simple map of names for "automated" rules.
       _lowerCamelCase    	=					{
           '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.
       _lowerCamelCase    	=					None

       # Extract the layers.
       for key, val in transformer.items():
              # Match the name.
              _lowerCamelCase    	=					layer_re.match(snake_case		)

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

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

              # The name of the layer.
              _lowerCamelCase    	=					f'bert.encoder.layer.{layer_idx}'

              # For layernorm(s), simply store the layer norm.
              if op_name.endswith('layernorm'		):
                     _lowerCamelCase    	=					'attention.ln' if op_name.startswith('input'		) else 'ln'
                     _lowerCamelCase    	=					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, ""

                     _lowerCamelCase    	=					fix_query_key_value_ordering(snake_case						,  snake_case						,  3						,  snake_case						,  snake_case		)
                     # Store the tensor as we need the bias as well to interleave QKV and biases.
                     _lowerCamelCase    	=					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.
                     _lowerCamelCase    	=					attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
                     _lowerCamelCase    	=					attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
                     _lowerCamelCase    	=					attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

                     _lowerCamelCase    	=					fix_query_key_value_ordering(snake_case						,  snake_case						,  3						,  snake_case						,  snake_case		)
                     # Split the bias.
                     _lowerCamelCase    	=					out_val[0 * config.hidden_size : 1 * config.hidden_size]
                     _lowerCamelCase    	=					out_val[1 * config.hidden_size : 2 * config.hidden_size]
                     _lowerCamelCase    	=					out_val[2 * config.hidden_size : 3 * config.hidden_size]

                     # Store.
                     _lowerCamelCase    	=					q
                     _lowerCamelCase    	=					q_bias
                     _lowerCamelCase    	=					k
                     _lowerCamelCase    	=					k_bias
                     _lowerCamelCase    	=					v
                     _lowerCamelCase    	=					v_bias

                     # Clear the stored tensor.
                     _lowerCamelCase    	=					None

              # Copy weights and biases as is.
              elif weight_or_bias in ["weight", "bias"]:
                     _lowerCamelCase    	=					megatron_to_transformers[op_name]
                     _lowerCamelCase    	=					val

    # The final layernorm.
       _lowerCamelCase    	=					transformer['final_layernorm.weight']
       _lowerCamelCase    	=					transformer['final_layernorm.bias']

       # The pooler.
       _lowerCamelCase    	=					lm['pooler']

       # Store the matrix and the bias.
       _lowerCamelCase    	=					pooler['dense.weight']
       _lowerCamelCase    	=					pooler['dense.bias']

       # The LM head from Megatron (for RACE).
       _lowerCamelCase    	=					model['lm_head']

       # The transform matrix.
       _lowerCamelCase    	=					lm_head['dense.weight']
       _lowerCamelCase    	=					lm_head['dense.bias']

       # The transform LN.
       _lowerCamelCase    	=					lm_head['layernorm.weight']
       _lowerCamelCase    	=					lm_head['layernorm.bias']

       # For the decoder, we replicate the weights.
       _lowerCamelCase    	=					word_embeddings
       _lowerCamelCase    	=					lm_head['bias']

       # The classifier from Megatron (for MLNI).
       _lowerCamelCase    	=					model['binary_head']

       # Store the classifier.
       _lowerCamelCase    	=					binary_head['weight']
       _lowerCamelCase    	=					binary_head['bias']

       # It should be done!
       return output_state_dict





def SCREAMING_SNAKE_CASE_   (       )->		Tuple:
       # Create the argument parser.
       _lowerCamelCase    	=					argparse.ArgumentParser()
       parser.add_argument('--print-checkpoint-structure'						,  action='store_true'		)
       parser.add_argument('path_to_checkpoint'						,  type=snake_case						,  help='Path to the ZIP file containing the checkpoint'		)
       parser.add_argument(
           '--config_file'						,  default=''						,  type=snake_case						,  help='An optional config json file describing the pre-trained model.'						,  )
       _lowerCamelCase    	=					parser.parse_args()

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

       if args.config_file == "":
              # Default config of megatron-bert 345m
              _lowerCamelCase    	=					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
              _lowerCamelCase    	=					input_state_dict['model']['lm_head']['bias'].numel()
       else:
              _lowerCamelCase    	=					MegatronBertConfig.from_json_file(args.config_file		)

       # Convert.
       print('Converting'		)
       _lowerCamelCase    	=					convert_megatron_checkpoint(snake_case						,  snake_case						,  snake_case		)

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

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

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


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

if __name__ == "__main__":
      main()

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






