





import copy
import os
from typing import Dict, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings, logging
from ..auto import AutoConfig


_UpperCAmelCase										=					logging.get_logger(__name__)


_UpperCAmelCase										=					{
    'suno/bark-small': 'https://huggingface.co/suno/bark-small/resolve/main/config.json',
    'suno/bark': 'https://huggingface.co/suno/bark/resolve/main/config.json',
}

_UpperCAmelCase										=					'\n    This is the configuration class to store the configuration of a [`{model}`]. It is used to instantiate the model\n    according to the specified arguments, defining the model architecture. Instantiating a configuration with the\n    defaults will yield a similar configuration to that of the Bark [suno/bark](https://huggingface.co/suno/bark)\n    architecture.\n\n    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the\n    documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        block_size (`int`, *optional*, defaults to 1024):\n            The maximum sequence length that this model might ever be used with. Typically set this to something large\n            just in case (e.g., 512 or 1024 or 2048).\n        input_vocab_size (`int`, *optional*, defaults to 10_048):\n            Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the\n            `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with\n            regards to the chosen sub-model.\n        output_vocab_size (`int`, *optional*, defaults to 10_048):\n            Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented\n            by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought\n            with regards to the chosen sub-model.\n        num_layers (`int`, *optional*, defaults to 12):\n            Number of hidden layers in the given sub-model.\n        num_heads (`int`, *optional*, defaults to 12):\n            Number of attention heads for each attention layer in the Transformer architecture.\n        hidden_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the "intermediate" (often named feed-forward) layer in the architecture.\n        dropout (`float`, *optional*, defaults to 0.0):\n            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.\n        bias (`bool`, *optional*, defaults to `True`):\n            Whether or not to use bias in the linear layers and layer norm layers.\n        initializer_range (`float`, *optional*, defaults to 0.02):\n            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n'



class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Optional[Any]       = '''bark_module'''
       _UpperCamelCase					:       Optional[int]       = ['''past_key_values''']

       _UpperCamelCase					:       Tuple       = {
           '''num_attention_heads''': '''num_heads''',
           '''num_hidden_layers''': '''num_layers''',
           '''vocab_size''': '''input_vocab_size''',
           '''window_size''': '''block_size''',
       }

       def __init__(		self:			List[str]       , _SCREAMING_SNAKE_CASE:			List[str]=1024       , _SCREAMING_SNAKE_CASE:			int=10048       , _SCREAMING_SNAKE_CASE:			Union[str, Any]=10048       , _SCREAMING_SNAKE_CASE:			Tuple=12       , _SCREAMING_SNAKE_CASE:			Any=12       , _SCREAMING_SNAKE_CASE:			str=768       , _SCREAMING_SNAKE_CASE:			int=0.0       , _SCREAMING_SNAKE_CASE:			List[str]=True       , _SCREAMING_SNAKE_CASE:			List[str]=0.02       , _SCREAMING_SNAKE_CASE:			Any=True       , **_SCREAMING_SNAKE_CASE:			Union[str, Any]       , )							->			Dict:




              """simple docstring"""


              UpperCamelCase_ 							=	block_size
              UpperCamelCase_ 							=	input_vocab_size
              UpperCamelCase_ 							=	output_vocab_size
              UpperCamelCase_ 							=	num_layers
              UpperCamelCase_ 							=	num_heads
              UpperCamelCase_ 							=	hidden_size
              UpperCamelCase_ 							=	dropout
              UpperCamelCase_ 							=	bias
              UpperCamelCase_ 							=	use_cache
              UpperCamelCase_ 							=	initializer_range

              super().__init__(**_SCREAMING_SNAKE_CASE		)



       @classmethod
       def lowercase			(		cls:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			Union[str, os.PathLike]       , _SCREAMING_SNAKE_CASE:			Optional[Union[str, os.PathLike]] = None       , _SCREAMING_SNAKE_CASE:			bool = False       , _SCREAMING_SNAKE_CASE:			bool = False       , _SCREAMING_SNAKE_CASE:			Optional[Union[str, bool]] = None       , _SCREAMING_SNAKE_CASE:			str = "main"       , **_SCREAMING_SNAKE_CASE:			str       , )							->			"PretrainedConfig":




              """simple docstring"""


              UpperCamelCase_ 							=	cache_dir
              UpperCamelCase_ 							=	force_download
              UpperCamelCase_ 							=	local_files_only
              UpperCamelCase_ 							=	revision

              cls._set_token_in_kwargs(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE		)

              UpperCamelCase_							,			UpperCamelCase_ 							=	cls.get_config_dict(_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE		)

              # get the config dict if we are loading from Bark
              if config_dict.get("model_type"		) == "bark":
                     UpperCamelCase_ 							=	config_dict[f'''{cls.model_type}_config''']

              if "model_type" in config_dict and hasattr(cls       , "model_type"		) and config_dict["model_type"] != cls.model_type:
                     logger.warning(
                         f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
                         f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

              return cls.from_dict(_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE		)



@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config='''BarkSemanticConfig'''	,    model='''BarkSemanticModel'''  )	,    '''
    Example:

    ```python
    >>> from transformers import BarkSemanticConfig, BarkSemanticModel

    >>> # Initializing a Bark sub-module style configuration
    >>> configuration = BarkSemanticConfig()

    >>> # Initializing a model (with random weights) from the suno/bark style configuration
    >>> model = BarkSemanticModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```'''	,    )
