





'''simple docstring'''




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:					Any							=				[
    """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:					Any							=				[
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     Dict:
 __UpperCAmelCase					=   {
     '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					=   int(re.match(r'.*layer_(\d*).*'       ,    _lowerCAmelCase    )[1]    )
 layer_number -= 3
 return F'h.{layer_number}.' + key







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







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase    )->     int:
 # Construct model
 if bloom_config_file == "":
  __UpperCAmelCase					=   BloomConfig()
 else:
  __UpperCAmelCase					=   BloomConfig.from_json_file(_lowerCAmelCase    )

 if shard_model:
  __UpperCAmelCase					=   os.listdir(_lowerCAmelCase    )
  __UpperCAmelCase					=   sorted(filter(lambda _lowerCAmelCase    : s.startswith('layer'    ) and "model_00" in s       ,    _lowerCAmelCase    )    )

  __UpperCAmelCase					=   {'weight_map': {}, 'metadata': {}}
  __UpperCAmelCase					=   0

  __UpperCAmelCase					=   None

  __UpperCAmelCase					=   BloomConfig()

  for j, file in enumerate(_lowerCAmelCase    ):
   print('Processing file: {}'.format(_lowerCAmelCase    )    )
   __UpperCAmelCase					=   None

   for i in range(_lowerCAmelCase    ):
    # load all TP files
    __UpperCAmelCase					=   file.replace('model_00'       ,    F'model_0{i}'    )
    __UpperCAmelCase					=   torch.load(os.path.join(_lowerCAmelCase       ,    _lowerCAmelCase    )       ,    map_location='cpu'    )

    # Rename keys in the transformers names
    __UpperCAmelCase					=   list(temp.keys()    )
    for key in keys:
     __UpperCAmelCase					=   temp.pop(_lowerCAmelCase    )

    if tensors is None:
     __UpperCAmelCase					=   temp
    else:
     for key in tensors.keys():
      if any(key.endswith(_lowerCAmelCase    ) 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					=   1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN    ) else 0
       # We concatenate these weights accross TP ranks
       __UpperCAmelCase					=   torch.cat([tensors[key], temp[key]]       ,    dim=_lowerCAmelCase    )

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

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

  __UpperCAmelCase					=   BloomConfig()
  __UpperCAmelCase					=   pytorch_dump_folder_path + '/' + CONFIG_NAME
  __UpperCAmelCase					=   total_size
  with open(_lowerCAmelCase       ,    'w'       ,    encoding='utf-8'    ) as f:
   f.write(config.to_json_string()    )
  with open(os.path.join(_lowerCAmelCase       ,    WEIGHTS_NAME + '.index.json'    )       ,    'w'       ,    encoding='utf-8'    ) as f:
   __UpperCAmelCase					=   json.dumps(_lowerCAmelCase       ,    indent=2       ,    sort_keys=_lowerCAmelCase    ) + '\n'
   f.write(_lowerCAmelCase    )
 else:
  __UpperCAmelCase					=   BloomModel(_lowerCAmelCase    )

  __UpperCAmelCase					=   os.listdir(_lowerCAmelCase    )
  __UpperCAmelCase					=   sorted(filter(lambda _lowerCAmelCase    : s.startswith('layer'    ) and "model_00" in s       ,    _lowerCAmelCase    )    )

  __UpperCAmelCase					=   None
  for i, file in enumerate(_lowerCAmelCase    ):
   __UpperCAmelCase					=   None
   for i in range(_lowerCAmelCase    ):
    # load all TP files
    __UpperCAmelCase					=   file.replace('model_00'       ,    F'model_0{i}'    )
    __UpperCAmelCase					=   torch.load(os.path.join(_lowerCAmelCase       ,    _lowerCAmelCase    )       ,    map_location='cpu'    )

    # Rename keys in the transformers names
    __UpperCAmelCase					=   list(temp.keys()    )
    for key in keys:
     __UpperCAmelCase					=   temp.pop(_lowerCAmelCase    )

    if tensors is None:
     __UpperCAmelCase					=   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(_lowerCAmelCase    ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH    ):
       tensors[key] += temp[key]
      else:
       # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
       __UpperCAmelCase					=   1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN    ) else 0
       # We concatenate these weights accross TP ranks
       __UpperCAmelCase					=   torch.cat([tensors[key], temp[key]]       ,    dim=_lowerCAmelCase    )

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

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

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

  # Save pytorch-model
  os.makedirs(_lowerCAmelCase       ,    exist_ok=_lowerCAmelCase    )
  __UpperCAmelCase					=   pytorch_dump_folder_path + '/' + WEIGHTS_NAME
  __UpperCAmelCase					=   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					=   model.to(config.torch_dtype    )
  torch.save(model.state_dict()       ,    _lowerCAmelCase    )
  print(F'Save configuration file to {pytorch_config_dump_path}'    )
  with open(_lowerCAmelCase       ,    '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:					Dict							=				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,
    )












'''simple docstring'''




import json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin






@require_tokenizers
class UpperCAmelCase							(		UpperCAmelCase_					,			unittest.TestCase							):
    _A :			List[str]						   =       MgpstrTokenizer
    _A :			Any						   =       False
    _A :			List[Any]						   =       {}
    _A :			int						   =       False



    def 					__lowerCamelCase						(  self					):
     super().setUp()

     # fmt: off
     __UpperCAmelCase					=   ['[GO]', '[s]', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
     # fmt: on
     __UpperCAmelCase					=   dict(zip(__A     ,			range(len(__A					)					)					)					)

     __UpperCAmelCase					=   os.path.join(self.tmpdirname     ,			VOCAB_FILES_NAMES['vocab_file']					)
     with open(self.vocab_file     ,			'w'     ,			encoding='utf-8'					) as fp:
      fp.write(json.dumps(__A					) + '\n'					)



    def 					__lowerCamelCase						(  self     ,			**__A					):
     return MgpstrTokenizer.from_pretrained(self.tmpdirname     ,			**__A					)



    def 					__lowerCamelCase						(  self     ,			__A					):
     __UpperCAmelCase					=   'tester'
     __UpperCAmelCase					=   'tester'
     return input_text, output_text



    @unittest.skip('MGP-STR always lower cases letters.'					)
    def 					__lowerCamelCase						(  self					):
     pass



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.get_tokenizers(do_lower_case=__A					)
     for tokenizer in tokenizers:
      with self.subTest(f'{tokenizer.__class__.__name__}'					):
       __UpperCAmelCase					=   '[SPECIAL_TOKEN]'

       tokenizer.add_special_tokens({'cls_token': special_token}					)
       __UpperCAmelCase					=   tokenizer.encode([special_token]     ,			add_special_tokens=__A					)
       self.assertEqual(len(__A					)     ,			1					)

       __UpperCAmelCase					=   tokenizer.decode(__A     ,			skip_special_tokens=__A					)
       self.assertTrue(special_token not in decoded					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.get_tokenizers()
     for tokenizer in tokenizers:
      with self.subTest(f'{tokenizer.__class__.__name__}'					):
       __UpperCAmelCase					,    __UpperCAmelCase					=   self.get_input_output_texts(__A					)

       __UpperCAmelCase					=   tokenizer.tokenize(__A					)
       __UpperCAmelCase					=   tokenizer.convert_tokens_to_ids(__A					)
       __UpperCAmelCase					=   tokenizer.encode(__A     ,			add_special_tokens=__A					)
       self.assertListEqual(__A     ,			__A					)

       __UpperCAmelCase					=   tokenizer.convert_ids_to_tokens(__A					)
       self.assertNotEqual(len(__A					)     ,			0					)
       __UpperCAmelCase					=   tokenizer.decode(__A					)
       self.assertIsInstance(__A     ,			__A					)

       self.assertEqual(text_a.replace(' '     ,			''					)     ,			__A					)



    @unittest.skip('MGP-STR tokenizer only handles one sequence.'					)
    def 					__lowerCamelCase						(  self					):
     pass




    @unittest.skip('inputs cannot be pretokenized in MgpstrTokenizer'					)
    def 					__lowerCamelCase						(  self					):
     pass





