





'''simple docstring'''







import argparse
import os
import re
import zipfile

import torch

from transformers import MegatronBertConfig

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[Any]       ,  lowerCAmelCase__     :	Any       ,  lowerCAmelCase__     :	Optional[Any]=0)			->						Any:

     '''simple docstring'''




     if name is None:
          _lowercase							:      Union[str, Any]  						=							None
     else:
          _lowercase							:      Optional[Any]  						=							'.' * max(0       ,  spaces - 2) + '# {:' + str(50 - spaces) + 's}'
          _lowercase							:      str  						=							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 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	int       ,  lowerCAmelCase__     :	Optional[Any]       ,  lowerCAmelCase__     :	Optional[int]       ,  lowerCAmelCase__     :	Optional[Any]       ,  lowerCAmelCase__     :	Union[str, Any])			->						str:

     '''simple docstring'''




     _lowercase							:      str  						=							param.size()
     if checkpoint_version == 1.0:
          # version 1.0 stores [num_heads * hidden_size * num_splits, :]
          _lowercase							:      Tuple  						=							(num_heads, hidden_size, num_splits) + input_shape[1:]
          _lowercase							:      Dict  						=							param.view(*lowerCAmelCase__)
          _lowercase							:      Tuple  						=							param.transpose(0       ,  2)
          _lowercase							:      Dict  						=							param.transpose(1       ,  2).contiguous()
     elif checkpoint_version >= 2.0:
          # other versions store [num_heads * num_splits * hidden_size, :]
          _lowercase							:      Tuple  						=							(num_heads, num_splits, hidden_size) + input_shape[1:]
          _lowercase							:      Optional[Any]  						=							param.view(*lowerCAmelCase__)
          _lowercase							:      Union[str, Any]  						=							param.transpose(0       ,  1).contiguous()
     _lowercase							:      str  						=							param.view(*lowerCAmelCase__)
     return param

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[Any]       ,  lowerCAmelCase__     :	str       ,  lowerCAmelCase__     :	Tuple)			->						Dict:

     '''simple docstring'''




     _lowercase							:      List[Any]  						=							{}

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

          _lowercase							:      Optional[Any]  						=							ds_args.tokenizer_type
          _lowercase							:      Optional[int]  						=							ds_args.padded_vocab_size
          _lowercase							:      int  						=							ds_args.max_position_embeddings
          _lowercase							:      int  						=							ds_args.hidden_size
          _lowercase							:      List[str]  						=							ds_args.num_layers
          _lowercase							:      Optional[Any]  						=							ds_args.num_attention_heads
          _lowercase							:      Union[str, Any]  						=							ds_args.ffn_hidden_size if 'ffn_hidden_size' in ds_args else 4 * ds_args.hidden_size
          # pprint(config)

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

     # The model.
     _lowercase							:      Any  						=							input_state_dict['model']
     # The language model.
     _lowercase							:      Optional[Any]  						=							model['language_model']
     # The embeddings.
     _lowercase							:      Union[str, Any]  						=							lm['embedding']

     # The word embeddings.
     _lowercase							:      List[Any]  						=							embeddings['word_embeddings']['weight']
     # Truncate the embedding table to vocab_size rows.
     _lowercase							:      Optional[Any]  						=							word_embeddings[: config.vocab_size, :]
     # Store the word embeddings.
     _lowercase							:      Dict  						=							word_embeddings

     # The position embeddings.
     _lowercase							:      List[Any]  						=							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.
     _lowercase							:      Union[str, Any]  						=							pos_embeddings

     # The token-type embeddings.
     _lowercase							:      str  						=							embeddings['tokentype_embeddings']['weight']
     # Store the position embeddings.
     _lowercase							:      int  						=							tokentype_embeddings

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

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

     # The simple map of names for "automated" rules.
     _lowercase							:      Dict  						=							{
         '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.
     _lowercase							:      Union[str, Any]  						=							None

     # Extract the layers.
     for key, val in transformer.items():
          # Match the name.
          _lowercase							:      List[Any]  						=							layer_re.match(lowerCAmelCase__)

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

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

          # The name of the layer.
          _lowercase							:      str  						=							F'''bert.encoder.layer.{layer_idx}'''

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

               _lowercase							:      List[str]  						=							fix_query_key_value_ordering(lowerCAmelCase__       ,  lowerCAmelCase__       ,  3       ,  lowerCAmelCase__       ,  lowerCAmelCase__)
               # Store the tensor as we need the bias as well to interleave QKV and biases.
               _lowercase							:      Optional[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":
               # 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.
               _lowercase							:      str  						=							attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
               _lowercase							:      Optional[int]  						=							attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
               _lowercase							:      List[Any]  						=							attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

               _lowercase							:      List[str]  						=							fix_query_key_value_ordering(lowerCAmelCase__       ,  lowerCAmelCase__       ,  3       ,  lowerCAmelCase__       ,  lowerCAmelCase__)
               # Split the bias.
               _lowercase							:      List[Any]  						=							out_val[0 * config.hidden_size : 1 * config.hidden_size]
               _lowercase							:      List[str]  						=							out_val[1 * config.hidden_size : 2 * config.hidden_size]
               _lowercase							:      List[str]  						=							out_val[2 * config.hidden_size : 3 * config.hidden_size]

               # Store.
               _lowercase							:      List[Any]  						=							q
               _lowercase							:      Optional[int]  						=							q_bias
               _lowercase							:      int  						=							k
               _lowercase							:      Any  						=							k_bias
               _lowercase							:      List[str]  						=							v
               _lowercase							:      Any  						=							v_bias

               # Clear the stored tensor.
               _lowercase							:      Tuple  						=							None

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

    # The final layernorm.
     _lowercase							:      List[Any]  						=							transformer['final_layernorm.weight']
     _lowercase							:      Any  						=							transformer['final_layernorm.bias']

     # The pooler.
     _lowercase							:      Dict  						=							lm['pooler']

     # Store the matrix and the bias.
     _lowercase							:      List[Any]  						=							pooler['dense.weight']
     _lowercase							:      int  						=							pooler['dense.bias']

     # The LM head from Megatron (for RACE).
     _lowercase							:      List[Any]  						=							model['lm_head']

     # The transform matrix.
     _lowercase							:      int  						=							lm_head['dense.weight']
     _lowercase							:      Optional[Any]  						=							lm_head['dense.bias']

     # The transform LN.
     _lowercase							:      Dict  						=							lm_head['layernorm.weight']
     _lowercase							:      List[Any]  						=							lm_head['layernorm.bias']

     # For the decoder, we replicate the weights.
     _lowercase							:      List[Any]  						=							word_embeddings
     _lowercase							:      Optional[Any]  						=							lm_head['bias']

     # The classifier from Megatron (for MLNI).
     _lowercase							:      Union[str, Any]  						=							model['binary_head']

     # Store the classifier.
     _lowercase							:      List[Any]  						=							binary_head['weight']
     _lowercase							:      Optional[Any]  						=							binary_head['bias']

     # It should be done!
     return output_state_dict

def 		SCREAMING_SNAKE_CASE      (					)			->						Tuple:

     '''simple docstring'''




     _lowercase							:      Dict  						=							argparse.ArgumentParser()
     parser.add_argument('--print-checkpoint-structure'       ,  action='store_true')
     parser.add_argument('path_to_checkpoint'       ,  type=lowerCAmelCase__       ,  help='Path to the ZIP file containing the checkpoint')
     parser.add_argument(
         '--config_file'       ,  default=''       ,  type=lowerCAmelCase__       ,  help='An optional config json file describing the pre-trained model.'       ,  )
     _lowercase							:      List[str]  						=							parser.parse_args()

     # Extract the basename.
     _lowercase							:      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:
                    _lowercase							:      Optional[Any]  						=							torch.load(lowerCAmelCase__       ,  map_location='cpu')
     else:
          _lowercase							:      Any  						=							torch.load(args.path_to_checkpoint       ,  map_location='cpu')

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

     # Convert.
     print('Converting')
     _lowercase							:      str  						=							convert_megatron_checkpoint(lowerCAmelCase__       ,  lowerCAmelCase__       ,  lowerCAmelCase__)

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

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

     # Store the state_dict to file.
     _lowercase							:      Dict  						=							os.path.join(lowerCAmelCase__       ,  'pytorch_model.bin')
     print(F'''Saving checkpoint to "{output_checkpoint_file}"''')
     torch.save(lowerCAmelCase__       ,  lowerCAmelCase__)


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

if __name__ == "__main__":
    main()

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









'''simple docstring'''







import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


A						    =							logging.get_logger(__name__)

A						    =							{
    '''facebook/sam-vit-huge''': '''https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json''',
    '''facebook/sam-vit-large''': '''https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json''',
    '''facebook/sam-vit-base''': '''https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json''',
}



class 							__SCREAMING_SNAKE_CASE			(       lowerCAmelCase__      ):
      '''simple docstring'''



      def __init__(							self    : Union[str, Any]   ,UpperCamelCase    : Tuple=256   ,UpperCamelCase    : Union[str, Any]=1024   ,UpperCamelCase    : Any=16   ,UpperCamelCase    : List[str]=16   ,UpperCamelCase    : List[str]=4   ,UpperCamelCase    : Dict="gelu"   ,UpperCamelCase    : Any=1e-6   ,**UpperCamelCase    : List[Any]   ,)		->  str:
           super().__init__(**UpperCamelCase  )
           _lowercase							:      int  						=							hidden_size
           _lowercase							:      List[str]  						=							image_size
           _lowercase							:      Optional[Any]  						=							patch_size
           _lowercase							:      Any  						=							image_size // patch_size
           _lowercase							:      str  						=							mask_input_channels
           _lowercase							:      Optional[Any]  						=							num_point_embeddings
           _lowercase							:      Dict  						=							hidden_act
           _lowercase							:      int  						=							layer_norm_eps



class 							__SCREAMING_SNAKE_CASE			(       lowerCAmelCase__      ):
      '''simple docstring'''



      def __init__(							self    : Any   ,UpperCamelCase    : Any=256   ,UpperCamelCase    : Any="relu"   ,UpperCamelCase    : Dict=2048   ,UpperCamelCase    : Dict=2   ,UpperCamelCase    : List[str]=8   ,UpperCamelCase    : int=2   ,UpperCamelCase    : List[Any]=3   ,UpperCamelCase    : Any=3   ,UpperCamelCase    : Any=256   ,UpperCamelCase    : Any=1e-6   ,**UpperCamelCase    : Union[str, Any]   ,)		->  Tuple:
           super().__init__(**UpperCamelCase  )
           _lowercase							:      str  						=							hidden_size
           _lowercase							:      Union[str, Any]  						=							hidden_act
           _lowercase							:      Optional[Any]  						=							mlp_dim
           _lowercase							:      Tuple  						=							num_hidden_layers
           _lowercase							:      Tuple  						=							num_attention_heads
           _lowercase							:      List[Any]  						=							attention_downsample_rate
           _lowercase							:      Optional[Any]  						=							num_multimask_outputs
           _lowercase							:      List[str]  						=							iou_head_depth
           _lowercase							:      List[str]  						=							iou_head_hidden_dim
           _lowercase							:      Union[str, Any]  						=							layer_norm_eps



class 							__SCREAMING_SNAKE_CASE			(       lowerCAmelCase__      ):
      '''simple docstring'''



      def __init__(							self    : Optional[Any]   ,UpperCamelCase    : int=768   ,UpperCamelCase    : str=256   ,UpperCamelCase    : str=12   ,UpperCamelCase    : int=12   ,UpperCamelCase    : Optional[int]=3   ,UpperCamelCase    : str=1024   ,UpperCamelCase    : Dict=16   ,UpperCamelCase    : Tuple="gelu"   ,UpperCamelCase    : List[str]=1e-06   ,UpperCamelCase    : List[str]=0.0   ,UpperCamelCase    : int=1e-10   ,UpperCamelCase    : List[Any]=True   ,UpperCamelCase    : Optional[int]=4.0   ,UpperCamelCase    : List[Any]=True   ,UpperCamelCase    : Optional[int]=True   ,UpperCamelCase    : Tuple=14   ,UpperCamelCase    : Tuple=[2, 5, 8, 11]   ,UpperCamelCase    : List[str]=128   ,UpperCamelCase    : List[Any]=None   ,**UpperCamelCase    : List[str]   ,)		->  Optional[int]:
           super().__init__(**UpperCamelCase  )

           _lowercase							:      Optional[int]  						=							hidden_size
           _lowercase							:      List[Any]  						=							output_channels
           _lowercase							:      List[Any]  						=							num_hidden_layers
           _lowercase							:      Tuple  						=							num_attention_heads
           _lowercase							:      Union[str, Any]  						=							num_channels
           _lowercase							:      Optional[Any]  						=							image_size
           _lowercase							:      List[str]  						=							patch_size
           _lowercase							:      List[Any]  						=							hidden_act
           _lowercase							:      Dict  						=							layer_norm_eps
           _lowercase							:      List[str]  						=							attention_dropout
           _lowercase							:      Union[str, Any]  						=							initializer_range
           _lowercase							:      Any  						=							qkv_bias
           _lowercase							:      Dict  						=							mlp_ratio
           _lowercase							:      Optional[int]  						=							use_abs_pos
           _lowercase							:      int  						=							use_rel_pos
           _lowercase							:      Dict  						=							window_size
           _lowercase							:      Dict  						=							global_attn_indexes
           _lowercase							:      str  						=							num_pos_feats
           _lowercase							:      List[Any]  						=							int(hidden_size * mlp_ratio  ) if mlp_dim is None else mlp_dim



class 							__SCREAMING_SNAKE_CASE			(       lowerCAmelCase__      ):
      '''simple docstring'''



      lowerCAmelCase__       :      Any 		=		"sam"
      lowerCAmelCase__       :      Optional[int] 		=		True

      def __init__(							self    : Optional[Any]   ,UpperCamelCase    : Optional[int]=None   ,UpperCamelCase    : List[str]=None   ,UpperCamelCase    : List[Any]=None   ,UpperCamelCase    : int=0.0_2   ,**UpperCamelCase    : Any   ,)		->  Tuple:
           super().__init__(**UpperCamelCase  )
           _lowercase							:      str  						=							vision_config if vision_config is not None else {}
           _lowercase							:      int  						=							prompt_encoder_config if prompt_encoder_config is not None else {}
           _lowercase							:      List[Any]  						=							mask_decoder_config if mask_decoder_config is not None else {}

           if isinstance(UpperCamelCase   ,UpperCamelCase  ):
                _lowercase							:      Optional[Any]  						=							vision_config.to_dict()
           if isinstance(UpperCamelCase   ,UpperCamelCase  ):
                _lowercase							:      Union[str, Any]  						=							prompt_encoder_config.to_dict()
           if isinstance(UpperCamelCase   ,UpperCamelCase  ):
                _lowercase							:      List[Any]  						=							mask_decoder_config.to_dict()

           _lowercase							:      Optional[int]  						=							SamVisionConfig(**UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							SamPromptEncoderConfig(**UpperCamelCase  )
           _lowercase							:      Union[str, Any]  						=							SamMaskDecoderConfig(**UpperCamelCase  )
           _lowercase							:      Tuple  						=							initializer_range



      def   _lowerCamelCase		(							self    : List[str]  )		->  Optional[int]:
           _lowercase							:      List[Any]  						=							copy.deepcopy(self.__dict__  )
           _lowercase							:      Any  						=							self.vision_config.to_dict()
           _lowercase							:      Dict  						=							self.prompt_encoder_config.to_dict()
           _lowercase							:      Union[str, Any]  						=							self.mask_decoder_config.to_dict()
           _lowercase							:      List[Any]  						=							self.__class__.model_type
           return output