




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()

lowerCamelCase				:		Optional[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",
]

lowerCamelCase				:		List[Any]     =   [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Optional[int]	,							lowercase	:       Union[str, Any]     ):


      '''simple docstring'''



      lowerCamelCase_		     =    {
          '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
      lowerCamelCase_		     =    int(re.match(r'.*layer_(\d*).*'	,							lowercase     )[1]     )
      layer_number -= 3
      return f"""h.{layer_number}.""" + key

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Optional[Any]     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Tuple	,							lowercase	:       Dict	,							lowercase	:       Optional[int]	,							lowercase	:       str	,							lowercase	:       Dict     ):


      '''simple docstring'''



      if bloom_config_file == "":
            lowerCamelCase_		     =    BloomConfig()
      else:
            lowerCamelCase_		     =    BloomConfig.from_json_file(lowercase     )

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

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

            lowerCamelCase_		     =    None

            lowerCamelCase_		     =    BloomConfig()

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

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

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

                        if tensors is None:
                              lowerCamelCase_		     =    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
                                          lowerCamelCase_		     =    1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN     ) else 0
                                          # We concatenate these weights accross TP ranks
                                          lowerCamelCase_		     =    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     ):
                              lowerCamelCase_		     =    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():
                        lowerCamelCase_		     =    tensors[key]
                        total_size += value.numel() * get_dtype_size(value.dtype     )
                        if key not in index_dict["weight_map"]:
                              lowerCamelCase_		     =    'pytorch_model_{}-of-{}.bin'.format(
                                  str(j + 1     ).zfill(5     )	,							str(len(lowercase     )     ).zfill(5     )     )

            lowerCamelCase_		     =    BloomConfig()
            lowerCamelCase_		     =    pytorch_dump_folder_path + '/' + CONFIG_NAME
            lowerCamelCase_		     =    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:
                  lowerCamelCase_		     =    json.dumps(lowercase	,							indent=2	,							sort_keys=lowercase     ) + '\n'
                  f.write(lowercase     )
      else:
            lowerCamelCase_		     =    BloomModel(lowercase     )

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

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

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

                        if tensors is None:
                              lowerCamelCase_		     =    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
                                          lowerCamelCase_		     =    1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN     ) else 0
                                          # We concatenate these weights accross TP ranks
                                          lowerCamelCase_		     =    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     ):
                              lowerCamelCase_		     =    tensors[key] / pretraining_tp

                  lowerCamelCase_		     =    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:
                        lowerCamelCase_		     =    set(other_keys.missing_keys     )
                  else:
                        lowerCamelCase_		     =    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     )
            lowerCamelCase_		     =    pytorch_dump_folder_path + '/' + WEIGHTS_NAME
            lowerCamelCase_		     =    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:
                  lowerCamelCase_		     =    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__":
 lowerCamelCase				:		str     =   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",
 )
 lowerCamelCase				:		List[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,
 )









import warnings

from ...utils import logging
from .image_processing_donut import DonutImageProcessor


lowerCamelCase				:		Any     =   logging.get_logger(__name__)







class   A(						UpperCamelCase   ):



 '''simple docstring'''


 def __init__(				self					:  Dict						,							*A_					:  Optional[int]						,							**A_					:  int  )   ->	None:

       """simple docstring"""






       warnings.warn(
           'The class DonutFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'
           ' use DonutImageProcessor instead.'						,							A_						,							)
       super().__init__(*A_						,							**A_  )



