



'''simple docstring'''





import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
						from transformers import WhisperFeatureExtractor, WhisperProcessor


A			:   Optional[Any]			=      50_358
A			:   Union[str, Any]			=      50_362





@require_torch
@require_torchaudio
@require_sentencepiece
class 							lowerCAmelCase_		(	unittest.TestCase						):





				def 					__snake_case       (    self    :    List[str]     ):



							'''simple docstring'''







							snake_case   :							Tuple						  ='''openai/whisper-small.en'''
							snake_case   :							Optional[Any]						  =tempfile.mkdtemp()





				def 					__snake_case       (    self    :    Any,    **_snake_case    :    str     ):



							'''simple docstring'''







							return WhisperTokenizer.from_pretrained(self.checkpoint,    **_snake_case     )





				def 					__snake_case       (    self    :    Tuple,    **_snake_case    :    Any     ):



							'''simple docstring'''







							return WhisperFeatureExtractor.from_pretrained(self.checkpoint,    **_snake_case     )





				def 					__snake_case       (    self    :    str     ):



							'''simple docstring'''







							shutil.rmtree(self.tmpdirname     )





				def 					__snake_case       (    self    :    List[Any]     ):



							'''simple docstring'''







							snake_case   :							str						  =self.get_tokenizer()
							snake_case   :							str						  =self.get_feature_extractor()

							snake_case   :							Dict						  =WhisperProcessor(tokenizer=_snake_case,    feature_extractor=_snake_case     )

							processor.save_pretrained(self.tmpdirname     )
							snake_case   :							Any						  =WhisperProcessor.from_pretrained(self.tmpdirname     )

							self.assertEqual(processor.tokenizer.get_vocab(),    tokenizer.get_vocab()     )
							self.assertIsInstance(processor.tokenizer,    _snake_case     )

							self.assertEqual(processor.feature_extractor.to_json_string(),    feature_extractor.to_json_string()     )
							self.assertIsInstance(processor.feature_extractor,    _snake_case     )





				def 					__snake_case       (    self    :    int     ):



							'''simple docstring'''







							snake_case   :							Union[str, Any]						  =WhisperProcessor(tokenizer=self.get_tokenizer(),    feature_extractor=self.get_feature_extractor()     )
							processor.save_pretrained(self.tmpdirname     )

							snake_case   :							int						  =self.get_tokenizer(bos_token='''(BOS)''',    eos_token='''(EOS)'''     )
							snake_case   :							int						  =self.get_feature_extractor(do_normalize=_snake_case,    padding_value=1.0     )

							snake_case   :							str						  =WhisperProcessor.from_pretrained(
							    self.tmpdirname,    bos_token='''(BOS)''',    eos_token='''(EOS)''',    do_normalize=_snake_case,    padding_value=1.0     )

							self.assertEqual(processor.tokenizer.get_vocab(),    tokenizer_add_kwargs.get_vocab()     )
							self.assertIsInstance(processor.tokenizer,    _snake_case     )

							self.assertEqual(processor.feature_extractor.to_json_string(),    feature_extractor_add_kwargs.to_json_string()     )
							self.assertIsInstance(processor.feature_extractor,    _snake_case     )





				def 					__snake_case       (    self    :    Dict     ):



							'''simple docstring'''







							snake_case   :							Union[str, Any]						  =self.get_feature_extractor()
							snake_case   :							Optional[Any]						  =self.get_tokenizer()

							snake_case   :							List[Any]						  =WhisperProcessor(tokenizer=_snake_case,    feature_extractor=_snake_case     )

							snake_case   :							List[Any]						  =floats_list((3, 1_000)     )

							snake_case   :							str						  =feature_extractor(_snake_case,    return_tensors='''np'''     )
							snake_case   :							Union[str, Any]						  =processor(_snake_case,    return_tensors='''np'''     )

							for key in input_feat_extract.keys():
										self.assertAlmostEqual(input_feat_extract[key].sum(),    input_processor[key].sum(),    delta=1E-2     )





				def 					__snake_case       (    self    :    Optional[Any]     ):



							'''simple docstring'''







							snake_case   :							Dict						  =self.get_feature_extractor()
							snake_case   :							List[Any]						  =self.get_tokenizer()

							snake_case   :							Dict						  =WhisperProcessor(tokenizer=_snake_case,    feature_extractor=_snake_case     )

							snake_case   :							Tuple						  ='''This is a test string'''

							snake_case   :							str						  =processor(text=_snake_case     )

							snake_case   :							Any						  =tokenizer(_snake_case     )

							for key in encoded_tok.keys():
										self.assertListEqual(encoded_tok[key],    encoded_processor[key]     )





				def 					__snake_case       (    self    :    Any     ):



							'''simple docstring'''







							snake_case   :							Union[str, Any]						  =self.get_feature_extractor()
							snake_case   :							Optional[Any]						  =self.get_tokenizer()

							snake_case   :							int						  =WhisperProcessor(tokenizer=_snake_case,    feature_extractor=_snake_case     )

							snake_case   :							List[Any]						  =[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

							snake_case   :							Optional[Any]						  =processor.batch_decode(_snake_case     )
							snake_case   :							Any						  =tokenizer.batch_decode(_snake_case     )

							self.assertListEqual(_snake_case,    _snake_case     )





				def 					__snake_case       (    self    :    Dict     ):



							'''simple docstring'''







							snake_case   :							Optional[int]						  =self.get_feature_extractor()
							snake_case   :							int						  =self.get_tokenizer()

							snake_case   :							List[str]						  =WhisperProcessor(tokenizer=_snake_case,    feature_extractor=_snake_case     )

							self.assertListEqual(
							    processor.model_input_names,    feature_extractor.model_input_names,    msg='''`processor` and `feature_extractor` model input names do not match''',    )





				def 					__snake_case       (    self    :    List[Any]     ):



							'''simple docstring'''







							snake_case   :							Optional[Any]						  =self.get_feature_extractor()
							snake_case   :							List[str]						  =self.get_tokenizer()

							snake_case   :							Dict						  =WhisperProcessor(tokenizer=_snake_case,    feature_extractor=_snake_case     )
							snake_case   :							Dict						  =processor.get_decoder_prompt_ids(task='''transcribe''',    no_timestamps=_snake_case     )

							self.assertIsInstance(_snake_case,    _snake_case     )
							for ids in forced_decoder_ids:
										self.assertIsInstance(_snake_case,    (list, tuple)     )

							snake_case   :							int						  =[TRANSCRIBE, NOTIMESTAMPS]
							self.assertListEqual([ids[-1] for ids in forced_decoder_ids],    _snake_case     )





				def 					__snake_case       (    self    :    Optional[int]     ):



							'''simple docstring'''







							snake_case   :							Union[str, Any]						  =WhisperProcessor(tokenizer=self.get_tokenizer(),    feature_extractor=self.get_feature_extractor()     )
							snake_case   :							Optional[int]						  =processor.get_prompt_ids('''Mr. Quilter'''     )
							snake_case   :							List[str]						  =processor.tokenizer.decode(_snake_case     )

							self.assertListEqual(prompt_ids.tolist(),    [50_360, 1_770, 13, 2_264, 346, 353]     )
							self.assertEqual(_snake_case,    '''<|startofprev|> Mr. Quilter'''     )





				def 					__snake_case       (    self    :    Optional[int]     ):



							'''simple docstring'''







							snake_case   :							Union[str, Any]						  =WhisperProcessor(tokenizer=self.get_tokenizer(),    feature_extractor=self.get_feature_extractor()     )
							snake_case   :							str						  =processor.get_prompt_ids(''''''     )
							snake_case   :							Optional[int]						  =processor.tokenizer.decode(_snake_case     )

							self.assertListEqual(prompt_ids.tolist(),    [50_360, 220]     )
							self.assertEqual(_snake_case,    '''<|startofprev|> '''     )






				def 					__snake_case       (    self    :    Any     ):



							'''simple docstring'''







							snake_case   :							Any						  =WhisperProcessor(tokenizer=self.get_tokenizer(),    feature_extractor=self.get_feature_extractor()     )

							def _test_prompt_error_raised_helper(_snake_case    :    Dict,    _snake_case    :    List[Any]     ):
										with pytest.raises(_snake_case     ) as excinfo:
													processor.get_prompt_ids(_snake_case     )
										snake_case   :							List[str]						  =f'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
										self.assertEqual(_snake_case,    str(excinfo.value     )     )

							_test_prompt_error_raised_helper('''<|startofprev|> test''',    '''<|startofprev|>'''     )
							_test_prompt_error_raised_helper('''test <|notimestamps|>''',    '''<|notimestamps|>'''     )
							_test_prompt_error_raised_helper('''test <|zh|> test <|transcribe|>''',    '''<|zh|>'''     )










'''simple docstring'''





import argparse

import torch
from saprl.hub import distilhubert

from transformers import HubertConfig, HubertModel, WavaVecaFeatureExtractor, logging


logging.set_verbosity_info()
A			:   Dict			=      logging.get_logger(__name__)

A			:   Tuple			=      {
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """mask_emb""": """masked_spec_embed""",
}
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_							):
   for attribute in key.split('''.'''							):
      snake_case   :							Optional[Any]						  =getattr(lowerCamelCase_					,						lowerCamelCase_							)

   if weight_type is not None:
      snake_case   :							int						  =getattr(lowerCamelCase_					,						lowerCamelCase_							).shape
   else:
      snake_case   :							Any						  =hf_pointer.shape

   assert hf_shape == value.shape, (
       F'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
       F''' {value.shape} for {full_name}'''
   )

   if weight_type == "weight":
      snake_case   :							Dict						  =value
   elif weight_type == "weight_g":
      snake_case   :							Any						  =value
   elif weight_type == "weight_v":
      snake_case   :							int						  =value
   elif weight_type == "bias":
      snake_case   :							int						  =value
   else:
      snake_case   :							str						  =value

   logger.info(F'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''							)
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_							):
   snake_case   :							str						  =[]
   snake_case   :							Tuple						  =fairseq_model.state_dict()

   snake_case   :							Union[str, Any]						  =hf_model.feature_extractor

   for name, value in fairseq_dict.items():
      snake_case   :							Optional[int]						  =False
      if "conv_layers" in name:
         load_conv_layer(
             lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						hf_model.config.feat_extract_norm == '''group'''					,						)
         snake_case   :							int						  =True
      else:
         for key, mapped_key in MAPPING.items():
            snake_case   :							Dict						  =mapped_key

            if key in name:
               snake_case   :							str						  =True
               if "*" in mapped_key:
                  snake_case   :							Tuple						  =name.split(lowerCamelCase_							)[0].split('''.'''							)[-2]
                  snake_case   :							Optional[int]						  =mapped_key.replace('''*'''					,						lowerCamelCase_							)
               if "weight_g" in name:
                  snake_case   :							int						  ='''weight_g'''
               elif "weight_v" in name:
                  snake_case   :							List[str]						  ='''weight_v'''
               elif "weight" in name:
                  snake_case   :							Any						  ='''weight'''
               elif "bias" in name:
                  snake_case   :							List[str]						  ='''bias'''
               else:
                  snake_case   :							Tuple						  =None
               set_recursively(lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_							)
            continue
      if not is_used:
         unused_weights.append(lowerCamelCase_							)

   logger.warning(F'''Unused weights: {unused_weights}'''							)
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_							):
   snake_case   :							List[str]						  =full_name.split('''conv_layers.'''							)[-1]
   snake_case   :							Any						  =name.split('''.'''							)
   snake_case   :							Tuple						  =int(items[0]							)
   snake_case   :							int						  =int(items[1]							)

   if type_id == 0:
      if "bias" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
             F'''{full_name} has size {value.shape}, but'''
             F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
         )
         snake_case   :							Optional[int]						  =value
         logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''							)
      elif "weight" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
             F'''{full_name} has size {value.shape}, but'''
             F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
         )
         snake_case   :							str						  =value
         logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''							)
   elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
      if "bias" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
             F'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
             " found."
         )
         snake_case   :							Optional[int]						  =value
         logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
      elif "weight" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
             F'''{full_name} has size {value.shape}, but'''
             F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
         )
         snake_case   :							Union[str, Any]						  =value
         logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
   else:
      unused_weights.append(lowerCamelCase_							)
