

"""simple docstring"""





import hashlib
import unittest

from transformers import MODEL_FOR_DEPTH_ESTIMATION_MAPPING, is_torch_available, is_vision_available
from transformers.pipelines import DepthEstimationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_timm,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_torch_available():
							import torch

if is_vision_available():
							from PIL import Image
else:







							class 							SCREAMING_SNAKE_CASE   :






												"""simple docstring"""



												@staticmethod
												def 			__lowerCAmelCase					(			*lowercase_     :							Tuple  ,**lowercase_     :							List[Any]						):
													pass

def     __SCREAMING_SNAKE_CASE    ( A_					):
	lowerCAmelCase__      :		Tuple    							=      hashlib.mda(image.tobytes()					)
	return m.hexdigest()







@is_pipeline_test
@require_vision
@require_timm
@require_torch
class 							SCREAMING_SNAKE_CASE   (						unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			MODEL_FOR_DEPTH_ESTIMATION_MAPPING







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]  ,lowercase_     :							Optional[int]  ,lowercase_     :							List[str]  ,lowercase_     :							Dict						):
						lowerCAmelCase__      :		Optional[int]    							=      DepthEstimationPipeline(model=__snake_case  ,image_processor=__snake_case						)
						return depth_estimator, [
						    "./tests/fixtures/tests_samples/COCO/000000039769.png",
						    "./tests/fixtures/tests_samples/COCO/000000039769.png",
						]







					def 			__lowerCAmelCase					(			self     :							List[str]  ,lowercase_     :							Any  ,lowercase_     :							List[str]						):
						lowerCAmelCase__      :		Optional[int]    							=      depth_estimator('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
						self.assertEqual({'''predicted_depth''': ANY(torch.Tensor						), '''depth''': ANY(Image.Image						)}  ,__snake_case						)
						import datasets

						lowerCAmelCase__      :		Optional[Any]    							=      datasets.load_dataset('''hf-internal-testing/fixtures_image_utils'''  ,'''image'''  ,split='''test'''						)
						lowerCAmelCase__      :		Union[str, Any]    							=      depth_estimator(
						    [
						        Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						),
						        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
						        # RGBA
						        dataset[0]['''file'''],
						        # LA
						        dataset[1]['''file'''],
						        # L
						        dataset[2]['''file'''],
						    ]						)
						self.assertEqual(
						    [
						        {'''predicted_depth''': ANY(torch.Tensor						), '''depth''': ANY(Image.Image						)},
						        {'''predicted_depth''': ANY(torch.Tensor						), '''depth''': ANY(Image.Image						)},
						        {'''predicted_depth''': ANY(torch.Tensor						), '''depth''': ANY(Image.Image						)},
						        {'''predicted_depth''': ANY(torch.Tensor						), '''depth''': ANY(Image.Image						)},
						        {'''predicted_depth''': ANY(torch.Tensor						), '''depth''': ANY(Image.Image						)},
						    ]  ,__snake_case  ,)







					@require_tf
					@unittest.skip('''Depth estimation is not implemented in TF'''						)
					def 			__lowerCAmelCase					(			self     :							List[str]						):
						pass







					@slow
					@require_torch
					def 			__lowerCAmelCase					(			self     :							Dict						):
						lowerCAmelCase__      :		Optional[Any]    							=      """Intel/dpt-large"""
						lowerCAmelCase__      :		Any    							=      pipeline('''depth-estimation'''  ,model=__snake_case						)
						lowerCAmelCase__      :		Optional[int]    							=      depth_estimator('''http://images.cocodataset.org/val2017/000000039769.jpg'''						)
						lowerCAmelCase__      :		Optional[int]    							=      hashimage(outputs['''depth''']						)

						# This seems flaky.
						# self.assertEqual(outputs["depth"], "1a39394e282e9f3b0741a90b9f108977")
						self.assertEqual(nested_simplify(outputs['''predicted_depth'''].max().item()						)  ,29.304						)
						self.assertEqual(nested_simplify(outputs['''predicted_depth'''].min().item()						)  ,2.662						)







					@require_torch
					def 			__lowerCAmelCase					(			self     :							Any						):
						# This is highly irregular to have no small tests.
						self.skipTest('''There is not hf-internal-testing tiny model for either GLPN nor DPT'''						)












"""simple docstring"""





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



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



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



def   a    (     __UpperCAmelCase      :    Union[str, Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    List[Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    Tuple  )    ->							List[str]:
						# Construct model
						if bloom_config_file == "":
												__magic_name__:	Optional[Any]							=       BloomConfig()
						else:
												__magic_name__:	Tuple							=       BloomConfig.from_json_file(__UpperCAmelCase  )

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

												__magic_name__:	str							=       {"""weight_map""": {}, """metadata""": {}}
												__magic_name__:	Optional[Any]							=       0

												__magic_name__:	Dict							=       None

												__magic_name__:	Optional[int]							=       BloomConfig()

												for j, file in enumerate(__UpperCAmelCase  ):
																		print("""Processing file: {}""".format(__UpperCAmelCase  )  )
																		__magic_name__:	Tuple							=       None

																		for i in range(__UpperCAmelCase  ):
																								# load all TP files
																								__magic_name__:	int							=       file.replace("""model_00"""      ,   f'model_0{i}'  )
																								__magic_name__:	Optional[Any]							=       torch.load(os.path.join(__UpperCAmelCase      ,   __UpperCAmelCase  )      ,   map_location="""cpu"""  )

																								# Rename keys in the transformers names
																								__magic_name__:	str							=       list(temp.keys()  )
																								for key in keys:
																														__magic_name__:	Any							=       temp.pop(__UpperCAmelCase  )

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

												__magic_name__:	Tuple							=       BloomConfig()
												__magic_name__:	Optional[int]							=       pytorch_dump_folder_path + """/""" + CONFIG_NAME
												__magic_name__:	Optional[int]							=       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:
																		__magic_name__:	str							=       json.dumps(__UpperCAmelCase      ,   indent=2      ,   sort_keys=__UpperCAmelCase  ) + """\n"""
																		f.write(__UpperCAmelCase  )
						else:
												__magic_name__:	Optional[int]							=       BloomModel(__UpperCAmelCase  )

												__magic_name__:	List[str]							=       os.listdir(__UpperCAmelCase  )
												__magic_name__:	Optional[Any]							=       sorted(filter(lambda __UpperCAmelCase  : s.startswith("""layer"""  ) and "model_00" in s      ,   __UpperCAmelCase  )  )

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

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

																								if tensors is None:
																														__magic_name__:	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(__UpperCAmelCase  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
																																										tensors[key] += temp[key]
																																				else:
																																										# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																																										__magic_name__:	Tuple							=       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
																																										# We concatenate these weights accross TP ranks
																																										__magic_name__:	Dict							=       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  ):
																														__magic_name__:	Optional[int]							=       tensors[key] / pretraining_tp

																		__magic_name__:	Union[str, Any]							=       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:
																								__magic_name__:	Optional[Any]							=       set(other_keys.missing_keys  )
																		else:
																								__magic_name__:	Tuple							=       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  )
												__magic_name__:	int							=       pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
												__magic_name__:	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:
																		__magic_name__:	Optional[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							=  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							=  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,
							)



