




import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
lowerCamelCase__            = logging.get_logger(__name__)

lowerCamelCase__            = {
    '''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''',
    '''encoder.layer_norm_for_extract''': '''layer_norm_for_extract''',
    '''w2v_model.layer_norm''': '''feature_projection.layer_norm''',
    '''quantizer.weight_proj''': '''quantizer.weight_proj''',
    '''quantizer.vars''': '''quantizer.codevectors''',
    '''project_q''': '''project_q''',
    '''final_proj''': '''project_hid''',
    '''w2v_encoder.proj''': '''lm_head''',
    '''label_embs_concat''': '''label_embeddings_concat''',
    '''mask_emb''': '''masked_spec_embed''',
    '''spk_proj''': '''speaker_proj''',
}
lowerCamelCase__            = [
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
    '''label_embeddings_concat''',
    '''speaker_proj''',
    '''layer_norm_for_extract''',
]







def 	lowercase_		(     SCREAMING_SNAKE_CASE						:		Optional[Any]						,						SCREAMING_SNAKE_CASE						:		Dict						,						SCREAMING_SNAKE_CASE						:		int						,						SCREAMING_SNAKE_CASE						:		int						,						SCREAMING_SNAKE_CASE						:		List[str] ):






	"""simple docstring"""
	for attribute in key.split('''.''' ):
		snake_case__						:     str		=getattr(UpperCamelCase_						,						UpperCamelCase_ )

	if weight_type is not None:
		snake_case__						:     Optional[int]		=getattr(UpperCamelCase_						,						UpperCamelCase_ ).shape
	else:
		snake_case__						:     str		=hf_pointer.shape

	if hf_shape != value.shape:
		raise ValueError(
		    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__						:     int		=value
	elif weight_type == "weight_g":
		snake_case__						:     Optional[Any]		=value
	elif weight_type == "weight_v":
		snake_case__						:     Union[str, Any]		=value
	elif weight_type == "bias":
		snake_case__						:     str		=value
	else:
		snake_case__						:     Optional[int]		=value

	logger.info(F'''{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.''' )







def 	lowercase_		(     SCREAMING_SNAKE_CASE						:		str						,						SCREAMING_SNAKE_CASE						:		Optional[Any] ):






	"""simple docstring"""
	snake_case__						:     Any		=[]
	snake_case__						:     Optional[int]		=fairseq_model.state_dict()

	snake_case__						:     Tuple		=hf_model.unispeech_sat.feature_extractor

	for name, value in fairseq_dict.items():
		snake_case__						:     Optional[int]		=False
		if "conv_layers" in name:
			load_conv_layer(
			    UpperCamelCase_						,						UpperCamelCase_						,						UpperCamelCase_						,						UpperCamelCase_						,						hf_model.config.feat_extract_norm == '''group'''						,						)
			snake_case__						:     Optional[Any]		=True
		else:
			for key, mapped_key in MAPPING.items():
				snake_case__						:     Any		='''unispeech_sat.''' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
				if key in name or key.split('''w2v_model.''' )[-1] == name.split('''.''' )[0]:
					if "layer_norm_for_extract" in name and (".".join(name.split('''.''' )[:-1] ) != key):
						# special case since naming is very similar
						continue
					snake_case__						:     Any		=True
					if "*" in mapped_key:
						snake_case__						:     List[Any]		=name.split(UpperCamelCase_ )[0].split('''.''' )[-2]
						snake_case__						:     Optional[int]		=mapped_key.replace('''*'''						,						UpperCamelCase_ )
					if "weight_g" in name:
						snake_case__						:     Tuple		='''weight_g'''
					elif "weight_v" in name:
						snake_case__						:     List[str]		='''weight_v'''
					elif "bias" in name:
						snake_case__						:     Union[str, Any]		='''bias'''
					elif "weight" in name:
						# TODO: don't match quantizer.weight_proj
						snake_case__						:     List[str]		='''weight'''
					else:
						snake_case__						:     List[str]		=None
					set_recursively(UpperCamelCase_						,						UpperCamelCase_						,						UpperCamelCase_						,						UpperCamelCase_						,						UpperCamelCase_ )
				continue
		if not is_used:
			unused_weights.append(UpperCamelCase_ )

	logger.warning(F'''Unused weights: {unused_weights}''' )







def 	lowercase_		(     SCREAMING_SNAKE_CASE						:		List[str]						,						SCREAMING_SNAKE_CASE						:		Union[str, Any]						,						SCREAMING_SNAKE_CASE						:		Optional[Any]						,						SCREAMING_SNAKE_CASE						:		Optional[int]						,						SCREAMING_SNAKE_CASE						:		Any ):






	"""simple docstring"""
	snake_case__						:     Tuple		=full_name.split('''conv_layers.''' )[-1]
	snake_case__						:     Any		=name.split('''.''' )
	snake_case__						:     Any		=int(items[0] )
	snake_case__						:     Optional[int]		=int(items[1] )

	if type_id == 0:
		if "bias" in name:
			if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
				raise ValueError(
				    F'''{full_name} has size {value.shape}, but'''
				    F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.''' )
			snake_case__						:     Optional[Any]		=value
			logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.''' )
		elif "weight" in name:
			if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
				raise ValueError(
				    F'''{full_name} has size {value.shape}, but'''
				    F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.''' )
			snake_case__						:     Dict		=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:
			if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
				raise ValueError(
				    F'''{full_name} has size {value.shape}, but'''
				    F''' {feature_extractor[layer_id].layer_norm.bias.data.shape} was found.''' )
			snake_case__						:     str		=value
			logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
		elif "weight" in name:
			if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
				raise ValueError(
				    F'''{full_name} has size {value.shape}, but'''
				    F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.''' )
			snake_case__						:     Optional[Any]		=value
			logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.''' )
	else:
		unused_weights.append(UpperCamelCase_ )







@torch.no_grad()
def 	lowercase_		(     SCREAMING_SNAKE_CASE						:		str						,						SCREAMING_SNAKE_CASE						:		List[str]						,						SCREAMING_SNAKE_CASE						:		Union[str, Any]=None						,						SCREAMING_SNAKE_CASE						:		Optional[int]=None						,						SCREAMING_SNAKE_CASE						:		Union[str, Any]=True ):






	"""simple docstring"""
	if config_path is not None:
		snake_case__						:     Union[str, Any]		=UniSpeechSatConfig.from_pretrained(UpperCamelCase_ )
	else:
		snake_case__						:     Optional[Any]		=UniSpeechSatConfig()

	snake_case__						:     Optional[int]		=''''''

	if is_finetuned:
		snake_case__						:     List[str]		=UniSpeechSatForCTC(UpperCamelCase_ )
	else:
		snake_case__						:     List[str]		=UniSpeechSatForPreTraining(UpperCamelCase_ )

	snake_case__,    snake_case__,    snake_case__						:     List[str]		=fairseq.checkpoint_utils.load_model_ensemble_and_task(
	    [checkpoint_path]						,						arg_overrides={'''data''': '''/'''.join(dict_path.split('''/''' )[:-1] )} )
	snake_case__						:     int		=model[0].eval()

	recursively_load_weights(UpperCamelCase_						,						UpperCamelCase_ )

	hf_wavavec.save_pretrained(UpperCamelCase_ )


if __name__ == "__main__":
		lowerCamelCase__            = argparse.ArgumentParser()
		parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
		parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to fairseq checkpoint''')
		parser.add_argument('''--dict_path''', default=None, type=str, help='''Path to dict of fine-tuned model''')
		parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
		parser.add_argument(
		    '''--not_finetuned''', action='''store_true''', help='''Whether the model to convert is a fine-tuned model or not'''
		)
		lowerCamelCase__            = parser.parse_args()
		convert_unispeech_sat_checkpoint(
		    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
		)