class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Union[str, Any]       = '''semantic'''



@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config='''BarkCoarseConfig'''	,    model='''BarkCoarseModel'''  )	,    '''
    Example:

    ```python
    >>> from transformers import BarkCoarseConfig, BarkCoarseModel

    >>> # Initializing a Bark sub-module style configuration
    >>> configuration = BarkCoarseConfig()

    >>> # Initializing a model (with random weights) from the suno/bark style configuration
    >>> model = BarkCoarseModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```'''	,    )
class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Optional[Any]       = '''coarse_acoustics'''



@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config='''BarkFineConfig'''	,    model='''BarkFineModel'''  )	,    '''
        n_codes_total (`int`, *optional*, defaults to 8):
            The total number of audio codebooks predicted. Used in the fine acoustics sub-model.
        n_codes_given (`int`, *optional*, defaults to 1):
            The number of audio codebooks predicted in the coarse acoustics sub-model. Used in the acoustics
            sub-models.
    Example:

    ```python
    >>> from transformers import BarkFineConfig, BarkFineModel

    >>> # Initializing a Bark sub-module style configuration
    >>> configuration = BarkFineConfig()

    >>> # Initializing a model (with random weights) from the suno/bark style configuration
    >>> model = BarkFineModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```'''	,    )
class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Tuple       = '''fine_acoustics'''

       def __init__(		self:			List[str]       , _SCREAMING_SNAKE_CASE:			List[str]=True       , _SCREAMING_SNAKE_CASE:			Optional[Any]=8       , _SCREAMING_SNAKE_CASE:			Dict=1       , **_SCREAMING_SNAKE_CASE:			List[Any]		)							->			List[str]:




              """simple docstring"""


              UpperCamelCase_ 							=	n_codes_total
              UpperCamelCase_ 							=	n_codes_given

              super().__init__(tie_word_embeddings=_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE		)



