





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



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



def 	A_       (  _lowerCAmelCase  )		->				Union[str, Any]:
 if dtype == torch.bool:
  return 1 / 8
 UpperCamelCase			:      Tuple 			=     re.search(r"[^\d](\d+)$"							,     str(UpperCAmelCase__  )  )
 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 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase  )		->				Optional[Any]:
 # Construct model
 if bloom_config_file == "":
  UpperCamelCase			:      Union[str, Any] 			=     BloomConfig()
 else:
  UpperCamelCase			:      int 			=     BloomConfig.from_json_file(UpperCAmelCase__  )

 if shard_model:
  UpperCamelCase			:      Optional[Any] 			=     os.listdir(UpperCAmelCase__  )
  UpperCamelCase			:      Union[str, Any] 			=     sorted(filter(lambda _lowerCAmelCase  : s.startswith("layer"  ) and "model_00" in s							,     UpperCAmelCase__  )  )

  UpperCamelCase			:      Union[str, Any] 			=     {"""weight_map""": {}, """metadata""": {}}
  UpperCamelCase			:      Optional[int] 			=     0

  UpperCamelCase			:      int 			=     None

  UpperCamelCase			:      str 			=     BloomConfig()

  for j, file in enumerate(UpperCAmelCase__  ):
   print("Processing file: {}".format(UpperCAmelCase__  )  )
   UpperCamelCase			:      Any 			=     None

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

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

    if tensors is None:
     UpperCamelCase			:      str 			=     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			:      Union[str, 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			:      Any 			=     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			:      Dict 			=     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			:      Optional[int] 			=     tensors[key]
    total_size += value.numel() * get_dtype_size(value.dtype  )
    if key not in index_dict["weight_map"]:
     UpperCamelCase			:      Optional[int] 			=     """pytorch_model_{}-of-{}.bin""".format(
         str(j + 1  ).zfill(5  )							,     str(len(UpperCAmelCase__  )  ).zfill(5  )  )

  UpperCamelCase			:      Dict 			=     BloomConfig()
  UpperCamelCase			:      int 			=     pytorch_dump_folder_path + """/""" + CONFIG_NAME
  UpperCamelCase			:      List[Any] 			=     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			:      List[Any] 			=     json.dumps(UpperCAmelCase__							,     indent=2							,     sort_keys=UpperCAmelCase__  ) + """\n"""
   f.write(UpperCAmelCase__  )
 else:
  UpperCamelCase			:      Optional[int] 			=     BloomModel(UpperCAmelCase__  )

  UpperCamelCase			:      int 			=     os.listdir(UpperCAmelCase__  )
  UpperCamelCase			:      Optional[int] 			=     sorted(filter(lambda _lowerCAmelCase  : s.startswith("layer"  ) and "model_00" in s							,     UpperCAmelCase__  )  )

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

    # Rename keys in the transformers names
    UpperCamelCase			:      int 			=     list(temp.keys()  )
    for key in keys:
     UpperCamelCase			:      int 			=     temp.pop(UpperCAmelCase__  )

    if tensors is None:
     UpperCamelCase			:      Dict 			=     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			:      Dict 			=     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[Any] 			=     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			:      int 			=     tensors[key] / pretraining_tp

   UpperCamelCase			:      Optional[int] 			=     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			:      Union[str, Any] 			=     set(other_keys.missing_keys  )
   else:
    UpperCamelCase			:      str 			=     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			:      Union[str, Any] 			=     pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
  UpperCamelCase			:      Any 			=     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			:      Union[str, Any] 			=     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__":
    __lowerCamelCase :  Union[str, Any]										=  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[str]										=  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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCamelCase :  Tuple										=  logging.get_logger(__name__)

__lowerCamelCase :  str										=  {
    """camembert-base""": """https://huggingface.co/camembert-base/resolve/main/config.json""",
    """umberto-commoncrawl-cased-v1""": (
        """https://huggingface.co/Musixmatch/umberto-commoncrawl-cased-v1/resolve/main/config.json"""
    ),
    """umberto-wikipedia-uncased-v1""": (
        """https://huggingface.co/Musixmatch/umberto-wikipedia-uncased-v1/resolve/main/config.json"""
    ),
}



class       A__     ( __snake_case    ):
      _UpperCAmelCase			:Union[str, Any]							=     'camembert'





      def __init__(				self ,							A_=3_0522 ,							A_=768 ,							A_=12 ,							A_=12 ,							A_=3072 ,							A_="gelu" ,							A_=0.1 ,							A_=0.1 ,							A_=512 ,							A_=2 ,							A_=0.02 ,							A_=1e-12 ,							A_=1 ,							A_=0 ,							A_=2 ,							A_="absolute" ,							A_=True ,							A_=None ,							**A_ ,							):
       '''simple docstring'''




       super().__init__(pad_token_id=A_ ,							bos_token_id=A_ ,							eos_token_id=A_ ,							**A_     )

       UpperCamelCase			:      List[str] 			=     vocab_size
       UpperCamelCase			:      Union[str, Any] 			=     hidden_size
       UpperCamelCase			:      Any 			=     num_hidden_layers
       UpperCamelCase			:      Union[str, Any] 			=     num_attention_heads
       UpperCamelCase			:      Dict 			=     hidden_act
       UpperCamelCase			:      str 			=     intermediate_size
       UpperCamelCase			:      str 			=     hidden_dropout_prob
       UpperCamelCase			:      Dict 			=     attention_probs_dropout_prob
       UpperCamelCase			:      Union[str, Any] 			=     max_position_embeddings
       UpperCamelCase			:      Optional[Any] 			=     type_vocab_size
       UpperCamelCase			:      int 			=     initializer_range
       UpperCamelCase			:      List[str] 			=     layer_norm_eps
       UpperCamelCase			:      Dict 			=     position_embedding_type
       UpperCamelCase			:      int 			=     use_cache
       UpperCamelCase			:      List[str] 			=     classifier_dropout



class       A__     ( __snake_case    ):





      @property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       if self.task == "multiple-choice":
        UpperCamelCase			:      Optional[int] 			=     {0: "batch", 1: "choice", 2: "sequence"}
       else:
        UpperCamelCase			:      str 			=     {0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
           ]     )