import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_SCREAMING_SNAKE_CASE 		=      get_tests_dir("""fixtures""")
_SCREAMING_SNAKE_CASE 		=      get_tests_dir("""fixtures/dummy_feature_extractor_config.json""")
_SCREAMING_SNAKE_CASE 		=      get_tests_dir("""fixtures/dummy-config.json""")
class SCREAMING_SNAKE_CASE_		(      unittest.TestCase		):







   def 						lowerCamelCase_	(     self						:						Tuple				):





    """simple docstring"""


    UpperCamelCase			  =      0







   def 						lowerCamelCase_	(     self						:						int				):





    """simple docstring"""


    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained("""facebook/wav2vec2-base-960h"""				)
    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						Optional[Any]				):





    """simple docstring"""


    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)
    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						Tuple				):





    """simple docstring"""


    with tempfile.TemporaryDirectory() as tmpdirname:
     UpperCamelCase			  =      WavaVecaConfig()

     # remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				).to_dict()

     config_dict.pop("""feature_extractor_type"""				)
     UpperCamelCase			  =      WavaVecaFeatureExtractor(**lowerCamelCase_				)

     # save in new folder
     model_config.save_pretrained(lowerCamelCase_				)
     config.save_pretrained(lowerCamelCase_				)

     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)

     # make sure private variable is not incorrectly saved
     UpperCamelCase			  =      json.loads(config.to_json_string()				)
     self.assertTrue("""_processor_class""" not in dict_as_saved				)

    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						List[str]				):





    """simple docstring"""


    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)
    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						int				):





    """simple docstring"""


    with self.assertRaisesRegex(
        lowerCamelCase_   ,   """bert-base is not a local folder and is not a valid model identifier"""				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained("""bert-base"""				)







   def 						lowerCamelCase_	(     self						:						str				):





    """simple docstring"""


    with self.assertRaisesRegex(
        lowerCamelCase_   ,   R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_   ,   revision="""aaaaaa"""				)







   def 						lowerCamelCase_	(     self						:						Any				):





    """simple docstring"""


    with self.assertRaisesRegex(
        lowerCamelCase_   ,   """hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."""   ,   ):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained("""hf-internal-testing/config-no-model"""				)







   def 						lowerCamelCase_	(     self						:						List[Any]				):





    """simple docstring"""


    with self.assertRaises(lowerCamelCase_				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""				)
    # If remote code is disabled, we can't load this config.
    with self.assertRaises(lowerCamelCase_				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)

    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
        """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)
    self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)

    # Test feature extractor can be reloaded.
    with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(lowerCamelCase_				)
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_   ,   trust_remote_code=lowerCamelCase_				)
    self.assertEqual(reloaded_feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)







   def 						lowerCamelCase_	(     self						:						Union[str, Any]				):





    """simple docstring"""


    try:
     AutoConfig.register("""custom"""   ,   lowerCamelCase_				)
     AutoFeatureExtractor.register(lowerCamelCase_   ,   lowerCamelCase_				)
     # Trying to register something existing in the Transformers library will raise an error
     with self.assertRaises(lowerCamelCase_				):
      AutoFeatureExtractor.register(lowerCamelCase_   ,   lowerCamelCase_				)

     # Now that the config is registered, it can be used as any other config with the auto-API
     UpperCamelCase			  =      CustomFeatureExtractor.from_pretrained(lowerCamelCase_				)
     with tempfile.TemporaryDirectory() as tmp_dir:
      feature_extractor.save_pretrained(lowerCamelCase_				)
      UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)
      self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)

    finally:
     if "custom" in CONFIG_MAPPING._extra_content:
      del CONFIG_MAPPING._extra_content["custom"]
     if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
      del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]







   def 						lowerCamelCase_	(     self						:						Dict				):





    """simple docstring"""


    class SCREAMING_SNAKE_CASE_		(      __lowerCAmelCase		):
       __lowerCAmelCase           =				True

    try:
     AutoConfig.register("""custom"""   ,   lowerCamelCase_				)
     AutoFeatureExtractor.register(lowerCamelCase_   ,   lowerCamelCase_				)
     # If remote code is not set, the default is to use local
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""				)
     self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)
     self.assertTrue(feature_extractor.is_local				)

     # If remote code is disabled, we load the local one.
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)
     self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)
     self.assertTrue(feature_extractor.is_local				)

     # If remote is enabled, we load from the Hub
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)
     self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)
     self.assertTrue(not hasattr(lowerCamelCase_   ,   """is_local"""				)				)

    finally:
     if "custom" in CONFIG_MAPPING._extra_content:
      del CONFIG_MAPPING._extra_content["custom"]
     if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
      del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
