






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

__snake_case      :      Dict             =   [
    """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""",
]

__snake_case      :      str             =   [
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]







def 		_UpperCamelCase				(      UpperCamelCase_    :   str      ,    UpperCamelCase_    :   Dict						)					->				Any:

						"""simple docstring"""



						lowerCAmelCase__        =  {
						    '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
						lowerCAmelCase__        =  int(re.match(r'.*layer_(\d*).*'      ,    UpperCamelCase_						)[1]						)
						layer_number -= 3
						return F"h.{layer_number}." + key







def 		_UpperCamelCase				(      UpperCamelCase_    :   Union[str, Any]						)					->				List[str]:

						"""simple docstring"""



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







def 		_UpperCamelCase				(      UpperCamelCase_    :   Union[str, Any]      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   Tuple      ,    UpperCamelCase_    :   List[Any]						)					->				List[Any]:

						"""simple docstring"""



						if bloom_config_file == "":
												lowerCAmelCase__        =  BloomConfig()
						else:
												lowerCAmelCase__        =  BloomConfig.from_json_file(UpperCamelCase_						)

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

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

												lowerCAmelCase__        =  None

												lowerCAmelCase__        =  BloomConfig()

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

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

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

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

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

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

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

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

																								if tensors is None:
																														lowerCAmelCase__        =  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
																																										lowerCAmelCase__        =  1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
																																										# We concatenate these weights accross TP ranks
																																										lowerCAmelCase__        =  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						):
																														lowerCAmelCase__        =  tensors[key] / pretraining_tp

																		lowerCAmelCase__        =  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:
																								lowerCAmelCase__        =  set(other_keys.missing_keys						)
																		else:
																								lowerCAmelCase__        =  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_						)
												lowerCAmelCase__        =  pytorch_dump_folder_path + '/' + WEIGHTS_NAME
												lowerCAmelCase__        =  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:
																		lowerCAmelCase__        =  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__":
					__snake_case      :      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""",
					)
					__snake_case      :      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 __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import TFResNetForImageClassification, TFResNetModel
					from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor


class       __SCREAMING_SNAKE_CASE   :






						def __init__(		self       ,  _UpperCamelCase       ,  _UpperCamelCase=3       ,  _UpperCamelCase=32       ,  _UpperCamelCase=3       ,  _UpperCamelCase=10       ,  _UpperCamelCase=[10, 20, 30, 40]       ,  _UpperCamelCase=[1, 1, 2, 1]       ,  _UpperCamelCase=True       ,  _UpperCamelCase=True       ,  _UpperCamelCase="relu"       ,  _UpperCamelCase=3       ,  _UpperCamelCase=None       ,  ):






												"""simple docstring"""
												lowerCAmelCase__        =  parent
												lowerCAmelCase__        =  batch_size
												lowerCAmelCase__        =  image_size
												lowerCAmelCase__        =  num_channels
												lowerCAmelCase__        =  embeddings_size
												lowerCAmelCase__        =  hidden_sizes
												lowerCAmelCase__        =  depths
												lowerCAmelCase__        =  is_training
												lowerCAmelCase__        =  use_labels
												lowerCAmelCase__        =  hidden_act
												lowerCAmelCase__        =  num_labels
												lowerCAmelCase__        =  scope
												lowerCAmelCase__        =  len(_UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

												lowerCAmelCase__        =  None
												if self.use_labels:
																		lowerCAmelCase__        =  ids_tensor([self.batch_size]       ,  self.num_labels				)

												lowerCAmelCase__        =  self.get_config()

												return config, pixel_values, labels






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												return ResNetConfig(
												    num_channels=self.num_channels       ,  embeddings_size=self.embeddings_size       ,  hidden_sizes=self.hidden_sizes       ,  depths=self.depths       ,  hidden_act=self.hidden_act       ,  num_labels=self.num_labels       ,  image_size=self.image_size       ,  )






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												lowerCAmelCase__        =  TFResNetModel(config=_UpperCamelCase				)
												lowerCAmelCase__        =  model(_UpperCamelCase				)
												# expected last hidden states: B, C, H // 32, W // 32
												self.parent.assertEqual(
												    result.last_hidden_state.shape       ,  (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,  )






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												lowerCAmelCase__        =  self.num_labels
												lowerCAmelCase__        =  TFResNetForImageClassification(_UpperCamelCase				)
												lowerCAmelCase__        =  model(_UpperCamelCase       ,  labels=_UpperCamelCase				)
												self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_labels)				)







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  self.prepare_config_and_inputs()
												lowerCAmelCase__  ,    lowerCAmelCase__  ,    lowerCAmelCase__        =  config_and_inputs
												lowerCAmelCase__        =  {'pixel_values': pixel_values}
												return config, inputs_dict




@require_tf
class       __SCREAMING_SNAKE_CASE   (				__lowercase						,   __lowercase						,   unittest.TestCase):
						_SCREAMING_SNAKE_CASE		:    Any						=   (TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
						_SCREAMING_SNAKE_CASE		:    List[str]						=   (
						    {'''feature-extraction''': TFResNetModel, '''image-classification''': TFResNetForImageClassification}
						    if is_tf_available()
						    else {}
						)

						_SCREAMING_SNAKE_CASE		:    Union[str, Any]						=   False
						_SCREAMING_SNAKE_CASE		:    int						=   False
						_SCREAMING_SNAKE_CASE		:    Optional[int]						=   False
						_SCREAMING_SNAKE_CASE		:    List[str]						=   False
						_SCREAMING_SNAKE_CASE		:    Optional[Any]						=   False






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  TFResNetModelTester(self				)
												lowerCAmelCase__        =  ConfigTester(self       ,  config_class=_UpperCamelCase       ,  has_text_modality=_UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												self.create_and_test_config_common_properties()
												self.config_tester.create_and_test_config_to_json_string()
												self.config_tester.create_and_test_config_to_json_file()
												self.config_tester.create_and_test_config_from_and_save_pretrained()
												self.config_tester.create_and_test_config_with_num_labels()
												self.config_tester.check_config_can_be_init_without_params()
												self.config_tester.check_config_arguments_init()






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												return






						@unittest.skip(reason='ResNet does not use inputs_embeds'				)
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												pass






						@unittest.skip(reason='ResNet does not support input and output embeddings'				)
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												pass






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__  ,    lowerCAmelCase__        =  self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		lowerCAmelCase__        =  model_class(_UpperCamelCase				)
																		lowerCAmelCase__        =  inspect.signature(model.call				)
																		# signature.parameters is an OrderedDict => so arg_names order is deterministic
																		lowerCAmelCase__        =  [*signature.parameters.keys()]

																		lowerCAmelCase__        =  ['pixel_values']
																		self.assertListEqual(arg_names[:1]       ,  _UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												def check_hidden_states_output(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):
																		lowerCAmelCase__        =  model_class(_UpperCamelCase				)
																		lowerCAmelCase__        =  model(**self._prepare_for_class(_UpperCamelCase       ,  _UpperCamelCase				)				)

																		lowerCAmelCase__        =  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																		lowerCAmelCase__        =  self.model_tester.num_stages
																		self.assertEqual(len(_UpperCamelCase				)       ,  expected_num_stages + 1				)

																		# ResNet's feature maps are of shape (batch_size, num_channels, height, width)
																		self.assertListEqual(
																		    list(hidden_states[0].shape[-2:]				)       ,  [self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,  )

												lowerCAmelCase__  ,    lowerCAmelCase__        =  self.model_tester.prepare_config_and_inputs_for_common()
												lowerCAmelCase__        =  ['basic', 'bottleneck']
												for model_class in self.all_model_classes:
																		for layer_type in layers_type:
																								lowerCAmelCase__        =  layer_type
																								lowerCAmelCase__        =  True
																								check_hidden_states_output(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				)

																								# check that output_hidden_states also work using config
																								del inputs_dict["output_hidden_states"]
																								lowerCAmelCase__        =  True

																								check_hidden_states_output(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase				)







						@slow
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase__        =  TFResNetModel.from_pretrained(_UpperCamelCase				)
																		self.assertIsNotNone(_UpperCamelCase				)







def 		_UpperCamelCase				(      )					->				List[str]:

						"""simple docstring"""



						lowerCAmelCase__        =  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
						return image


@require_tf
@require_vision
class       __SCREAMING_SNAKE_CASE   (				unittest.TestCase):






						@cached_property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												return (
												    AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]				)
												    if is_vision_available()
												    else None
												)







						@slow
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]				)

												lowerCAmelCase__        =  self.default_image_processor
												lowerCAmelCase__        =  prepare_img()
												lowerCAmelCase__        =  image_processor(images=_UpperCamelCase       ,  return_tensors='tf'				)

												# forward pass
												lowerCAmelCase__        =  model(**_UpperCamelCase				)

												# verify the logits
												lowerCAmelCase__        =  tf.TensorShape((1, 10_00)				)
												self.assertEqual(outputs.logits.shape       ,  _UpperCamelCase				)

												lowerCAmelCase__        =  tf.constant([-11.10_69, -9.78_77, -8.37_77]				)

												self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()       ,  _UpperCamelCase       ,  atol=1E-4				)				)





