




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 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_      )       ->	Dict:

    """simple docstring"""


    __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*).*'''					,     lowercase_      )[1]      )
    layer_number -= 3
    return F"h.{layer_number}." + key




def 					__SCREAMING_SNAKE_CASE   (  lowercase_      )       ->	Union[str, Any]:

    """simple docstring"""


    if dtype == torch.bool:
        return 1 / 8
    __UpperCamelCase  		=						re.search(r'''[^\d](\d+)$'''					,     str(lowercase_      )      )
    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 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_					,     lowercase_					,     lowercase_					,     lowercase_      )       ->	Any:

    """simple docstring"""


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

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

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

        __UpperCamelCase  		=						None

        __UpperCamelCase  		=						BloomConfig()

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

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

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

                if tensors is None:
                    __UpperCamelCase  		=						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  		=						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=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  		=						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  		=						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(lowercase_      )      ).zfill(5      )      )

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

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

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

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

                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(lowercase_      ) 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=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  		=						tensors[key] / pretraining_tp

            __UpperCamelCase  		=						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  		=						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(lowercase_					,     exist_ok=lowercase_      )
        __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()					,     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,
 )












from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a_   		=     {
    "configuration_longformer": [
        "LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "LongformerConfig",
        "LongformerOnnxConfig",
    ],
    "tokenization_longformer": ["LongformerTokenizer"],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a_   		=     ["LongformerTokenizerFast"]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a_   		=     [
     "LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
     "LongformerForMaskedLM",
     "LongformerForMultipleChoice",
     "LongformerForQuestionAnswering",
     "LongformerForSequenceClassification",
     "LongformerForTokenClassification",
     "LongformerModel",
     "LongformerPreTrainedModel",
     "LongformerSelfAttention",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a_   		=     [
     "TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
     "TFLongformerForMaskedLM",
     "TFLongformerForMultipleChoice",
     "TFLongformerForQuestionAnswering",
     "TFLongformerForSequenceClassification",
     "TFLongformerForTokenClassification",
     "TFLongformerModel",
     "TFLongformerPreTrainedModel",
     "TFLongformerSelfAttention",
 ]


if TYPE_CHECKING:
 from .configuration_longformer import (
     LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
     LongformerConfig,
     LongformerOnnxConfig,
 )
 from .tokenization_longformer import LongformerTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_longformer_fast import LongformerTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_longformer import (
      LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      LongformerForMaskedLM,
      LongformerForMultipleChoice,
      LongformerForQuestionAnswering,
      LongformerForSequenceClassification,
      LongformerForTokenClassification,
      LongformerModel,
      LongformerPreTrainedModel,
      LongformerSelfAttention,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_longformer import (
      TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFLongformerForMaskedLM,
      TFLongformerForMultipleChoice,
      TFLongformerForQuestionAnswering,
      TFLongformerForSequenceClassification,
      TFLongformerForTokenClassification,
      TFLongformerModel,
      TFLongformerPreTrainedModel,
      TFLongformerSelfAttention,
  )

else:
 import sys

 a_   		=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






