




import json
import os
import shutil
import warnings
from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import List

from ..utils import logging
from . import BaseTransformersCLICommand


try:
     from cookiecutter.main import cookiecutter

     __a							     =					True
except ImportError:
     __a							     =					False

__a							     =					logging.get_logger(__name__)  # pylint: disable=invalid-name







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     return AddNewModelCommand(args.testing					,		args.testing_file					,		path=args.path )



class     __a(	_a  ):






    """simple docstring"""





    @staticmethod
    def 				a__ ( _SCREAMING_SNAKE_CASE    )					->						Tuple:
         UpperCAmelCase_     :				Tuple  =     parser.add_parser('''add-new-model'''    )
         add_new_model_parser.add_argument('''--testing'''  ,action='''store_true'''  ,help='''If in testing mode.'''    )
         add_new_model_parser.add_argument('''--testing_file'''  ,type=_SCREAMING_SNAKE_CASE  ,help='''Configuration file on which to run.'''    )
         add_new_model_parser.add_argument(
             '''--path'''  ,type=_SCREAMING_SNAKE_CASE  ,help='''Path to cookiecutter. Should only be used for testing purposes.'''    )
         add_new_model_parser.set_defaults(func=_SCREAMING_SNAKE_CASE    )





    def __init__( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE=None  ,*_SCREAMING_SNAKE_CASE    )					->						Optional[int]:
         UpperCAmelCase_     :				str  =     testing
         UpperCAmelCase_     :				str  =     testing_file
         UpperCAmelCase_     :				Dict  =     path





    def 				a__ ( self    )					->						Optional[int]:
         warnings.warn(
             '''The command `transformers-cli add-new-model` is deprecated and will be removed in v5 of Transformers. '''
             '''It is not actively maintained anymore, so might give a result that won\'t pass all tests and quality '''
             '''checks, you should use `transformers-cli add-new-model-like` instead.'''    )
         if not _has_cookiecutter:
              raise ImportError(
                  '''Model creation dependencies are required to use the `add_new_model` command. Install them by running '''
                  '''the following at the root of your `transformers` clone:\n\n\t$ pip install -e .[modelcreation]\n'''    )
         # Ensure that there is no other `cookiecutter-template-xxx` directory in the current working directory
         UpperCAmelCase_     :				List[Any]  =     [directory for directory in os.listdir() if '''cookiecutter-template-''' == directory[:22]]
         if len(_SCREAMING_SNAKE_CASE    ) > 0:
              raise ValueError(
                  '''Several directories starting with `cookiecutter-template-` in current working directory. '''
                  '''Please clean your directory by removing all folders starting with `cookiecutter-template-` or '''
                  '''change your working directory.'''    )

         UpperCAmelCase_     :				Union[str, Any]  =     (
             Path(_SCREAMING_SNAKE_CASE    ).parent.parent.parent.parent if self._path is None else Path(self._path    ).parent.parent
         )
         UpperCAmelCase_     :				int  =     path_to_transformer_root / '''templates''' / '''adding_a_new_model'''

         # Execute cookiecutter
         if not self._testing:
              cookiecutter(str(_SCREAMING_SNAKE_CASE    )    )
         else:
              with open(self._testing_file  ,'''r'''    ) as configuration_file:
                   UpperCAmelCase_     :				str  =     json.load(_SCREAMING_SNAKE_CASE    )

              cookiecutter(
                  str(path_to_cookiecutter if self._path is None else self._path    )  ,no_input=_SCREAMING_SNAKE_CASE  ,extra_context=_SCREAMING_SNAKE_CASE  ,)

         UpperCAmelCase_     :				str  =     [directory for directory in os.listdir() if '''cookiecutter-template-''' in directory[:22]][0]

         # Retrieve configuration
         with open(directory + '''/configuration.json'''  ,'''r'''    ) as configuration_file:
              UpperCAmelCase_     :				List[Any]  =     json.load(_SCREAMING_SNAKE_CASE    )

         UpperCAmelCase_     :				Tuple  =     configuration['''lowercase_modelname''']
         UpperCAmelCase_     :				Optional[int]  =     configuration['''generate_tensorflow_pytorch_and_flax''']
         os.remove(f'''{directory}/configuration.json'''    )

         UpperCAmelCase_     :				List[str]  =     '''PyTorch''' in generate_tensorflow_pytorch_and_flax
         UpperCAmelCase_     :				Optional[int]  =     '''TensorFlow''' in generate_tensorflow_pytorch_and_flax
         UpperCAmelCase_     :				List[str]  =     '''Flax''' in generate_tensorflow_pytorch_and_flax

         UpperCAmelCase_     :				Optional[int]  =     f'''{path_to_transformer_root}/src/transformers/models/{lowercase_model_name}'''
         os.makedirs(_SCREAMING_SNAKE_CASE  ,exist_ok=_SCREAMING_SNAKE_CASE    )
         os.makedirs(f'''{path_to_transformer_root}/tests/models/{lowercase_model_name}'''  ,exist_ok=_SCREAMING_SNAKE_CASE    )

         # Tests require submodules as they have parent imports
         with open(f'''{path_to_transformer_root}/tests/models/{lowercase_model_name}/__init__.py'''  ,'''w'''    ):
              pass

         shutil.move(
             f'''{directory}/__init__.py'''  ,f'''{model_dir}/__init__.py'''  ,)
         shutil.move(
             f'''{directory}/configuration_{lowercase_model_name}.py'''  ,f'''{model_dir}/configuration_{lowercase_model_name}.py'''  ,)

         def remove_copy_lines(_SCREAMING_SNAKE_CASE    ):
              with open(_SCREAMING_SNAKE_CASE  ,'''r'''    ) as f:
                   UpperCAmelCase_     :				Optional[int]  =     f.readlines()
              with open(_SCREAMING_SNAKE_CASE  ,'''w'''    ) as f:
                   for line in lines:
                        if "# Copied from transformers." not in line:
                             f.write(_SCREAMING_SNAKE_CASE    )

         if output_pytorch:
              if not self._testing:
                   remove_copy_lines(f'''{directory}/modeling_{lowercase_model_name}.py'''    )

              shutil.move(
                  f'''{directory}/modeling_{lowercase_model_name}.py'''  ,f'''{model_dir}/modeling_{lowercase_model_name}.py'''  ,)

              shutil.move(
                  f'''{directory}/test_modeling_{lowercase_model_name}.py'''  ,f'''{path_to_transformer_root}/tests/models/{lowercase_model_name}/test_modeling_{lowercase_model_name}.py'''  ,)
         else:
              os.remove(f'''{directory}/modeling_{lowercase_model_name}.py'''    )
              os.remove(f'''{directory}/test_modeling_{lowercase_model_name}.py'''    )

         if output_tensorflow:
              if not self._testing:
                   remove_copy_lines(f'''{directory}/modeling_tf_{lowercase_model_name}.py'''    )

              shutil.move(
                  f'''{directory}/modeling_tf_{lowercase_model_name}.py'''  ,f'''{model_dir}/modeling_tf_{lowercase_model_name}.py'''  ,)

              shutil.move(
                  f'''{directory}/test_modeling_tf_{lowercase_model_name}.py'''  ,f'''{path_to_transformer_root}/tests/models/{lowercase_model_name}/test_modeling_tf_{lowercase_model_name}.py'''  ,)
         else:
              os.remove(f'''{directory}/modeling_tf_{lowercase_model_name}.py'''    )
              os.remove(f'''{directory}/test_modeling_tf_{lowercase_model_name}.py'''    )

         if output_flax:
              if not self._testing:
                   remove_copy_lines(f'''{directory}/modeling_flax_{lowercase_model_name}.py'''    )

              shutil.move(
                  f'''{directory}/modeling_flax_{lowercase_model_name}.py'''  ,f'''{model_dir}/modeling_flax_{lowercase_model_name}.py'''  ,)

              shutil.move(
                  f'''{directory}/test_modeling_flax_{lowercase_model_name}.py'''  ,f'''{path_to_transformer_root}/tests/models/{lowercase_model_name}/test_modeling_flax_{lowercase_model_name}.py'''  ,)
         else:
              os.remove(f'''{directory}/modeling_flax_{lowercase_model_name}.py'''    )
              os.remove(f'''{directory}/test_modeling_flax_{lowercase_model_name}.py'''    )

         shutil.move(
             f'''{directory}/{lowercase_model_name}.md'''  ,f'''{path_to_transformer_root}/docs/source/en/model_doc/{lowercase_model_name}.md'''  ,)

         shutil.move(
             f'''{directory}/tokenization_{lowercase_model_name}.py'''  ,f'''{model_dir}/tokenization_{lowercase_model_name}.py'''  ,)

         shutil.move(
             f'''{directory}/tokenization_fast_{lowercase_model_name}.py'''  ,f'''{model_dir}/tokenization_{lowercase_model_name}_fast.py'''  ,)

         from os import fdopen, remove
         from shutil import copymode, move
         from tempfile import mkstemp

         def replace(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    ):
              # Create temp file
              UpperCAmelCase_,       UpperCAmelCase_     :				Any  =     mkstemp()
              UpperCAmelCase_     :				List[Any]  =     False
              with fdopen(_SCREAMING_SNAKE_CASE  ,'''w'''    ) as new_file:
                   with open(_SCREAMING_SNAKE_CASE    ) as old_file:
                        for line in old_file:
                             new_file.write(_SCREAMING_SNAKE_CASE    )
                             if line_to_copy_below in line:
                                  UpperCAmelCase_     :				List[str]  =     True
                                  for line_to_copy in lines_to_copy:
                                       new_file.write(_SCREAMING_SNAKE_CASE    )

              if not line_found:
                   raise ValueError(f'''Line {line_to_copy_below} was not found in file.'''    )

              # Copy the file permissions from the old file to the new file
              copymode(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )
              # Remove original file
              remove(_SCREAMING_SNAKE_CASE    )
              # Move new file
              move(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )

         def skip_units(_SCREAMING_SNAKE_CASE    ):
              return (
                  ("generating PyTorch" in line and not output_pytorch)
                  or ("generating TensorFlow" in line and not output_tensorflow)
                  or ("generating Flax" in line and not output_flax)
              )

         def replace_in_files(_SCREAMING_SNAKE_CASE    ):
              with open(_SCREAMING_SNAKE_CASE    ) as datafile:
                   UpperCAmelCase_     :				Optional[Any]  =     []
                   UpperCAmelCase_     :				Optional[Any]  =     False
                   UpperCAmelCase_     :				int  =     False
                   for line in datafile:
                        if "# To replace in: " in line and "##" not in line:
                             UpperCAmelCase_     :				List[Any]  =     line.split('''"'''    )[1]
                             UpperCAmelCase_     :				Tuple  =     skip_units(_SCREAMING_SNAKE_CASE    )
                        elif "# Below: " in line and "##" not in line:
                             UpperCAmelCase_     :				Optional[int]  =     line.split('''"'''    )[1]
                             UpperCAmelCase_     :				List[str]  =     skip_units(_SCREAMING_SNAKE_CASE    )
                        elif "# End." in line and "##" not in line:
                             if not skip_file and not skip_snippet:
                                  replace(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )

                             UpperCAmelCase_     :				Dict  =     []
                        elif "# Replace with" in line and "##" not in line:
                             UpperCAmelCase_     :				Dict  =     []
                        elif "##" not in line:
                             lines_to_copy.append(_SCREAMING_SNAKE_CASE    )

              remove(_SCREAMING_SNAKE_CASE    )

         replace_in_files(f'''{directory}/to_replace_{lowercase_model_name}.py'''    )
         os.rmdir(_SCREAMING_SNAKE_CASE    )





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