class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       List[Any]       = '''bark'''
       _UpperCamelCase					:       List[Any]       = True

       def __init__(		self:			Tuple       , _SCREAMING_SNAKE_CASE:			Dict = None       , _SCREAMING_SNAKE_CASE:			Dict = None       , _SCREAMING_SNAKE_CASE:			Dict = None       , _SCREAMING_SNAKE_CASE:			Dict = None       , _SCREAMING_SNAKE_CASE:			Tuple=0.02       , **_SCREAMING_SNAKE_CASE:			Optional[int]       , )							->			Any:




              """simple docstring"""


              if semantic_config is None:
                     UpperCamelCase_ 							=	{}
                     logger.info("semantic_config is None. initializing the semantic model with default values."		)

              if coarse_acoustics_config is None:
                     UpperCamelCase_ 							=	{}
                     logger.info("coarse_acoustics_config is None. initializing the coarse model with default values."		)

              if fine_acoustics_config is None:
                     UpperCamelCase_ 							=	{}
                     logger.info("fine_acoustics_config is None. initializing the fine model with default values."		)

              if codec_config is None:
                     UpperCamelCase_ 							=	{}
                     logger.info("codec_config is None. initializing the codec model with default values."		)

              UpperCamelCase_ 							=	BarkSemanticConfig(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	BarkCoarseConfig(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	BarkFineConfig(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	AutoConfig.for_model(**_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_ 							=	initializer_range

              super().__init__(**_SCREAMING_SNAKE_CASE		)

       @classmethod
       def lowercase			(		cls:			Any       , _SCREAMING_SNAKE_CASE:			BarkSemanticConfig       , _SCREAMING_SNAKE_CASE:			BarkCoarseConfig       , _SCREAMING_SNAKE_CASE:			BarkFineConfig       , _SCREAMING_SNAKE_CASE:			AutoConfig       , **_SCREAMING_SNAKE_CASE:			List[Any]       , )							->			Dict:




              """simple docstring"""


              return cls(
                  semantic_config=semantic_config.to_dict()       , coarse_acoustics_config=coarse_acoustics_config.to_dict()       , fine_acoustics_config=fine_acoustics_config.to_dict()       , codec_config=codec_config.to_dict()       , **_SCREAMING_SNAKE_CASE       , )



       def lowercase			(		self:			str		)							->			Optional[int]:




              """simple docstring"""


              UpperCamelCase_ 							=	copy.deepcopy(self.__dict__		)

              UpperCamelCase_ 							=	self.semantic_config.to_dict()
              UpperCamelCase_ 							=	self.coarse_acoustics_config.to_dict()
              UpperCamelCase_ 							=	self.fine_acoustics_config.to_dict()
              UpperCamelCase_ 							=	self.codec_config.to_dict()

              UpperCamelCase_ 							=	self.__class__.model_type
              return output








# 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[int] ,     snake_case_		:						int ,     snake_case_		:						int=0						):
     # Format the message.
     if name is None:
          __magic_name__				=       None
     else:
          __magic_name__				=       '''.''' * max(0 ,     spaces - 2						) + '''# {:''' + str(50 - spaces						) + '''s}'''
          __magic_name__				=       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_		:						int ,     snake_case_		:						Optional[int] ,     snake_case_		:						Any ,     snake_case_		:						Tuple ,     snake_case_		:						Tuple						):
     # 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.
     __magic_name__				=       param.size()
     if checkpoint_version == 1.0:
          # version 1.0 stores [num_heads * hidden_size * num_splits, :]
          __magic_name__				=       (num_heads, hidden_size, num_splits) + input_shape[1:]
          __magic_name__				=       param.view(*snake_case_						)
          __magic_name__				=       param.transpose(0 ,     2						)
          __magic_name__				=       param.transpose(1 ,     2						).contiguous()
     elif checkpoint_version >= 2.0:
          # other versions store [num_heads * num_splits * hidden_size, :]
          __magic_name__				=       (num_heads, num_splits, hidden_size) + input_shape[1:]
          __magic_name__				=       param.view(*snake_case_						)
          __magic_name__				=       param.transpose(0 ,     1						).contiguous()
     __magic_name__				=       param.view(*snake_case_						)
     return param




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						int ,     snake_case_		:						Optional[int] ,     snake_case_		:						Union[str, Any]						):
     # The converted output model.
     __magic_name__				=       {}

     # old versions did not store training args
     __magic_name__				=       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))

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

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

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

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

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

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

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

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

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

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

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

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

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

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

               __magic_name__				=       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.
               __magic_name__				=       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.
               __magic_name__				=       attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
               __magic_name__				=       attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
               __magic_name__				=       attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

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

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

               # Clear the stored tensor.
               __magic_name__				=       None

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

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

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

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

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

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

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

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

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

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

     # It should be done!
     return output_state_dict




def   _SCREAMING_SNAKE_CASE      (	):
     # Create the argument parser.
     __magic_name__				=       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.''' ,     )
     __magic_name__				=       parser.parse_args()

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

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

     # Convert.
     print('''Converting'''						)
     __magic_name__				=       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.
     __magic_name__				=       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()

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