


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






def 						a__			(				A__,      A__				):
			SCREAMING_SNAKE_CASE_		:     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
			SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					int(re.match(r'.*layer_(\d*).*',      A__				)[1]				)
			layer_number -= 3
			return F'''h.{layer_number}.''' + key






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






def 						a__			(				A__,      A__,      A__,      A__,      A__				):
			# Construct model
			if bloom_config_file == "":
						SCREAMING_SNAKE_CASE_		:     Optional[int]					=					BloomConfig()
			else:
						SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					BloomConfig.from_json_file(A__				)

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

						SCREAMING_SNAKE_CASE_		:     Tuple					=					{'weight_map': {}, 'metadata': {}}
						SCREAMING_SNAKE_CASE_		:     int					=					0

						SCREAMING_SNAKE_CASE_		:     str					=					None

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

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

									for i in range(A__				):
												# load all TP files
												SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					file.replace('model_00',      F'''model_0{i}'''				)
												SCREAMING_SNAKE_CASE_		:     List[Any]					=					torch.load(os.path.join(A__,      A__				),      map_location='cpu'				)

												# Rename keys in the transformers names
												SCREAMING_SNAKE_CASE_		:     Dict					=					list(temp.keys()				)
												for key in keys:
															SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					temp.pop(A__				)

												if tensors is None:
															SCREAMING_SNAKE_CASE_		:     Tuple					=					temp
												else:
															for key in tensors.keys():
																		if any(key.endswith(A__				) 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_		:     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[int]					=					torch.cat([tensors[key], temp[key]],      dim=A__				)

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

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

						SCREAMING_SNAKE_CASE_		:     Any					=					BloomConfig()
						SCREAMING_SNAKE_CASE_		:     str					=					pytorch_dump_folder_path + '/' + CONFIG_NAME
						SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					total_size
						with open(A__,      'w',      encoding='utf-8'				) as f:
									f.write(config.to_json_string()				)
						with open(os.path.join(A__,      WEIGHTS_NAME + '.index.json'				),      'w',      encoding='utf-8'				) as f:
									SCREAMING_SNAKE_CASE_		:     Dict					=					json.dumps(A__,      indent=2,      sort_keys=A__				) + '\n'
									f.write(A__				)
			else:
						SCREAMING_SNAKE_CASE_		:     Tuple					=					BloomModel(A__				)

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

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

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

												if tensors is None:
															SCREAMING_SNAKE_CASE_		:     List[Any]					=					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(A__				) 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_		:     Dict					=					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_		:     Union[str, Any]					=					torch.cat([tensors[key], temp[key]],      dim=A__				)

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

									SCREAMING_SNAKE_CASE_		:     Any					=					model.load_state_dict(A__,      strict=A__				)
									assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
									if missing_keys is None:
												SCREAMING_SNAKE_CASE_		:     Tuple					=					set(other_keys.missing_keys				)
									else:
												SCREAMING_SNAKE_CASE_		:     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(A__,      exist_ok=A__				)
						SCREAMING_SNAKE_CASE_		:     Any					=					pytorch_dump_folder_path + '/' + WEIGHTS_NAME
						SCREAMING_SNAKE_CASE_		:     List[str]					=					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_		:     Union[str, Any]					=					model.to(config.torch_dtype				)
						torch.save(model.state_dict(),      A__				)
						print(F'''Save configuration file to {pytorch_config_dump_path}'''				)
						with open(A__,      'w',      encoding='utf-8'				) as f:
									f.write(config.to_json_string()				)


if __name__ == "__main__":
				lowerCAmelCase__			:      Optional[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,
				)













'''simple docstring'''







# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from packaging import version

from .. import __version__
from .constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD
from .doc import (
    add_code_sample_docstrings,
    add_end_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    copy_func,
    replace_return_docstrings,
)
from .generic import (
    ContextManagers,
    ExplicitEnum,
    ModelOutput,
    PaddingStrategy,
    TensorType,
    add_model_info_to_auto_map,
    cached_property,
    can_return_loss,
    expand_dims,
    find_labels,
    flatten_dict,
    infer_framework,
    is_jax_tensor,
    is_numpy_array,
    is_tensor,
    is_tf_symbolic_tensor,
    is_tf_tensor,
    is_torch_device,
    is_torch_dtype,
    is_torch_tensor,
    reshape,
    squeeze,
    strtobool,
    tensor_size,
    to_numpy,
    to_py_obj,
    transpose,
    working_or_temp_dir,
)
from .hub import (
    CLOUDFRONT_DISTRIB_PREFIX,
    DISABLE_TELEMETRY,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_PREFIX,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
    S3_BUCKET_PREFIX,
    TRANSFORMERS_CACHE,
    TRANSFORMERS_DYNAMIC_MODULE_NAME,
    EntryNotFoundError,
    PushToHubMixin,
    RepositoryNotFoundError,
    RevisionNotFoundError,
    cached_file,
    default_cache_path,
    define_sagemaker_information,
    download_url,
    extract_commit_hash,
    get_cached_models,
    get_file_from_repo,
    get_full_repo_name,
    has_file,
    http_user_agent,
    is_offline_mode,
    is_remote_url,
    move_cache,
    send_example_telemetry,
    try_to_load_from_cache,
)
from .import_utils import (
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    TORCH_FX_REQUIRED_VERSION,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    _LazyModule,
    ccl_version,
    direct_transformers_import,
    get_torch_version,
    is_accelerate_available,
    is_apex_available,
    is_bitsandbytes_available,
    is_bsa_available,
    is_coloredlogs_available,
    is_cython_available,
    is_datasets_available,
    is_decord_available,
    is_detectrona_available,
    is_faiss_available,
    is_flax_available,
    is_ftfy_available,
    is_in_notebook,
    is_ipex_available,
    is_jieba_available,
    is_jumanpp_available,
    is_kenlm_available,
    is_keras_nlp_available,
    is_librosa_available,
    is_natten_available,
    is_ninja_available,
    is_onnx_available,
    is_openai_available,
    is_optimum_available,
    is_pandas_available,
    is_peft_available,
    is_phonemizer_available,
    is_protobuf_available,
    is_psutil_available,
    is_pyanvml_available,
    is_pyctcdecode_available,
    is_pytesseract_available,
    is_pytest_available,
    is_pytorch_quantization_available,
    is_rjieba_available,
    is_sacremoses_available,
    is_safetensors_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_scipy_available,
    is_sentencepiece_available,
    is_seqio_available,
    is_sklearn_available,
    is_soundfile_availble,
    is_spacy_available,
    is_speech_available,
    is_sudachi_available,
    is_tensorflow_probability_available,
    is_tensorflow_text_available,
    is_tfaonnx_available,
    is_tf_available,
    is_timm_available,
    is_tokenizers_available,
    is_torch_available,
    is_torch_bfaa_available,
    is_torch_bfaa_cpu_available,
    is_torch_bfaa_gpu_available,
    is_torch_compile_available,
    is_torch_cuda_available,
    is_torch_fx_available,
    is_torch_fx_proxy,
    is_torch_mps_available,
    is_torch_neuroncore_available,
    is_torch_tensorrt_fx_available,
    is_torch_tfaa_available,
    is_torch_tpu_available,
    is_torchaudio_available,
    is_torchdistx_available,
    is_torchdynamo_available,
    is_torchvision_available,
    is_training_run_on_sagemaker,
    is_vision_available,
    requires_backends,
    torch_only_method,
)


__lowerCAmelCase        =							"pytorch_model.bin"
__lowerCAmelCase        =							"pytorch_model.bin.index.json"
__lowerCAmelCase        =							"adapter_config.json"
__lowerCAmelCase        =							"adapter_model.bin"
__lowerCAmelCase        =							"adapter_model.safetensors"
__lowerCAmelCase        =							"tf_model.h5"
__lowerCAmelCase        =							"tf_model.h5.index.json"
__lowerCAmelCase        =							"model.ckpt"
__lowerCAmelCase        =							"flax_model.msgpack"
__lowerCAmelCase        =							"flax_model.msgpack.index.json"
__lowerCAmelCase        =							"model.safetensors"
__lowerCAmelCase        =							"model.safetensors.index.json"
__lowerCAmelCase        =							"config.json"
__lowerCAmelCase        =							"preprocessor_config.json"
__lowerCAmelCase        =							FEATURE_EXTRACTOR_NAME
__lowerCAmelCase        =							"generation_config.json"
__lowerCAmelCase        =							"modelcard.json"

__lowerCAmelCase        =							"▁"
__lowerCAmelCase        =							SENTENCEPIECE_UNDERLINE  # Kept for backward compatibility

__lowerCAmelCase        =							[
    [[0, 1, 0, 1], [1, 0, 0, 1]]
] * 2  # Needs to have 0s and 1s only since XLM uses it for langs too.
__lowerCAmelCase        =							[[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]]
__lowerCAmelCase        =							[[1, 1, 1, 1, 1], [1, 1, 1, 0, 0], [0, 0, 0, 1, 1]]


def       __UpperCamelCase					(      lowercase_       :  Optional[int]   ):
       """simple docstring"""




       if version.parse(lowercase_   ) < version.parse(lowercase_   ):
              if "dev" in min_version:
                     a_						     =			(
                         'This example requires a source install from HuggingFace Transformers (see '
                         '`https://huggingface.co/docs/transformers/installation#install-from-source`),'
                     )
              else:
                     a_						     =			F'This example requires a minimum version of {min_version},'
              error_message += F' but the version found is {__version__}.\n'
              raise ImportError(
                  error_message
                  + 'Check out https://github.com/huggingface/transformers/tree/main/examples#important-note for the examples corresponding to other '
                  'versions of HuggingFace Transformers.'   )




