






import unittest

from transformers import MPNetConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       MPNetForMaskedLM,
       MPNetForMultipleChoice,
       MPNetForQuestionAnswering,
       MPNetForSequenceClassification,
       MPNetForTokenClassification,
       MPNetModel,
   )




class     _lowerCAmelCase:





       """simple docstring"""







       def __init__( self  ,      _lowerCamelCase  ,      _lowerCamelCase=1_3  ,      _lowerCamelCase=7  ,      _lowerCamelCase=True  ,      _lowerCamelCase=True  ,      _lowerCamelCase=False  ,      _lowerCamelCase=True  ,      _lowerCamelCase=9_9  ,      _lowerCamelCase=6_4  ,      _lowerCamelCase=5  ,      _lowerCamelCase=4  ,      _lowerCamelCase=6_4  ,      _lowerCamelCase="gelu"  ,      _lowerCamelCase=0.1  ,      _lowerCamelCase=0.1  ,      _lowerCamelCase=5_1_2  ,      _lowerCamelCase=1_6  ,      _lowerCamelCase=2  ,      _lowerCamelCase=0.0_2  ,      _lowerCamelCase=3  ,      _lowerCamelCase=4  ,      _lowerCamelCase=None  ,      ):
         UpperCamelCase_: Optional[int]   				=    parent
         UpperCamelCase_: Dict   				=    batch_size
         UpperCamelCase_: List[Any]   				=    seq_length
         UpperCamelCase_: List[str]   				=    is_training
         UpperCamelCase_: List[str]   				=    use_input_mask
         UpperCamelCase_: List[str]   				=    use_token_type_ids
         UpperCamelCase_: Optional[Any]   				=    use_labels
         UpperCamelCase_: List[Any]   				=    vocab_size
         UpperCamelCase_: Optional[Any]   				=    hidden_size
         UpperCamelCase_: Any   				=    num_hidden_layers
         UpperCamelCase_: Union[str, Any]   				=    num_attention_heads
         UpperCamelCase_: Tuple   				=    intermediate_size
         UpperCamelCase_: List[Any]   				=    hidden_act
         UpperCamelCase_: Union[str, Any]   				=    hidden_dropout_prob
         UpperCamelCase_: List[str]   				=    attention_probs_dropout_prob
         UpperCamelCase_: List[Any]   				=    max_position_embeddings
         UpperCamelCase_: Optional[Any]   				=    type_vocab_size
         UpperCamelCase_: Optional[Any]   				=    type_sequence_label_size
         UpperCamelCase_: Dict   				=    initializer_range
         UpperCamelCase_: Optional[int]   				=    num_labels
         UpperCamelCase_: List[str]   				=    num_choices
         UpperCamelCase_: Dict   				=    scope





       def       _a   ( self							):
         return MPNetConfig.from_pretrained('microsoft/mpnet-base'							)





       def       _a   ( self							):
         UpperCamelCase_: List[Any]   				=    ids_tensor([self.batch_size, self.seq_length]  ,      self.vocab_size							)

         UpperCamelCase_: Optional[int]   				=    None
         if self.use_input_mask:
           UpperCamelCase_: str   				=    random_attention_mask([self.batch_size, self.seq_length]							)

         UpperCamelCase_: Dict   				=    None
         UpperCamelCase_: Dict   				=    None
         UpperCamelCase_: Union[str, Any]   				=    None
         if self.use_labels:
           UpperCamelCase_: str   				=    ids_tensor([self.batch_size]  ,      self.type_sequence_label_size							)
           UpperCamelCase_: Tuple   				=    ids_tensor([self.batch_size, self.seq_length]  ,      self.num_labels							)
           UpperCamelCase_: Union[str, Any]   				=    ids_tensor([self.batch_size]  ,      self.num_choices							)

         UpperCamelCase_: Optional[int]   				=    self.get_config()
         return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels





       def       _a   ( self							):
         return MPNetConfig(
             vocab_size=self.vocab_size  ,      hidden_size=self.hidden_size  ,      num_hidden_layers=self.num_hidden_layers  ,      num_attention_heads=self.num_attention_heads  ,      intermediate_size=self.intermediate_size  ,      hidden_act=self.hidden_act  ,      hidden_dropout_prob=self.hidden_dropout_prob  ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,      max_position_embeddings=self.max_position_embeddings  ,      initializer_range=self.initializer_range  ,      )





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: Optional[int]   				=    MPNetModel(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         UpperCamelCase_: Any   				=    model(_lowerCamelCase  ,      _lowerCamelCase							)
         UpperCamelCase_: Tuple   				=    model(_lowerCamelCase							)
         self.parent.assertEqual(result.last_hidden_state.shape  ,      (self.batch_size, self.seq_length, self.hidden_size)							)
         self.parent.assertEqual(result.pooler_output.shape  ,      (self.batch_size, self.hidden_size)							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: Optional[Any]   				=    MPNetForQuestionAnswering(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         UpperCamelCase_: str   				=    model(
             _lowerCamelCase  ,      attention_mask=_lowerCamelCase  ,      start_positions=_lowerCamelCase  ,      end_positions=_lowerCamelCase  ,      )
         self.parent.assertEqual(result.start_logits.shape  ,      (self.batch_size, self.seq_length)							)
         self.parent.assertEqual(result.end_logits.shape  ,      (self.batch_size, self.seq_length)							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: str   				=    self.num_labels
         UpperCamelCase_: int   				=    MPNetForSequenceClassification(_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         UpperCamelCase_: List[str]   				=    model(_lowerCamelCase  ,      attention_mask=_lowerCamelCase  ,      labels=_lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.num_labels)							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: List[Any]   				=    self.num_choices
         UpperCamelCase_: Union[str, Any]   				=    MPNetForMultipleChoice(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         UpperCamelCase_: Union[str, Any]   				=    input_ids.unsqueeze(1							).expand(-1  ,      self.num_choices  ,      -1							).contiguous()
         UpperCamelCase_: Optional[int]   				=    input_mask.unsqueeze(1							).expand(-1  ,      self.num_choices  ,      -1							).contiguous()
         UpperCamelCase_: Optional[Any]   				=    model(
             _lowerCamelCase  ,      attention_mask=_lowerCamelCase  ,      labels=_lowerCamelCase  ,      )
         self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.num_choices)							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: Dict   				=    self.num_labels
         UpperCamelCase_: Dict   				=    MPNetForTokenClassification(config=_lowerCamelCase							)
         model.to(_lowerCamelCase							)
         model.eval()
         UpperCamelCase_: Tuple   				=    model(_lowerCamelCase  ,      attention_mask=_lowerCamelCase  ,      labels=_lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.seq_length, self.num_labels)							)





       def       _a   ( self							):
         UpperCamelCase_: int   				=    self.prepare_config_and_inputs()
         ((UpperCamelCase_)				,(UpperCamelCase_)				,(UpperCamelCase_)				,(UpperCamelCase_)				,(UpperCamelCase_)				,(UpperCamelCase_)): List[Any]   				=    config_and_inputs
         UpperCamelCase_: int   				=    {'input_ids': input_ids, 'attention_mask': input_mask}
         return config, inputs_dict







@require_torch
class     _lowerCAmelCase(	UpperCAmelCase_							,     UpperCAmelCase_							,     unittest.TestCase ):





       """simple docstring"""







       a    :		Dict									=(
           (
               MPNetForMaskedLM,
               MPNetForMultipleChoice,
               MPNetForQuestionAnswering,
               MPNetForSequenceClassification,
               MPNetForTokenClassification,
               MPNetModel,
           )
           if is_torch_available()
           else ()
       )
       a    :		Tuple									=(
           {
               '''feature-extraction''': MPNetModel,
               '''fill-mask''': MPNetForMaskedLM,
               '''question-answering''': MPNetForQuestionAnswering,
               '''text-classification''': MPNetForSequenceClassification,
               '''token-classification''': MPNetForTokenClassification,
               '''zero-shot''': MPNetForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       a    :		str									=False
       a    :		Union[str, Any]									=True





       def       _a   ( self							):
         UpperCamelCase_: Any   				=    MPNetModelTester(self							)
         UpperCamelCase_: List[Any]   				=    ConfigTester(self  ,      config_class=_lowerCamelCase  ,      hidden_size=3_7							)





       def       _a   ( self							):
         self.config_tester.run_common_tests()





       def       _a   ( self							):
         UpperCamelCase_: Tuple   				=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mpnet_model(*_lowerCamelCase							)





       def       _a   ( self							):
         UpperCamelCase_: List[str]   				=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mpnet_for_sequence_classification(*_lowerCamelCase							)





       def       _a   ( self							):
         UpperCamelCase_: Dict   				=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mpnet_for_multiple_choice(*_lowerCamelCase							)





       def       _a   ( self							):
         UpperCamelCase_: int   				=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mpnet_for_token_classification(*_lowerCamelCase							)





       def       _a   ( self							):
         UpperCamelCase_: int   				=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mpnet_for_question_answering(*_lowerCamelCase							)







@require_torch
class     _lowerCAmelCase(	unittest.TestCase ):





       """simple docstring"""







       @slow
       def       _a   ( self							):
         UpperCamelCase_: Optional[Any]   				=    MPNetModel.from_pretrained('microsoft/mpnet-base'							)
         UpperCamelCase_: Union[str, Any]   				=    torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]							)
         UpperCamelCase_: Tuple   				=    model(_lowerCamelCase							)[0]
         UpperCamelCase_: Optional[int]   				=    torch.Size((1, 1_1, 7_6_8)							)
         self.assertEqual(output.shape  ,      _lowerCamelCase							)
         UpperCamelCase_: str   				=    torch.tensor(
             [[[-0.0_5_5_0, 0.1_9_4_3, -0.0_7_4_0], [-0.0_5_6_2, 0.2_2_1_1, -0.0_5_7_9], [-0.0_4_3_7, 0.3_3_3_7, -0.0_6_4_1]]]							)
         # compare the actual values for a slice.
         self.assertTrue(torch.allclose(output[:, :3, :3]  ,      _lowerCamelCase  ,      atol=1e-4							)							)







import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

A_		:		str		     =			[
    'word_embeddings_layernorm.weight',
    'word_embeddings_layernorm.bias',
    'input_layernorm.weight',
    'input_layernorm.bias',
    'post_attention_layernorm.weight',
    'post_attention_layernorm.bias',
    'self_attention.dense.bias',
    'mlp.dense_4h_to_h.bias',
    'ln_f.weight',
    'ln_f.bias',
]

A_		:		Optional[int]		     =			[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]

def   snake_case (UpperCAmelCase__     ,   UpperCAmelCase__  )	->    List[str]:
  UpperCamelCase_: Dict   				=    {
      'word_embeddings.weight': 'word_embeddings.weight',
      'word_embeddings.norm.weight': 'word_embeddings_layernorm.weight',
      'word_embeddings.norm.bias': 'word_embeddings_layernorm.bias',
      'weight': 'ln_f.weight',
      'bias': 'ln_f.bias',
  }

  if key in layer_rename_map:
    return layer_rename_map[key]

  # Handle transformer blocks
  UpperCamelCase_: Tuple   				=    int(re.match(R'.*layer_(\d*).*'     ,   UpperCAmelCase__  )[1]  )
  layer_number -= 3
  return F'''h.{layer_number}.''' + key

def   snake_case (UpperCAmelCase__  )	->    List[str]:
  if dtype == torch.bool:
    return 1 / 8
  UpperCamelCase_: Optional[Any]   				=    re.search(R'[^\d](\d+)$'     ,   str(UpperCAmelCase__  )  )
  if bit_search is None:
    raise ValueError(F'''`dtype` is not a valid dtype: {dtype}.'''  )
  UpperCamelCase_: List[Any]   				=    int(bit_search.groups()[0]  )
  return bit_size // 8

def   snake_case (UpperCAmelCase__     ,   UpperCAmelCase__     ,   UpperCAmelCase__     ,   UpperCAmelCase__     ,   UpperCAmelCase__  )	->    Any:
  # Construct model
  if bloom_config_file == "":
    UpperCamelCase_: List[str]   				=    BloomConfig()
  else:
    UpperCamelCase_: List[str]   				=    BloomConfig.from_json_file(UpperCAmelCase__  )

  if shard_model:
    UpperCamelCase_: str   				=    os.listdir(UpperCAmelCase__  )
    UpperCamelCase_: List[str]   				=    sorted(filter(lambda UpperCAmelCase__  : s.startswith('layer'  ) and "model_00" in s     ,   UpperCAmelCase__  )  )

    UpperCamelCase_: Optional[int]   				=    {'weight_map': {}, 'metadata': {}}
    UpperCamelCase_: str   				=    0

    UpperCamelCase_: Optional[Any]   				=    None

    UpperCamelCase_: int   				=    BloomConfig()

    for j, file in enumerate(UpperCAmelCase__  ):
      print('Processing file: {}'.format(UpperCAmelCase__  )  )
      UpperCamelCase_: Tuple   				=    None

      for i in range(UpperCAmelCase__  ):
        # load all TP files
        UpperCamelCase_: List[Any]   				=    file.replace('model_00'     ,   F'''model_0{i}'''  )
        UpperCamelCase_: List[str]   				=    torch.load(os.path.join(UpperCAmelCase__     ,   UpperCAmelCase__  )     ,   map_location='cpu'  )

        # Rename keys in the transformers names
        UpperCamelCase_: Optional[int]   				=    list(temp.keys()  )
        for key in keys:
          UpperCamelCase_: List[Any]   				=    temp.pop(UpperCAmelCase__  )

        if tensors is None:
          UpperCamelCase_: Dict   				=    temp
        else:
          for key in tensors.keys():
            if any(key.endswith(UpperCAmelCase__  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
              # We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
              tensors[key] += temp[key]
            else:
              # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
              UpperCamelCase_: List[Any]   				=    1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
              # We concatenate these weights accross TP ranks
              UpperCamelCase_: Dict   				=    torch.cat([tensors[key], temp[key]]     ,   dim=UpperCAmelCase__  )

            # Divide by the number of TP the weights we want to average
      for key in tensors.keys():
        if any(key.endswith(UpperCAmelCase__  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
          UpperCamelCase_: Optional[int]   				=    tensors[key] / pretraining_tp
      torch.save(
          UpperCAmelCase__     ,   os.path.join(
              UpperCAmelCase__     ,   'pytorch_model_{}-of-{}.bin'.format(str(j + 1  ).zfill(5  )     ,   str(len(UpperCAmelCase__  )  ).zfill(5  )  )     ,   )     ,   )

      for key in tensors.keys():
        UpperCamelCase_: int   				=    tensors[key]
        total_size += value.numel() * get_dtype_size(value.dtype  )
        if key not in index_dict["weight_map"]:
          UpperCamelCase_: Dict   				=    'pytorch_model_{}-of-{}.bin'.format(
              str(j + 1  ).zfill(5  )     ,   str(len(UpperCAmelCase__  )  ).zfill(5  )  )

    UpperCamelCase_: Union[str, Any]   				=    BloomConfig()
    UpperCamelCase_: Any   				=    pytorch_dump_folder_path + '/' + CONFIG_NAME
    UpperCamelCase_: Optional[int]   				=    total_size
    with open(UpperCAmelCase__     ,   'w'     ,   encoding='utf-8'  ) as f:
      f.write(config.to_json_string()  )
    with open(os.path.join(UpperCAmelCase__     ,   WEIGHTS_NAME + '.index.json'  )     ,   'w'     ,   encoding='utf-8'  ) as f:
      UpperCamelCase_: Tuple   				=    json.dumps(UpperCAmelCase__     ,   indent=2     ,   sort_keys=UpperCAmelCase__  ) + '\n'
      f.write(UpperCAmelCase__  )
  else:
    UpperCamelCase_: Optional[Any]   				=    BloomModel(UpperCAmelCase__  )

    UpperCamelCase_: Tuple   				=    os.listdir(UpperCAmelCase__  )
    UpperCamelCase_: Tuple   				=    sorted(filter(lambda UpperCAmelCase__  : s.startswith('layer'  ) and "model_00" in s     ,   UpperCAmelCase__  )  )

    UpperCamelCase_: Tuple   				=    None
    for i, file in enumerate(UpperCAmelCase__  ):
      UpperCamelCase_: Union[str, Any]   				=    None
      for i in range(UpperCAmelCase__  ):
        # load all TP files
        UpperCamelCase_: Any   				=    file.replace('model_00'     ,   F'''model_0{i}'''  )
        UpperCamelCase_: Union[str, Any]   				=    torch.load(os.path.join(UpperCAmelCase__     ,   UpperCAmelCase__  )     ,   map_location='cpu'  )

        # Rename keys in the transformers names
        UpperCamelCase_: Dict   				=    list(temp.keys()  )
        for key in keys:
          UpperCamelCase_: Any   				=    temp.pop(UpperCAmelCase__  )

        if tensors is None:
          UpperCamelCase_: Any   				=    temp
        else:
          for key in tensors.keys():
            # We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
            if any(key.endswith(UpperCAmelCase__  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
              tensors[key] += temp[key]
            else:
              # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
              UpperCamelCase_: int   				=    1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
              # We concatenate these weights accross TP ranks
              UpperCamelCase_: Optional[int]   				=    torch.cat([tensors[key], temp[key]]     ,   dim=UpperCAmelCase__  )

            # Divide by the number of TP the weights we want to average
      for key in tensors.keys():
        if any(key.endswith(UpperCAmelCase__  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
          UpperCamelCase_: Tuple   				=    tensors[key] / pretraining_tp

      UpperCamelCase_: Any   				=    model.load_state_dict(UpperCAmelCase__     ,   strict=UpperCAmelCase__  )
      assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
      if missing_keys is None:
        UpperCamelCase_: Any   				=    set(other_keys.missing_keys  )
      else:
        UpperCamelCase_: int   				=    missing_keys.intersection(set(other_keys.missing_keys  )  )

    assert not missing_keys, F'''The keys {missing_keys} are missing'''

    # Save pytorch-model
    os.makedirs(UpperCAmelCase__     ,   exist_ok=UpperCAmelCase__  )
    UpperCamelCase_: Optional[Any]   				=    pytorch_dump_folder_path + '/' + WEIGHTS_NAME
    UpperCamelCase_: str   				=    pytorch_dump_folder_path + '/' + CONFIG_NAME
    print(F'''Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}'''  )
    if config.torch_dtype is not None:
      UpperCamelCase_: Tuple   				=    model.to(config.torch_dtype  )
    torch.save(model.state_dict()     ,   UpperCAmelCase__  )
    print(F'''Save configuration file to {pytorch_config_dump_path}'''  )
    with open(UpperCAmelCase__     ,   'w'     ,   encoding='utf-8'  ) as f:
      f.write(config.to_json_string()  )


if __name__ == "__main__":
   A_		:		Dict		     =			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--bloom_checkpoint_path',
       default=None,
       type=str,
       required=True,
       help='Path to the Megatron-LM checkpoint path.',
   )
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
   )
   parser.add_argument(
       '--bloom_config_file',
       default='',
       type=str,
       help=(
           'An optional config json file corresponding to the pre-trained model. \n'
           'This specifies the model architecture.'
       ),
   )
   parser.add_argument(
       '--shard_model',
       action='store_true',
       help='An optional setting to shard the output model \nThis enables sharding the converted checkpoint',
   )
   parser.add_argument(
       '--pretraining_tp',
       default=4,
       type=int,
       help='Pretraining TP rank that has been used when training the model in Megatron-LM \n',
   )
   A_		:		Any		     =			parser.parse_args()
   convert_bloom_checkpoint_to_pytorch(
       args.bloom_checkpoint_path,
       args.bloom_config_file,
       args.pytorch_dump_folder_path,
       args.shard_model,
       args.pretraining_tp,
   )