


'''simple docstring'''






import itertools
import os
import random
import tempfile
import unittest

import numpy as np

from transformers import TvltFeatureExtractor, is_datasets_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
				import torch

if is_datasets_available():
				from datasets import load_dataset

__UpperCAmelCase       =							random.Random()


def  _snake_case       (					A			,   A=1.0			,   A=None			,   A=None	)						->   List[Any]:
							if rng is None:
														lowerCAmelCase__	       =     global_rng

							lowerCAmelCase__	       =     []
							for batch_idx in range(shape[0]	):
														values.append([]	)
														for _ in range(shape[1]	):
																					values[-1].append(rng.random() * scale	)

							return values




class 		a__    (							unittest.TestCase  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_=7 ,   lowerCamelCase_=4_00 ,   lowerCamelCase_=20_00 ,   lowerCamelCase_=20_48 ,   lowerCamelCase_=1_28 ,   lowerCamelCase_=1 ,   lowerCamelCase_=5_12 ,   lowerCamelCase_=30 ,   lowerCamelCase_=4_41_00 ,   )	->							Any:
														lowerCAmelCase__	       =     parent
														lowerCAmelCase__	       =     batch_size
														lowerCAmelCase__	       =     min_seq_length
														lowerCAmelCase__	       =     max_seq_length
														lowerCAmelCase__	       =     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
														lowerCAmelCase__	       =     spectrogram_length
														lowerCAmelCase__	       =     feature_size
														lowerCAmelCase__	       =     num_audio_channels
														lowerCAmelCase__	       =     hop_length
														lowerCAmelCase__	       =     chunk_length
														lowerCAmelCase__	       =     sampling_rate







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														return {
														    "spectrogram_length": self.spectrogram_length,
														    "feature_size": self.feature_size,
														    "num_audio_channels": self.num_audio_channels,
														    "hop_length": self.hop_length,
														    "chunk_length": self.chunk_length,
														    "sampling_rate": self.sampling_rate,
														}







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_=False ,   lowerCamelCase_=False       )	->							Dict:
														def _flatten(lowerCamelCase_       ):
																					return list(itertools.chain(*lowerCamelCase_       )       )

														if equal_length:
																					lowerCAmelCase__	       =     [floats_list((self.max_seq_length, self.feature_size)       ) for _ in range(self.batch_size       )]
														else:
																					# make sure that inputs increase in size
																					lowerCAmelCase__	       =     [
																					    floats_list((x, self.feature_size)       )
																					    for x in range(self.min_seq_length ,   self.max_seq_length ,   self.seq_length_diff       )
																					]
														if numpify:
																					lowerCAmelCase__	       =     [np.asarray(lowerCamelCase_       ) for x in speech_inputs]
														return speech_inputs




@require_torch
@require_torchaudio
class 		a__    (							a__    ,    unittest.TestCase  ):





							'''simple docstring'''






							lowercase__      :       int     						=				TvltFeatureExtractor







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[str]:
														lowerCAmelCase__	       =     TvltFeatureExtractionTester(self       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''spectrogram_length'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''feature_size'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''num_audio_channels'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''hop_length'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''chunk_length'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''sampling_rate'''       )       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							int:
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )

														with tempfile.TemporaryDirectory() as tmpdirname:
																					lowerCAmelCase__	       =     feat_extract_first.save_pretrained(lowerCamelCase_       )[0]
																					check_json_file_has_correct_format(lowerCamelCase_       )
																					lowerCAmelCase__	       =     self.feature_extraction_class.from_pretrained(lowerCamelCase_       )

														lowerCAmelCase__	       =     feat_extract_first.to_dict()
														lowerCAmelCase__	       =     feat_extract_second.to_dict()
														lowerCAmelCase__	       =     dict_first.pop('''mel_filters'''       )
														lowerCAmelCase__	       =     dict_second.pop('''mel_filters'''       )
														self.assertTrue(np.allclose(lowerCamelCase_ ,   lowerCamelCase_       )       )
														self.assertEqual(lowerCamelCase_ ,   lowerCamelCase_       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )

														with tempfile.TemporaryDirectory() as tmpdirname:
																					lowerCAmelCase__	       =     os.path.join(lowerCamelCase_ ,   '''feat_extract.json'''       )
																					feat_extract_first.to_json_file(lowerCamelCase_       )
																					lowerCAmelCase__	       =     self.feature_extraction_class.from_json_file(lowerCamelCase_       )

														lowerCAmelCase__	       =     feat_extract_first.to_dict()
														lowerCAmelCase__	       =     feat_extract_second.to_dict()
														lowerCAmelCase__	       =     dict_first.pop('''mel_filters'''       )
														lowerCAmelCase__	       =     dict_second.pop('''mel_filters'''       )
														self.assertTrue(np.allclose(lowerCamelCase_ ,   lowerCamelCase_       )       )
														self.assertEqual(lowerCamelCase_ ,   lowerCamelCase_       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														# Initialize feature_extractor
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )

														# create three inputs of length 800, 1000, and 1200
														lowerCAmelCase__	       =     [floats_list((1, x)       )[0] for x in range(8_00 ,   14_00 ,   2_00       )]
														lowerCAmelCase__	       =     [np.asarray(lowerCamelCase_       ) for speech_input in speech_inputs]

														# Test not batched input
														lowerCAmelCase__	       =     feature_extractor(np_speech_inputs[0] ,   return_tensors='''np''' ,   sampling_rate=4_41_00       ).audio_values

														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )

														# Test batched
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   return_tensors='''np''' ,   sampling_rate=4_41_00       ).audio_values

														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )

														# Test audio masking
														lowerCAmelCase__	       =     feature_extractor(
														    lowerCamelCase_ ,   return_tensors='''np''' ,   sampling_rate=4_41_00 ,   mask_audio=lowerCamelCase_       ).audio_values

														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )

														# Test 2-D numpy arrays are batched.
														lowerCAmelCase__	       =     [floats_list((1, x)       )[0] for x in (8_00, 8_00, 8_00)]
														lowerCAmelCase__	       =     np.asarray(lowerCamelCase_       )
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   return_tensors='''np''' ,   sampling_rate=4_41_00       ).audio_values
														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_       )	->							List[str]:
														lowerCAmelCase__	       =     load_dataset('''hf-internal-testing/librispeech_asr_dummy''' ,   '''clean''' ,   split='''validation'''       )
														# automatic decoding with librispeech
														lowerCAmelCase__	       =     ds.sort('''id'''       ).select(range(lowerCamelCase_       )       )[:num_samples]['''audio''']

														return [x["array"] for x in speech_samples]







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[str]:
														lowerCAmelCase__	       =     self._load_datasamples(1       )
														lowerCAmelCase__	       =     TvltFeatureExtractor()
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   return_tensors='''pt'''       ).audio_values

														self.assertEquals(audio_values.shape ,   (1, 1, 1_92, 1_28)       )

														lowerCAmelCase__	       =     torch.tensor([[-0.3_032, -0.2_708], [-0.4_434, -0.4_007]]       )
														self.assertTrue(torch.allclose(audio_values[0, 0, :2, :2] ,   lowerCamelCase_ ,   atol=1e-4       )       )



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