






import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
_A			:    List[Any]					  =logging.get_logger(__name__)

_A			:    Optional[int]					  ={
    '''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''',
    '''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''': '''ctc_proj''',
    '''mask_emb''': '''masked_spec_embed''',
}
_A			:    Any					  =[
    '''ctc_proj''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]





def 				SCREAMING_SNAKE_CASE_ (UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase   )   -> Union[str, Any]:


							for attribute in key.split("""."""   ):
														if is_finetuned:
																					if attribute in ["quantizer", "project_q", "project_hid"]:
																												# those layers are only relevant for pretraining and should be dropped
																												return

																					if attribute == "ctc_proj":
																												# we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
																												lowerCamelCase__      :			Optional[Any] 						= """lm_head"""

														lowerCamelCase__      :			Any 						= getattr(UpperCamelCase       ,    UpperCamelCase   )

							if weight_type is not None:
														lowerCamelCase__      :			int 						= getattr(UpperCamelCase       ,    UpperCamelCase   ).shape
							else:
														lowerCamelCase__      :			str 						= 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":
														lowerCamelCase__      :			int 						= value
							elif weight_type == "weight_g":
														lowerCamelCase__      :			str 						= value
							elif weight_type == "weight_v":
														lowerCamelCase__      :			Tuple 						= value
							elif weight_type == "bias":
														lowerCamelCase__      :			List[Any] 						= value
							else:
														lowerCamelCase__      :			List[Any] 						= value

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





def 				SCREAMING_SNAKE_CASE_ (UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase   )   -> Optional[int]:
							lowerCamelCase__      :			List[Any] 						= []
							lowerCamelCase__      :			Union[str, Any] 						= fairseq_model.state_dict()

							lowerCamelCase__      :			List[str] 						= hf_model.unispeech.feature_extractor

							for name, value in fairseq_dict.items():
														lowerCamelCase__      :			Union[str, Any] 						= False
														if "conv_layers" in name:
																					load_conv_layer(
																					    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    hf_model.config.feat_extract_norm == """group"""       ,    )
																					lowerCamelCase__      :			Dict 						= True
														else:
																					for key, mapped_key in MAPPING.items():
																												lowerCamelCase__      :			Any 						= """unispeech.""" + 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]:
																																			lowerCamelCase__      :			Dict 						= True
																																			if "*" in mapped_key:
																																										lowerCamelCase__      :			Union[str, Any] 						= name.split(UpperCamelCase   )[0].split("""."""   )[-2]
																																										lowerCamelCase__      :			int 						= mapped_key.replace("""*"""       ,    UpperCamelCase   )
																																			if "weight_g" in name:
																																										lowerCamelCase__      :			Dict 						= """weight_g"""
																																			elif "weight_v" in name:
																																										lowerCamelCase__      :			Optional[Any] 						= """weight_v"""
																																			elif "bias" in name:
																																										lowerCamelCase__      :			Dict 						= """bias"""
																																			elif "weight" in name:
																																										# TODO: don't match quantizer.weight_proj
																																										lowerCamelCase__      :			Optional[int] 						= """weight"""
																																			else:
																																										lowerCamelCase__      :			List[str] 						= None
																																			set_recursively(UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase   )
																												continue
														if not is_used:
																					unused_weights.append(UpperCamelCase   )

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





def 				SCREAMING_SNAKE_CASE_ (UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase   )   -> Optional[int]:
							lowerCamelCase__      :			Tuple 						= full_name.split("""conv_layers."""   )[-1]
							lowerCamelCase__      :			Union[str, Any] 						= name.split("""."""   )
							lowerCamelCase__      :			Optional[int] 						= int(items[0]   )
							lowerCamelCase__      :			Optional[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.'''
																					)
																					lowerCamelCase__      :			Optional[Any] 						= 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.'''
																					)
																					lowerCamelCase__      :			Optional[Any] 						= 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."
																					)
																					lowerCamelCase__      :			List[Any] 						= 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.'''
																					)
																					lowerCamelCase__      :			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(UpperCamelCase   )





