




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







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Optional[int]:



       """simple docstring"""

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







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	Tuple:



       """simple docstring"""

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







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_			,    lowercase_			,    lowercase_			,    lowercase_						)       ->	Optional[Any]:



       """simple docstring"""

       if bloom_config_file == "":
              A__							=   BloomConfig()
       else:
              A__							=   BloomConfig.from_json_file(lowercase_						)

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

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

              A__							=   None

              A__							=   BloomConfig()

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

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

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

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

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

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

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

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

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

                     A__							=   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:
                            A__							=   set(other_keys.missing_keys						)
                     else:
                            A__							=   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_						)
              A__							=   pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
              A__							=   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:
                     A__							=   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     :		Optional[int]								=       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,
  )








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


_lowerCamelCase     :		Optional[Any]								=       {
    """configuration_trocr""": ["""TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TrOCRConfig"""],
    """processing_trocr""": ["""TrOCRProcessor"""],
}


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCamelCase     :		Any								=       [
      """TROCR_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TrOCRForCausalLM""",
      """TrOCRPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
  from .processing_trocr import TrOCRProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
  import sys

  _lowerCamelCase     :		List[str]								=       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


