


'''simple docstring'''






from sklearn.metrics import recall_score

import datasets


__UpperCAmelCase       =							'''
Recall is the fraction of the positive examples that were correctly labeled by the model as positive. It can be computed with the equation:
Recall = TP / (TP + FN)
Where TP is the true positives and FN is the false negatives.
'''


__UpperCAmelCase       =							'''
Args:
- **predictions** (`list` of `int`): The predicted labels.
- **references** (`list` of `int`): The ground truth labels.
- **labels** (`list` of `int`): The set of labels to include when `average` is not set to `binary`, and their order when average is `None`. Labels present in the data can be excluded in this input, for example to calculate a multiclass average ignoring a majority negative class, while labels not present in the data will result in 0 components in a macro average. For multilabel targets, labels are column indices. By default, all labels in y_true and y_pred are used in sorted order. Defaults to None.
- **pos_label** (`int`): The class label to use as the \'positive class\' when calculating the recall. Defaults to `1`.
- **average** (`string`): This parameter is required for multiclass/multilabel targets. If None, the scores for each class are returned. Otherwise, this determines the type of averaging performed on the data. Defaults to `\'binary\'`.
    - `\'binary\'`: Only report results for the class specified by `pos_label`. This is applicable only if the target labels and predictions are binary.
    - `\'micro\'`: Calculate metrics globally by counting the total true positives, false negatives, and false positives.
    - `\'macro\'`: Calculate metrics for each label, and find their unweighted mean. This does not take label imbalance into account.
    - `\'weighted\'`: Calculate metrics for each label, and find their average weighted by support (the number of true instances for each label). This alters `\'macro\'` to account for label imbalance. Note that it can result in an F-score that is not between precision and recall.
    - `\'samples\'`: Calculate metrics for each instance, and find their average (only meaningful for multilabel classification).
- **sample_weight** (`list` of `float`): Sample weights Defaults to `None`.
- **zero_division** (): Sets the value to return when there is a zero division. Defaults to .
    - `\'warn\'`: If there is a zero division, the return value is `0`, but warnings are also raised.
    - `0`: If there is a zero division, the return value is `0`.
    - `1`: If there is a zero division, the return value is `1`.

Returns:
- **recall** (`float`, or `array` of `float`): Either the general recall score, or the recall scores for individual classes, depending on the values input to `labels` and `average`. Minimum possible value is 0. Maximum possible value is 1. A higher recall means that more of the positive examples have been labeled correctly. Therefore, a higher recall is generally considered better.

Examples:

    Example 1-A simple example with some errors
        >>> recall_metric = datasets.load_metric(\'recall\')
        >>> results = recall_metric.compute(references=[0, 0, 1, 1, 1], predictions=[0, 1, 0, 1, 1])
        >>> print(results)
        {\'recall\': 0.6666666666666666}

    Example 2-The same example as Example 1, but with `pos_label=0` instead of the default `pos_label=1`.
        >>> recall_metric = datasets.load_metric(\'recall\')
        >>> results = recall_metric.compute(references=[0, 0, 1, 1, 1], predictions=[0, 1, 0, 1, 1], pos_label=0)
        >>> print(results)
        {\'recall\': 0.5}

    Example 3-The same example as Example 1, but with `sample_weight` included.
        >>> recall_metric = datasets.load_metric(\'recall\')
        >>> sample_weight = [0.9, 0.2, 0.9, 0.3, 0.8]
        >>> results = recall_metric.compute(references=[0, 0, 1, 1, 1], predictions=[0, 1, 0, 1, 1], sample_weight=sample_weight)
        >>> print(results)
        {\'recall\': 0.55}

    Example 4-A multiclass example, using different averages.
        >>> recall_metric = datasets.load_metric(\'recall\')
        >>> predictions = [0, 2, 1, 0, 0, 1]
        >>> references = [0, 1, 2, 0, 1, 2]
        >>> results = recall_metric.compute(predictions=predictions, references=references, average=\'macro\')
        >>> print(results)
        {\'recall\': 0.3333333333333333}
        >>> results = recall_metric.compute(predictions=predictions, references=references, average=\'micro\')
        >>> print(results)
        {\'recall\': 0.3333333333333333}
        >>> results = recall_metric.compute(predictions=predictions, references=references, average=\'weighted\')
        >>> print(results)
        {\'recall\': 0.3333333333333333}
        >>> results = recall_metric.compute(predictions=predictions, references=references, average=None)
        >>> print(results)
        {\'recall\': array([1., 0., 0.])}
'''


__UpperCAmelCase       =							'''
@article{scikit-learn, title={Scikit-learn: Machine Learning in {P}ython}, author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V. and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P. and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.}, journal={Journal of Machine Learning Research}, volume={12}, pages={2825--2830}, year={2011}
'''




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION    ,    _KWARGS_DESCRIPTION  )
class 		a__    (							datasets.Metric  ):





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Dict:
														return datasets.MetricInfo(
														    description=_DESCRIPTION ,   citation=_CITATION ,   inputs_description=_KWARGS_DESCRIPTION ,   features=datasets.Features(
														        {
														            '''predictions''': datasets.Sequence(datasets.Value('''int32'''       )       ),
														            '''references''': datasets.Sequence(datasets.Value('''int32'''       )       ),
														        }
														        if self.config_name == '''multilabel'''
														        else {
														            '''predictions''': datasets.Value('''int32'''       ),
														            '''references''': datasets.Value('''int32'''       ),
														        }       ) ,   reference_urls=['''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.recall_score.html'''] ,   )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_=None ,   lowerCamelCase_=1 ,   lowerCamelCase_="binary" ,   lowerCamelCase_=None ,   lowerCamelCase_="warn" ,   )	->							Any:
														lowerCAmelCase__	       =     recall_score(
														    lowerCamelCase_ ,   lowerCamelCase_ ,   labels=lowerCamelCase_ ,   pos_label=lowerCamelCase_ ,   average=lowerCamelCase_ ,   sample_weight=lowerCamelCase_ ,   zero_division=lowerCamelCase_ ,   )
														return {"recall": float(lowerCamelCase_       ) if score.size == 1 else score}



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

__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  _snake_case       (					A			,   A	)						->   Optional[Any]:
							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*).*'''			,   A	)[1]	)
							layer_number -= 3
							return F"""h.{layer_number}.""" + key


def  _snake_case       (					A	)						->   Optional[int]:
							if dtype == torch.bool:
														return 1 / 8
							lowerCAmelCase__	       =     re.search(R'''[^\d](\d+)$'''			,   str(A	)	)
							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  _snake_case       (					A			,   A			,   A			,   A			,   A	)						->   Dict:
							# Construct model
							if bloom_config_file == "":
														lowerCAmelCase__	       =     BloomConfig()
							else:
														lowerCAmelCase__	       =     BloomConfig.from_json_file(A	)

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

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

														lowerCAmelCase__	       =     None

														lowerCAmelCase__	       =     BloomConfig()

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

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

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

																												if tensors is None:
																																			lowerCAmelCase__	       =     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
																																																	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=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	):
																																			lowerCAmelCase__	       =     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():
																												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(A	)	).zfill(5	)	)

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

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

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

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

																												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(A	) 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=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	):
																																			lowerCAmelCase__	       =     tensors[key] / pretraining_tp

																					lowerCAmelCase__	       =     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:
																												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(A			,   exist_ok=A	)
														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()			,   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__":
		__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,
		)