





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

UpperCamelCase__      =[
    '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',
]

UpperCamelCase__      =[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase	):
     _SCREAMING_SNAKE_CASE       :	Optional[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
     _SCREAMING_SNAKE_CASE       :	List[Any]						    =							int(re.match(R".*layer_(\d*).*",       __lowerCamelCase	)[1]	)
     layer_number -= 3
     return f"""h.{layer_number}.""" + key






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






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase	):
     # Construct model
     if bloom_config_file == "":
          _SCREAMING_SNAKE_CASE       :	int						    =							BloomConfig()
     else:
          _SCREAMING_SNAKE_CASE       :	List[str]						    =							BloomConfig.from_json_file(__lowerCamelCase	)

     if shard_model:
          _SCREAMING_SNAKE_CASE       :	List[str]						    =							os.listdir(__lowerCamelCase	)
          _SCREAMING_SNAKE_CASE       :	Any						    =							sorted(filter(lambda __lowerCamelCase	: s.startswith("layer"	) and "model_00" in s,       __lowerCamelCase	)	)

          _SCREAMING_SNAKE_CASE       :	Tuple						    =							{"weight_map": {}, "metadata": {}}
          _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							0

          _SCREAMING_SNAKE_CASE       :	Tuple						    =							None

          _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							BloomConfig()

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

               for i in range(__lowerCamelCase	):
                    # load all TP files
                    _SCREAMING_SNAKE_CASE       :	List[Any]						    =							file.replace("model_00",       f"""model_0{i}"""	)
                    _SCREAMING_SNAKE_CASE       :	Dict						    =							torch.load(os.path.join(__lowerCamelCase,       __lowerCamelCase	),       map_location="cpu"	)

                    # Rename keys in the transformers names
                    _SCREAMING_SNAKE_CASE       :	str						    =							list(temp.keys()	)
                    for key in keys:
                         _SCREAMING_SNAKE_CASE       :	List[Any]						    =							temp.pop(__lowerCamelCase	)

                    if tensors is None:
                         _SCREAMING_SNAKE_CASE       :	List[Any]						    =							temp
                    else:
                         for key in tensors.keys():
                              if any(key.endswith(__lowerCamelCase	) 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
                                   _SCREAMING_SNAKE_CASE       :	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
                                   _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							torch.cat([tensors[key], temp[key]],       dim=__lowerCamelCase	)

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

               for key in tensors.keys():
                    _SCREAMING_SNAKE_CASE       :	Dict						    =							tensors[key]
                    total_size += value.numel() * get_dtype_size(value.dtype	)
                    if key not in index_dict["weight_map"]:
                         _SCREAMING_SNAKE_CASE       :	Tuple						    =							"pytorch_model_{}-of-{}.bin".format(
                             str(j + 1	).zfill(5	),       str(len(__lowerCamelCase	)	).zfill(5	)	)

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

          _SCREAMING_SNAKE_CASE       :	Tuple						    =							os.listdir(__lowerCamelCase	)
          _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							sorted(filter(lambda __lowerCamelCase	: s.startswith("layer"	) and "model_00" in s,       __lowerCamelCase	)	)

          _SCREAMING_SNAKE_CASE       :	List[str]						    =							None
          for i, file in enumerate(__lowerCamelCase	):
               _SCREAMING_SNAKE_CASE       :	Dict						    =							None
               for i in range(__lowerCamelCase	):
                    # load all TP files
                    _SCREAMING_SNAKE_CASE       :	List[str]						    =							file.replace("model_00",       f"""model_0{i}"""	)
                    _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							torch.load(os.path.join(__lowerCamelCase,       __lowerCamelCase	),       map_location="cpu"	)

                    # Rename keys in the transformers names
                    _SCREAMING_SNAKE_CASE       :	List[str]						    =							list(temp.keys()	)
                    for key in keys:
                         _SCREAMING_SNAKE_CASE       :	str						    =							temp.pop(__lowerCamelCase	)

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

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

               _SCREAMING_SNAKE_CASE       :	int						    =							model.load_state_dict(__lowerCamelCase,       strict=__lowerCamelCase	)
               assert not other_keys.unexpected_keys, f"""The keys {other_keys.unexpected_keys} are unexpected"""
               if missing_keys is None:
                    _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							set(other_keys.missing_keys	)
               else:
                    _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							missing_keys.intersection(set(other_keys.missing_keys	)	)

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

          # Save pytorch-model
          os.makedirs(__lowerCamelCase,       exist_ok=__lowerCamelCase	)
          _SCREAMING_SNAKE_CASE       :	Dict						    =							pytorch_dump_folder_path + "/" + WEIGHTS_NAME
          _SCREAMING_SNAKE_CASE       :	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:
               _SCREAMING_SNAKE_CASE       :	Any						    =							model.to(config.torch_dtype	)
          torch.save(model.state_dict(),       __lowerCamelCase	)
          print(f"""Save configuration file to {pytorch_config_dump_path}"""	)
          with open(__lowerCamelCase,       "w",       encoding="utf-8"	) as f:
               f.write(config.to_json_string()	)


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


UpperCamelCase__      ={
    'configuration_clipseg': [
        'CLIPSEG_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'CLIPSegConfig',
        'CLIPSegTextConfig',
        'CLIPSegVisionConfig',
    ],
    'processing_clipseg': ['CLIPSegProcessor'],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase__      =[
           'CLIPSEG_PRETRAINED_MODEL_ARCHIVE_LIST',
           'CLIPSegModel',
           'CLIPSegPreTrainedModel',
           'CLIPSegTextModel',
           'CLIPSegVisionModel',
           'CLIPSegForImageSegmentation',
       ]

if TYPE_CHECKING:
       from .configuration_clipseg import (
           CLIPSEG_PRETRAINED_CONFIG_ARCHIVE_MAP,
           CLIPSegConfig,
           CLIPSegTextConfig,
           CLIPSegVisionConfig,
       )
       from .processing_clipseg import CLIPSegProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_clipseg import (
                  CLIPSEG_PRETRAINED_MODEL_ARCHIVE_LIST,
                  CLIPSegForImageSegmentation,
                  CLIPSegModel,
                  CLIPSegPreTrainedModel,
                  CLIPSegTextModel,
                  CLIPSegVisionModel,
              )

else:
       import sys

       UpperCamelCase__      =_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)