




import itertools
import random
import unittest

import numpy as np

from transformers import WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST, WavaVecaConfig, WavaVecaFeatureExtractor
from transformers.testing_utils import require_torch, slow

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


__a							     =					random.Random()







def lowerCamelCase__  (     _lowercase					,		_lowercase=1.0					,		_lowercase=None					,		_lowercase=None ):

     '''simple docstring'''


     if rng is None:
          UpperCAmelCase_     :				int  =     global_rng

     UpperCAmelCase_     :				List[Any]  =     []
     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  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE=7  ,_SCREAMING_SNAKE_CASE=400  ,_SCREAMING_SNAKE_CASE=2_000  ,_SCREAMING_SNAKE_CASE=1  ,_SCREAMING_SNAKE_CASE=0.0  ,_SCREAMING_SNAKE_CASE=16_000  ,_SCREAMING_SNAKE_CASE=True  ,_SCREAMING_SNAKE_CASE=True  ,)					->						Dict:
         UpperCAmelCase_     :				int  =     parent
         UpperCAmelCase_     :				Tuple  =     batch_size
         UpperCAmelCase_     :				List[str]  =     min_seq_length
         UpperCAmelCase_     :				Union[str, Any]  =     max_seq_length
         UpperCAmelCase_     :				List[Any]  =     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
         UpperCAmelCase_     :				str  =     feature_size
         UpperCAmelCase_     :				Optional[Any]  =     padding_value
         UpperCAmelCase_     :				Dict  =     sampling_rate
         UpperCAmelCase_     :				List[Any]  =     return_attention_mask
         UpperCAmelCase_     :				Optional[int]  =     do_normalize





    def 				a__ ( self    )					->						Optional[int]:
         return {
             "feature_size": self.feature_size,
             "padding_value": self.padding_value,
             "sampling_rate": self.sampling_rate,
             "return_attention_mask": self.return_attention_mask,
             "do_normalize": self.do_normalize,
         }





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

         if equal_length:
              UpperCAmelCase_     :				int  =     floats_list((self.batch_size, self.max_seq_length)    )
         else:
              # make sure that inputs increase in size
              UpperCAmelCase_     :				Tuple  =     [
                  _flatten(floats_list((x, self.feature_size)    )    )
                  for x in range(self.min_seq_length  ,self.max_seq_length  ,self.seq_length_diff    )
              ]

         if numpify:
              UpperCAmelCase_     :				Tuple  =     [np.asarray(_SCREAMING_SNAKE_CASE    ) for x in speech_inputs]

         return speech_inputs