def 		_a   (  lowerCamelCase_							):
   snake_case   :							Any						  =HubertConfig()
   snake_case   :							List[str]						  =model.config

   snake_case   :							List[str]						  =fs_config.activation_dropout
   snake_case   :							List[Any]						  =False
   snake_case   :							Optional[Any]						  =fs_config.attention_dropout
   snake_case   :							Optional[int]						  =False
   snake_case   :							Any						  =eval(fs_config.extractor_conv_feature_layers							)
   snake_case   :							Optional[int]						  =[x[0] for x in conv_layers]
   snake_case   :							Any						  =[x[1] for x in conv_layers]
   snake_case   :							Optional[Any]						  =[x[2] for x in conv_layers]
   snake_case   :							Optional[Any]						  ='''gelu'''
   snake_case   :							Optional[Any]						  ='''layer''' if fs_config.extractor_mode == '''layer_norm''' else '''group'''
   snake_case   :							Optional[int]						  =False
   snake_case   :							Optional[Any]						  =0.0
   snake_case   :							Optional[int]						  =0.0
   snake_case   :							Dict						  =fs_config.activation_fn
   snake_case   :							Any						  =fs_config.dropout
   snake_case   :							Optional[int]						  =fs_config.encoder_embed_dim
   snake_case   :							Dict						  =0.02
   snake_case   :							Dict						  =fs_config.encoder_ffn_embed_dim
   snake_case   :							str						  =1e-5
   snake_case   :							List[Any]						  =0.0
   snake_case   :							Dict						  =fs_config.encoder_attention_heads
   snake_case   :							Union[str, Any]						  =fs_config.conv_pos_groups
   snake_case   :							Union[str, Any]						  =fs_config.conv_pos
   snake_case   :							Dict						  =len(lowerCamelCase_							)
   snake_case   :							str						  =fs_config.encoder_layers

   return config
@torch.no_grad()
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_=None							):
   snake_case   :							Union[str, Any]						  =distilhubert().model.model

   if config_path is not None:
      snake_case   :							int						  =HubertConfig.from_pretrained(lowerCamelCase_							)
   else:
      snake_case   :							Optional[Any]						  =convert_config(lowerCamelCase_							)
   snake_case   :							Tuple						  =model.eval()

   snake_case   :							Optional[Any]						  =WavaVecaFeatureExtractor(
       feature_size=1					,						sampling_rate=1_60_00					,						padding_value=0					,						do_normalize=lowerCamelCase_					,						return_attention_mask=lowerCamelCase_					,						)
   snake_case   :							Any						  =HubertModel(lowerCamelCase_							)

   recursively_load_weights(lowerCamelCase_					,						lowerCamelCase_							)

   feature_extractor.save_pretrained(lowerCamelCase_							)
   hf_model.save_pretrained(lowerCamelCase_							)


if __name__ == "__main__":
      A			:   Tuple			=      argparse.ArgumentParser()
      parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
      parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
      A			:   int			=      parser.parse_args()
      convert_hubert_checkpoint(args.pytorch_dump_folder_path, args.config_path)