@torch.no_grad()
def 				SCREAMING_SNAKE_CASE_ (UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase=None       ,    UpperCamelCase=None       ,    UpperCamelCase=True   )   -> Dict:


							if config_path is not None:
														lowerCamelCase__      :			Optional[int] 						= UniSpeechConfig.from_pretrained(UpperCamelCase   )
							else:
														lowerCamelCase__      :			str 						= UniSpeechConfig()

							if is_finetuned:
														if dict_path:
																					lowerCamelCase__      :			int 						= Dictionary.load_from_json(UpperCamelCase   )

																					# important change bos & pad token id since CTC symbol is <pad> and
																					# not <s> as in fairseq
																					lowerCamelCase__      :			str 						= target_dict.pad_index
																					lowerCamelCase__      :			Tuple 						= target_dict.bos_index
																					lowerCamelCase__      :			Optional[Any] 						= target_dict.eos_index
																					lowerCamelCase__      :			Dict 						= len(target_dict.symbols   )
																					lowerCamelCase__      :			int 						= os.path.join(UpperCamelCase       ,    """vocab.json"""   )
																					if not os.path.isdir(UpperCamelCase   ):
																												logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(UpperCamelCase   )   )
																												return
																					os.makedirs(UpperCamelCase       ,    exist_ok=UpperCamelCase   )
																					lowerCamelCase__      :			str 						= target_dict.indices

																					# fairseq has the <pad> and <s> switched
																					lowerCamelCase__      :			List[Any] 						= 42
																					lowerCamelCase__      :			Any 						= 43
																					with open(UpperCamelCase       ,    """w"""       ,    encoding="""utf-8"""   ) as vocab_handle:
																												json.dump(UpperCamelCase       ,    UpperCamelCase   )
																					lowerCamelCase__      :			List[str] 						= WavaVecaPhonemeCTCTokenizer(
																					    UpperCamelCase       ,    unk_token=target_dict.unk_word       ,    pad_token=target_dict.pad_word       ,    bos_token=target_dict.bos_word       ,    eos_token=target_dict.eos_word       ,    word_delimiter_token="""|"""       ,    do_lower_case=UpperCamelCase       ,    )
																					lowerCamelCase__      :			str 						= True if config.feat_extract_norm == """layer""" else False
																					lowerCamelCase__      :			str 						= WavaVecaFeatureExtractor(
																					    feature_size=1       ,    sampling_rate=16000       ,    padding_value=0       ,    do_normalize=UpperCamelCase       ,    return_attention_mask=UpperCamelCase       ,    )
																					lowerCamelCase__      :			Union[str, Any] 						= WavaVecaProcessor(feature_extractor=UpperCamelCase       ,    tokenizer=UpperCamelCase   )
																					processor.save_pretrained(UpperCamelCase   )

														lowerCamelCase__      :			List[str] 						= UniSpeechForCTC(UpperCamelCase   )
							else:
														lowerCamelCase__      :			List[Any] 						= UniSpeechForPreTraining(UpperCamelCase   )

							if is_finetuned:
														lowerCamelCase__						,	lowerCamelCase__						,	lowerCamelCase__      :			str 						= fairseq.checkpoint_utils.load_model_ensemble_and_task(
														    [checkpoint_path]       ,    arg_overrides={"""data""": """/""".join(dict_path.split("""/"""   )[:-1]   ), """w2v_path""": checkpoint_path}   )
							else:
														lowerCamelCase__						,	lowerCamelCase__						,	lowerCamelCase__      :			List[str] 						= fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]   )

							lowerCamelCase__      :			Dict 						= model[0].eval()

							recursively_load_weights(UpperCamelCase       ,    UpperCamelCase       ,    UpperCamelCase   )

							hf_unispeech.save_pretrained(UpperCamelCase   )


if __name__ == "__main__":
			_A			:    Optional[Any]					  =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'''
			)
			_A			:    Tuple					  =parser.parse_args()
			convert_unispeech_checkpoint(
			    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
			)









'''simple docstring'''







from ...configuration_utils import PretrainedConfig
from ...utils import logging


_A			:    Union[str, Any]					  =logging.get_logger(__name__)

_A			:    List[str]					  ={
    '''MIT/ast-finetuned-audioset-10-10-0.4593''': (
        '''https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json'''
    ),
}

class _lowercase			(       _lowercase	):
							a          = """audio-spectrogram-transformer"""




							def __init__(							self:       str    ,							UpperCamelCase__:       Any=768    ,							UpperCamelCase__:       Union[str, Any]=12    ,							UpperCamelCase__:       List[Any]=12    ,							UpperCamelCase__:       int=3_072    ,							UpperCamelCase__:       Optional[Any]="gelu"    ,							UpperCamelCase__:       Optional[int]=0.0    ,							UpperCamelCase__:       Tuple=0.0    ,							UpperCamelCase__:       Union[str, Any]=0.02    ,							UpperCamelCase__:       Dict=1e-12    ,							UpperCamelCase__:       List[str]=16    ,							UpperCamelCase__:       List[str]=True    ,							UpperCamelCase__:       Any=10    ,							UpperCamelCase__:       List[str]=10    ,							UpperCamelCase__:       Any=1_024    ,							UpperCamelCase__:       Optional[Any]=128    ,							**UpperCamelCase__:       Union[str, Any]    ,							):
														super().__init__(**UpperCamelCase__  )

														lowerCamelCase__      :			List[Any] 						= hidden_size
														lowerCamelCase__      :			int 						= num_hidden_layers
														lowerCamelCase__      :			List[str] 						= num_attention_heads
														lowerCamelCase__      :			Optional[int] 						= intermediate_size
														lowerCamelCase__      :			List[Any] 						= hidden_act
														lowerCamelCase__      :			List[Any] 						= hidden_dropout_prob
														lowerCamelCase__      :			str 						= attention_probs_dropout_prob
														lowerCamelCase__      :			Dict 						= initializer_range
														lowerCamelCase__      :			List[str] 						= layer_norm_eps
														lowerCamelCase__      :			List[Any] 						= patch_size
														lowerCamelCase__      :			List[str] 						= qkv_bias
														lowerCamelCase__      :			Dict 						= frequency_stride
														lowerCamelCase__      :			List[Any] 						= time_stride
														lowerCamelCase__      :			str 						= max_length
														lowerCamelCase__      :			Dict 						= num_mel_bins