class     __a(	_a							,		unittest.TestCase  ):






    """simple docstring"""





    lowerCAmelCase         = WavaVecaFeatureExtractor





    def 				a__ ( self    )					->						List[str]:
         UpperCAmelCase_     :				Optional[Any]  =     WavaVecaFeatureExtractionTester(self    )





    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE    )					->						Optional[int]:
         self.assertTrue(np.all(np.mean(_SCREAMING_SNAKE_CASE  ,axis=0    ) < 1e-3    )    )
         self.assertTrue(np.all(np.abs(np.var(_SCREAMING_SNAKE_CASE  ,axis=0    ) - 1    ) < 1e-3    )    )





    def 				a__ ( self    )					->						str:
         # Tests that all call wrap to encode_plus and batch_encode_plus
         UpperCAmelCase_     :				List[Any]  =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
         # create three inputs of length 800, 1000, and 1200
         UpperCAmelCase_     :				Optional[int]  =     [floats_list((1, x)    )[0] for x in range(800  ,1_400  ,200    )]
         UpperCAmelCase_     :				List[Any]  =     [np.asarray(_SCREAMING_SNAKE_CASE    ) for speech_input in speech_inputs]

         # Test not batched input
         UpperCAmelCase_     :				Union[str, Any]  =     feat_extract(speech_inputs[0]  ,return_tensors='''np'''    ).input_values
         UpperCAmelCase_     :				Optional[Any]  =     feat_extract(np_speech_inputs[0]  ,return_tensors='''np'''    ).input_values
         self.assertTrue(np.allclose(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,atol=1e-3    )    )

         # Test batched
         UpperCAmelCase_     :				Union[str, Any]  =     feat_extract(_SCREAMING_SNAKE_CASE  ,return_tensors='''np'''    ).input_values
         UpperCAmelCase_     :				List[Any]  =     feat_extract(_SCREAMING_SNAKE_CASE  ,return_tensors='''np'''    ).input_values
         for enc_seq_a, enc_seq_a in zip(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    ):
              self.assertTrue(np.allclose(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,atol=1e-3    )    )

         # Test 2-D numpy arrays are batched.
         UpperCAmelCase_     :				Optional[Any]  =     [floats_list((1, x)    )[0] for x in (800, 800, 800)]
         UpperCAmelCase_     :				int  =     np.asarray(_SCREAMING_SNAKE_CASE    )
         UpperCAmelCase_     :				int  =     feat_extract(_SCREAMING_SNAKE_CASE  ,return_tensors='''np'''    ).input_values
         UpperCAmelCase_     :				Optional[Any]  =     feat_extract(_SCREAMING_SNAKE_CASE  ,return_tensors='''np'''    ).input_values
         for enc_seq_a, enc_seq_a in zip(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    ):
              self.assertTrue(np.allclose(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,atol=1e-3    )    )





    def 				a__ ( self    )					->						Union[str, Any]:
         UpperCAmelCase_     :				List[str]  =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
         UpperCAmelCase_     :				Any  =     [floats_list((1, x)    )[0] for x in range(800  ,1_400  ,200    )]

         UpperCAmelCase_     :				Tuple  =     ['''longest''', '''max_length''', '''do_not_pad''']
         UpperCAmelCase_     :				List[str]  =     [None, 1_600, None]
         for max_length, padding in zip(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    ):
              UpperCAmelCase_     :				Union[str, Any]  =     feat_extract(_SCREAMING_SNAKE_CASE  ,padding=_SCREAMING_SNAKE_CASE  ,max_length=_SCREAMING_SNAKE_CASE  ,return_tensors='''np'''    )
              UpperCAmelCase_     :				Union[str, Any]  =     processed.input_values

              self._check_zero_mean_unit_variance(input_values[0][:800]    )
              self.assertTrue(input_values[0][800:].sum() < 1e-6    )
              self._check_zero_mean_unit_variance(input_values[1][:1_000]    )
              self.assertTrue(input_values[0][1_000:].sum() < 1e-6    )
              self._check_zero_mean_unit_variance(input_values[2][:1_200]    )





    def 				a__ ( self    )					->						Union[str, Any]:
         UpperCAmelCase_     :				List[str]  =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
         UpperCAmelCase_     :				Union[str, Any]  =     range(800  ,1_400  ,200    )
         UpperCAmelCase_     :				Tuple  =     [floats_list((1, x)    )[0] for x in lengths]

         UpperCAmelCase_     :				Any  =     ['''longest''', '''max_length''', '''do_not_pad''']
         UpperCAmelCase_     :				int  =     [None, 1_600, None]

         for max_length, padding in zip(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    ):
              UpperCAmelCase_     :				Optional[Any]  =     feat_extract(_SCREAMING_SNAKE_CASE  ,max_length=_SCREAMING_SNAKE_CASE  ,padding=_SCREAMING_SNAKE_CASE    )
              UpperCAmelCase_     :				Dict  =     processed.input_values

              self._check_zero_mean_unit_variance(input_values[0][:800]    )
              self._check_zero_mean_unit_variance(input_values[1][:1_000]    )
              self._check_zero_mean_unit_variance(input_values[2][:1_200]    )





    def 				a__ ( self    )					->						Any:
         UpperCAmelCase_     :				Union[str, Any]  =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
         UpperCAmelCase_     :				Optional[int]  =     [floats_list((1, x)    )[0] for x in range(800  ,1_400  ,200    )]
         UpperCAmelCase_     :				List[Any]  =     feat_extract(
             _SCREAMING_SNAKE_CASE  ,truncation=_SCREAMING_SNAKE_CASE  ,max_length=1_000  ,padding='''max_length'''  ,return_tensors='''np'''    )
         UpperCAmelCase_     :				Any  =     processed.input_values

         self._check_zero_mean_unit_variance(input_values[0, :800]    )
         self._check_zero_mean_unit_variance(input_values[1]    )
         self._check_zero_mean_unit_variance(input_values[2]    )





    def 				a__ ( self    )					->						List[Any]:
         UpperCAmelCase_     :				Any  =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
         UpperCAmelCase_     :				Optional[Any]  =     [floats_list((1, x)    )[0] for x in range(800  ,1_400  ,200    )]
         UpperCAmelCase_     :				Union[str, Any]  =     feat_extract(
             _SCREAMING_SNAKE_CASE  ,truncation=_SCREAMING_SNAKE_CASE  ,max_length=1_000  ,padding='''longest'''  ,return_tensors='''np'''    )
         UpperCAmelCase_     :				Union[str, Any]  =     processed.input_values

         self._check_zero_mean_unit_variance(input_values[0, :800]    )
         self._check_zero_mean_unit_variance(input_values[1, :1_000]    )
         self._check_zero_mean_unit_variance(input_values[2]    )

         # make sure that if max_length < longest -> then pad to max_length
         self.assertTrue(input_values.shape == (3, 1_000)    )

         UpperCAmelCase_     :				int  =     [floats_list((1, x)    )[0] for x in range(800  ,1_400  ,200    )]
         UpperCAmelCase_     :				Optional[Any]  =     feat_extract(
             _SCREAMING_SNAKE_CASE  ,truncation=_SCREAMING_SNAKE_CASE  ,max_length=2_000  ,padding='''longest'''  ,return_tensors='''np'''    )
         UpperCAmelCase_     :				str  =     processed.input_values

         self._check_zero_mean_unit_variance(input_values[0, :800]    )
         self._check_zero_mean_unit_variance(input_values[1, :1_000]    )
         self._check_zero_mean_unit_variance(input_values[2]    )

         # make sure that if max_length > longest -> then pad to longest
         self.assertTrue(input_values.shape == (3, 1_200)    )





    @require_torch
    def 				a__ ( self    )					->						Any:
         import torch

         UpperCAmelCase_     :				Dict  =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
         UpperCAmelCase_     :				List[Any]  =     np.random.rand(100    ).astype(np.floataa    )
         UpperCAmelCase_     :				List[str]  =     np_speech_inputs.tolist()

         for inputs in [py_speech_inputs, np_speech_inputs]:
              UpperCAmelCase_     :				Optional[Any]  =     feature_extractor.pad([{'''input_values''': inputs}]  ,return_tensors='''np'''    )
              self.assertTrue(np_processed.input_values.dtype == np.floataa    )
              UpperCAmelCase_     :				Union[str, Any]  =     feature_extractor.pad([{'''input_values''': inputs}]  ,return_tensors='''pt'''    )
              self.assertTrue(pt_processed.input_values.dtype == torch.floataa    )





    @slow
    @require_torch
    def 				a__ ( self    )					->						Dict:
         # this test makes sure that models that are using
         # group norm don't have their feature extractor return the
         # attention_mask
         for model_id in WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST:
              UpperCAmelCase_     :				str  =     WavaVecaConfig.from_pretrained(_SCREAMING_SNAKE_CASE    )
              UpperCAmelCase_     :				List[Any]  =     WavaVecaFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    )

              # only "layer" feature extraction norm should make use of
              # attention_mask
              self.assertEqual(feat_extract.return_attention_mask  ,config.feat_extract_norm == '''layer'''    )





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__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				List[Any]  =     {
         '''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_     :				Union[str, Any]  =     int(re.match(r'''.*layer_(\d*).*'''					,		_lowercase )[1] )
     layer_number -= 3
     return f'''h.{layer_number}.''' + key







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


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







def lowerCamelCase__  (     _lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase ):

     '''simple docstring'''


     if bloom_config_file == "":
          UpperCAmelCase_     :				Tuple  =     BloomConfig()
     else:
          UpperCAmelCase_     :				Optional[int]  =     BloomConfig.from_json_file(_lowercase )

     if shard_model:
          UpperCAmelCase_     :				Any  =     os.listdir(_lowercase )
          UpperCAmelCase_     :				Union[str, Any]  =     sorted(filter(lambda _lowercase : s.startswith('''layer''' ) and "model_00" in s					,		_lowercase ) )

          UpperCAmelCase_     :				Any  =     {'''weight_map''': {}, '''metadata''': {}}
          UpperCAmelCase_     :				List[str]  =     0

          UpperCAmelCase_     :				Any  =     None

          UpperCAmelCase_     :				Optional[int]  =     BloomConfig()

          for j, file in enumerate(_lowercase ):
               print('''Processing file: {}'''.format(_lowercase ) )
               UpperCAmelCase_     :				Optional[Any]  =     None

               for i in range(_lowercase ):
                    # load all TP files
                    UpperCAmelCase_     :				Tuple  =     file.replace('''model_00'''					,		f'''model_0{i}''' )
                    UpperCAmelCase_     :				Any  =     torch.load(os.path.join(_lowercase					,		_lowercase )					,		map_location='''cpu''' )

                    # Rename keys in the transformers names
                    UpperCAmelCase_     :				Dict  =     list(temp.keys() )
                    for key in keys:
                         UpperCAmelCase_     :				Union[str, Any]  =     temp.pop(_lowercase )

                    if tensors is None:
                         UpperCAmelCase_     :				Union[str, Any]  =     temp
                    else:
                         for key in tensors.keys():
                              if any(key.endswith(_lowercase ) 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_     :				int  =     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN ) else 0
                                   # We concatenate these weights accross TP ranks
                                   UpperCAmelCase_     :				Tuple  =     torch.cat([tensors[key], temp[key]]					,		dim=_lowercase )

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

               for key in tensors.keys():
                    UpperCAmelCase_     :				Union[str, Any]  =     tensors[key]
                    total_size += value.numel() * get_dtype_size(value.dtype )
                    if key not in index_dict["weight_map"]:
                         UpperCAmelCase_     :				List[str]  =     '''pytorch_model_{}-of-{}.bin'''.format(
                             str(j + 1 ).zfill(5 )					,		str(len(_lowercase ) ).zfill(5 ) )

          UpperCAmelCase_     :				List[Any]  =     BloomConfig()
          UpperCAmelCase_     :				Tuple  =     pytorch_dump_folder_path + '''/''' + CONFIG_NAME
          UpperCAmelCase_     :				List[str]  =     total_size
          with open(_lowercase					,		'''w'''					,		encoding='''utf-8''' ) as f:
               f.write(config.to_json_string() )
          with open(os.path.join(_lowercase					,		WEIGHTS_NAME + '''.index.json''' )					,		'''w'''					,		encoding='''utf-8''' ) as f:
               UpperCAmelCase_     :				Optional[Any]  =     json.dumps(_lowercase					,		indent=2					,		sort_keys=_lowercase ) + '''\n'''
               f.write(_lowercase )
     else:
          UpperCAmelCase_     :				Any  =     BloomModel(_lowercase )

          UpperCAmelCase_     :				Tuple  =     os.listdir(_lowercase )
          UpperCAmelCase_     :				Union[str, Any]  =     sorted(filter(lambda _lowercase : s.startswith('''layer''' ) and "model_00" in s					,		_lowercase ) )

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

                    # Rename keys in the transformers names
                    UpperCAmelCase_     :				str  =     list(temp.keys() )
                    for key in keys:
                         UpperCAmelCase_     :				Dict  =     temp.pop(_lowercase )

                    if tensors is None:
                         UpperCAmelCase_     :				int  =     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(_lowercase ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH ):
                                   tensors[key] += temp[key]
                              else:
                                   # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                   UpperCAmelCase_     :				Optional[int]  =     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN ) else 0
                                   # We concatenate these weights accross TP ranks
                                   UpperCAmelCase_     :				List[str]  =     torch.cat([tensors[key], temp[key]]					,		dim=_lowercase )

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

               UpperCAmelCase_     :				Tuple  =     model.load_state_dict(_lowercase					,		strict=_lowercase )
               assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
               if missing_keys is None:
                    UpperCAmelCase_     :				Union[str, Any]  =     set(other_keys.missing_keys )
               else:
                    UpperCAmelCase_     :				Dict  =     missing_keys.intersection(set(other_keys.missing_keys ) )

          assert not missing_keys, f'''The keys {missing_keys} are missing'''

          # Save pytorch-model
          os.makedirs(_lowercase					,		exist_ok=_lowercase )
          UpperCAmelCase_     :				str  =     pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
          UpperCAmelCase_     :				Dict  =     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_     :				Optional[int]  =     model.to(config.torch_dtype )
          torch.save(model.state_dict()					,		_lowercase )
          print(f'''Save configuration file to {pytorch_config_dump_path}''' )
          with open(_lowercase					,		'''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,
     )