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

_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       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		Optional[Any]:


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



def       lowerCamelCase__		(			__snake_case	)							->		List[Any]:


					"""simple docstring"""





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


					"""simple docstring"""





					if bloom_config_file == "":
										_UpperCamelCase							 =  BloomConfig()
					else:
										_UpperCamelCase							 =  BloomConfig.from_json_file(__snake_case	)

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

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

										_UpperCamelCase							 =  None

										_UpperCamelCase							 =  BloomConfig()

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

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

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

																				if tensors is None:
																									_UpperCamelCase							 =  temp
																				else:
																									for key in tensors.keys():
																														if any(key.endswith(__snake_case	) 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=__snake_case	)

            # Divide by the number of TP the weights we want to average
															for key in tensors.keys():
																				if any(key.endswith(__snake_case	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
																									_UpperCamelCase							 =  tensors[key] / pretraining_tp
															torch.save(
															    __snake_case,     os.path.join(
															        __snake_case,     '''pytorch_model_{}-of-{}.bin'''.format(str(j + 1	).zfill(5	),     str(len(__snake_case	)	).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(__snake_case	)	).zfill(5	)	)

										_UpperCamelCase							 =  BloomConfig()
										_UpperCamelCase							 =  pytorch_dump_folder_path + '''/''' + CONFIG_NAME
										_UpperCamelCase							 =  total_size
										with open(__snake_case,     '''w''',     encoding='''utf-8'''	) as f:
															f.write(config.to_json_string()	)
										with open(os.path.join(__snake_case,     WEIGHTS_NAME + '''.index.json'''	),     '''w''',     encoding='''utf-8'''	) as f:
															_UpperCamelCase							 =  json.dumps(__snake_case,     indent=2,     sort_keys=__snake_case	) + '''\n'''
															f.write(__snake_case	)
					else:
										_UpperCamelCase							 =  BloomModel(__snake_case	)

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

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

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

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

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

															_UpperCamelCase							 =  model.load_state_dict(__snake_case,     strict=__snake_case	)
															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(__snake_case,     exist_ok=__snake_case	)
										_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(),     __snake_case	)
										print(F'''Save configuration file to {pytorch_config_dump_path}'''	)
										with open(__snake_case,     '''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,
				)








"""simple docstring"""





from typing import List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


_a														=		logging.get_logger(__name__)






class     _UpperCAmelCase(    lowerCamelCase	):



		def 						UpperCAmelCase							(					self				,		__a)       ->       Tuple:






							'''simple docstring'''
							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [label.strip() for label in labels.split(''',''') if label.strip()]
							return labels




		def __call__(					self				,		__a				,		__a				,		__a)       ->       Optional[int]:






							'''simple docstring'''
							if len(__a) == 0 or len(__a) == 0:
												raise ValueError('''You must include at least one label and at least one sequence.''')
							if hypothesis_template.format(labels[0]) == hypothesis_template:
												raise ValueError(
												    (
												        '''The provided hypothesis_template "{}" was not able to be formatted with the target labels. '''
												        '''Make sure the passed template includes formatting syntax such as {{}} where the label should go.'''
												    ).format(__a))

							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [sequences]

							_UpperCamelCase							 =  []
							for sequence in sequences:
												sequence_pairs.extend([[sequence, hypothesis_template.format(__a)] for label in labels])

							return sequence_pairs, sequences






@add_end_docstrings(lowerCamelCase	)
class     _UpperCAmelCase(    lowerCamelCase	):



		def __init__(					self				,		__a=ZeroShotClassificationArgumentHandler()				,		*__a				,		**__a)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  args_parser
							super().__init__(*__a				,		**__a)
							if self.entailment_id == -1:
												logger.warning(
												    '''Failed to determine \'entailment\' label id from the label2id mapping in the model config. Setting to '''
												    '''-1. Define a descriptive label2id mapping in the model config to ensure correct outputs.''')



		@property
		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							for label, ind in self.model.config.labelaid.items():
												if label.lower().startswith('''entail'''):
																	return ind
							return -1



		def 						UpperCAmelCase							(					self				,		__a				,		__a=True				,		__a=True				,		__a=TruncationStrategy.ONLY_FIRST				,		**__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  self.framework
							if self.tokenizer.pad_token is None:
												# Override for tokenizers not supporting padding
												logger.error(
												    '''Tokenizer was not supporting padding necessary for zero-shot, attempting to use '''
												    ''' `pad_token=eos_token`''')
												_UpperCamelCase							 =  self.tokenizer.eos_token
							try:
												_UpperCamelCase							 =  self.tokenizer(
												    __a				,		add_special_tokens=__a				,		return_tensors=__a				,		padding=__a				,		truncation=__a				,		)
							except Exception as e:
												if "too short" in str(__a):
																	# tokenizers might yell that we want to truncate
																	# to a value that is not even reached by the input.
																	# In that case we don't want to truncate.
																	# It seems there's not a really better way to catch that
																	# exception.

																	_UpperCamelCase							 =  self.tokenizer(
																	    __a				,		add_special_tokens=__a				,		return_tensors=__a				,		padding=__a				,		truncation=TruncationStrategy.DO_NOT_TRUNCATE				,		)
												else:
																	raise e

							return inputs



		def 						UpperCAmelCase							(					self				,		**__a)       ->       Any:






							'''simple docstring'''
							if kwargs.get('''multi_class'''				,		__a) is not None:
												_UpperCamelCase							 =  kwargs['''multi_class''']
												logger.warning(
												    '''The `multi_class` argument has been deprecated and renamed to `multi_label`. '''
												    '''`multi_class` will be removed in a future version of Transformers.''')
							_UpperCamelCase							 =  {}
							if "candidate_labels" in kwargs:
												_UpperCamelCase							 =  self._args_parser._parse_labels(kwargs['''candidate_labels'''])
							if "hypothesis_template" in kwargs:
												_UpperCamelCase							 =  kwargs['''hypothesis_template''']

							_UpperCamelCase							 =  {}
							if "multi_label" in kwargs:
												_UpperCamelCase							 =  kwargs['''multi_label''']
							return preprocess_params, {}, postprocess_params



		def __call__(					self				,		__a				,		*__a				,		**__a				,		)       ->       int:






							'''simple docstring'''
							if len(__a) == 0:
												pass
							elif len(__a) == 1 and "candidate_labels" not in kwargs:
												_UpperCamelCase							 =  args[0]
							else:
												raise ValueError(F'''Unable to understand extra arguments {args}''')

							return super().__call__(__a				,		**__a)



		def 						UpperCAmelCase							(					self				,		__a				,		__a=None				,		__a="This example is {}.")       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self._args_parser(__a				,		__a				,		__a)

							for i, (candidate_label, sequence_pair) in enumerate(zip(__a				,		__a)):
												_UpperCamelCase							 =  self._parse_and_tokenize([sequence_pair])

												yield {
												    "candidate_label": candidate_label,
												    "sequence": sequences[0],
												    "is_last": i == len(__a) - 1,
												    **model_input,
												}



		def 						UpperCAmelCase							(					self				,		__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  inputs['''candidate_label''']
							_UpperCamelCase							 =  inputs['''sequence''']
							_UpperCamelCase							 =  {k: inputs[k] for k in self.tokenizer.model_input_names}
							_UpperCamelCase							 =  self.model(**__a)

							_UpperCamelCase							 =  {
							    '''candidate_label''': candidate_label,
							    '''sequence''': sequence,
							    '''is_last''': inputs['''is_last'''],
							    **outputs,
							}
							return model_outputs




		def 						UpperCAmelCase							(					self				,		__a				,		__a=False)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  [outputs['''candidate_label'''] for outputs in model_outputs]
							_UpperCamelCase							 =  [outputs['''sequence'''] for outputs in model_outputs]
							_UpperCamelCase							 =  np.concatenate([output['''logits'''].numpy() for output in model_outputs])
							_UpperCamelCase							 =  logits.shape[0]
							_UpperCamelCase							 =  len(__a)
							_UpperCamelCase							 =  N // n
							_UpperCamelCase							 =  logits.reshape((num_sequences, n, -1))

							if multi_label or len(__a) == 1:
												# softmax over the entailment vs. contradiction dim for each label independently
												_UpperCamelCase							 =  self.entailment_id
												_UpperCamelCase							 =  -1 if entailment_id == 0 else 0
												_UpperCamelCase							 =  reshaped_outputs[..., [contradiction_id, entailment_id]]
												_UpperCamelCase							 =  np.exp(__a) / np.exp(__a).sum(-1				,		keepdims=__a)
												_UpperCamelCase							 =  scores[..., 1]
							else:
												# softmax the "entailment" logits over all candidate labels
												_UpperCamelCase							 =  reshaped_outputs[..., self.entailment_id]
												_UpperCamelCase							 =  np.exp(__a) / np.exp(__a).sum(-1				,		keepdims=__a)

							_UpperCamelCase							 =  list(reversed(scores[0].argsort()))
							return {
							    "sequence": sequences[0],
							    "labels": [candidate_labels[i] for i in top_inds],
							    "scores": scores[0, top_inds].tolist(),
							}