def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				List[Any]  =     {
         '''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_     :				Union[str, Any]  =     int(re.match(r'''.*layer_(\d*).*'''					,		_lowercase )[1] )
     layer_number -= 3
     return f'''h.{layer_number}.''' + key







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     if dtype == torch.bool:
          return 1 / 8
     UpperCAmelCase_     :				Any  =     re.search(r'''[^\d](\d+)$'''					,		str(_lowercase ) )
     if bit_search is None:
          raise ValueError(f'''`dtype` is not a valid dtype: {dtype}.''' )
     UpperCAmelCase_     :				Optional[int]  =     int(bit_search.groups()[0] )
     return bit_size // 8







def lowerCamelCase__  (     _lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase ):

     '''simple docstring'''


     if bloom_config_file == "":
          UpperCAmelCase_     :				Tuple  =     BloomConfig()
     else:
          UpperCAmelCase_     :				Optional[int]  =     BloomConfig.from_json_file(_lowercase )

     if shard_model:
          UpperCAmelCase_     :				Any  =     os.listdir(_lowercase )
          UpperCAmelCase_     :				Union[str, Any]  =     sorted(filter(lambda _lowercase : s.startswith('''layer''' ) and "model_00" in s					,		_lowercase ) )

          UpperCAmelCase_     :				Any  =     {'''weight_map''': {}, '''metadata''': {}}
          UpperCAmelCase_     :				List[str]  =     0

          UpperCAmelCase_     :				Any  =     None

          UpperCAmelCase_     :				Optional[int]  =     BloomConfig()

          for j, file in enumerate(_lowercase ):
               print('''Processing file: {}'''.format(_lowercase ) )
               UpperCAmelCase_     :				Optional[Any]  =     None

               for i in range(_lowercase ):
                    # load all TP files
                    UpperCAmelCase_     :				Tuple  =     file.replace('''model_00'''					,		f'''model_0{i}''' )
                    UpperCAmelCase_     :				Any  =     torch.load(os.path.join(_lowercase					,		_lowercase )					,		map_location='''cpu''' )

                    # Rename keys in the transformers names
                    UpperCAmelCase_     :				Dict  =     list(temp.keys() )
                    for key in keys:
                         UpperCAmelCase_     :				Union[str, Any]  =     temp.pop(_lowercase )

                    if tensors is None:
                         UpperCAmelCase_     :				Union[str, Any]  =     temp
                    else:
                         for key in tensors.keys():
                              if any(key.endswith(_lowercase ) 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_     :				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_     :				Tuple  =     torch.cat([tensors[key], temp[key]]					,		dim=_lowercase )

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

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

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

          UpperCAmelCase_     :				Tuple  =     os.listdir(_lowercase )
          UpperCAmelCase_     :				Union[str, Any]  =     sorted(filter(lambda _lowercase : s.startswith('''layer''' ) and "model_00" in s					,		_lowercase ) )

          UpperCAmelCase_     :				Any  =     None
          for i, file in enumerate(_lowercase ):
               UpperCAmelCase_     :				Optional[Any]  =     None
               for i in range(_lowercase ):
                    # load all TP files
                    UpperCAmelCase_     :				List[Any]  =     file.replace('''model_00'''					,		f'''model_0{i}''' )
                    UpperCAmelCase_     :				Optional[int]  =     torch.load(os.path.join(_lowercase					,		_lowercase )					,		map_location='''cpu''' )

                    # Rename keys in the transformers names
                    UpperCAmelCase_     :				str  =     list(temp.keys() )
                    for key in keys:
                         UpperCAmelCase_     :				Dict  =     temp.pop(_lowercase )

                    if tensors is None:
                         UpperCAmelCase_     :				int  =     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(_lowercase ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH ):
                                   tensors[key] += temp[key]
                              else:
                                   # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                   UpperCAmelCase_     :				Optional[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_     :				List[str]  =     torch.cat([tensors[key], temp[key]]					,		dim=_lowercase )

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

               UpperCAmelCase_     :				Tuple  =     model.load_state_dict(_lowercase					,		strict=_lowercase )
               assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
               if missing_keys is None:
                    UpperCAmelCase_     :				Union[str, Any]  =     set(other_keys.missing_keys )
               else:
                    UpperCAmelCase_     :				Dict  =     missing_keys.intersection(set(other_keys.missing_keys ) )

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

          # Save pytorch-model
          os.makedirs(_lowercase					,		exist_ok=_lowercase )
          UpperCAmelCase_     :				str  =     pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
          UpperCAmelCase_     :				Dict  =     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_     :				Optional[int]  =     model.to(config.torch_dtype )
          torch.save(model.state_dict()					,		_lowercase )
          print(f'''Save configuration file to {pytorch_config_dump_path}''' )
          with open(_lowercase					,		'''w'''					,		encoding='''utf-8''' ) as f:
               f.write(config.to_json_string() )


if __name__ == "__main__":
     __a							     =					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							     =					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,
     )