



import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


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

lowercase_	       =			{
    '''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''',
    '''adapter_layer''': '''encoder.layers.*.adapter_layer''',
    '''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''',
    '''mask_emb''': '''masked_spec_embed''',
    '''pooling_layer.linear''': '''projector''',
    '''pooling_layer.projection''': '''classifier''',
}
lowercase_	       =			[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
    '''projector''',
    '''classifier''',
]






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Dict	)  ->						Dict:
					__lowerCAmelCase	   ={}
					with open(__lowerCamelCase				,   """r"""	) as file:
										for line_number, line in enumerate(__lowerCamelCase	):
															__lowerCAmelCase	   =line.strip()
															if line:
																				__lowerCAmelCase	   =line.split()
																				__lowerCAmelCase	   =line_number
																				__lowerCAmelCase	   =words[0]
																				__lowerCAmelCase	   =value
					return result






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    Any				,   __lowerCamelCase	:    List[Any]				,   __lowerCamelCase	:    Tuple				,   __lowerCamelCase	:    Optional[Any]	)  ->						List[str]:
					for attribute in key.split("""."""	):
										__lowerCAmelCase	   =getattr(__lowerCamelCase				,   __lowerCamelCase	)

					__lowerCAmelCase	   =None
					for param_key in PARAM_MAPPING.keys():
										if full_name.endswith(__lowerCamelCase	):
															__lowerCAmelCase	   =PARAM_MAPPING[full_name.split("""."""	)[-1]]
															__lowerCAmelCase	   ="""param"""

					if weight_type is not None and weight_type != "param":
										__lowerCAmelCase	   =getattr(__lowerCamelCase				,   __lowerCamelCase	).shape
					elif weight_type is not None and weight_type == "param":
										__lowerCAmelCase	   =hf_pointer
										for attribute in hf_param_name.split("""."""	):
															__lowerCAmelCase	   =getattr(__lowerCamelCase				,   __lowerCamelCase	)
										__lowerCAmelCase	   =shape_pointer.shape

										# let's reduce dimension
										__lowerCAmelCase	   =value[0]
					else:
										__lowerCAmelCase	   =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":
										__lowerCAmelCase	   =value
					elif weight_type == "weight_g":
										__lowerCAmelCase	   =value
					elif weight_type == "weight_v":
										__lowerCAmelCase	   =value
					elif weight_type == "bias":
										__lowerCAmelCase	   =value
					elif weight_type == "param":
										for attribute in hf_param_name.split("""."""	):
															__lowerCAmelCase	   =getattr(__lowerCamelCase				,   __lowerCamelCase	)
										__lowerCAmelCase	   =value
					else:
										__lowerCAmelCase	   =value

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






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    int				,   __lowerCamelCase	:    Any				,   __lowerCamelCase	:    str				,   __lowerCamelCase	:    Optional[Any]	)  ->						Union[str, Any]:
					__lowerCAmelCase	   =None
					for param_key in PARAM_MAPPING.keys():
										if full_name.endswith(__lowerCamelCase	):
															__lowerCAmelCase	   =PARAM_MAPPING[full_name.split("""."""	)[-1]]
															__lowerCAmelCase	   ="""param"""

					if weight_type is not None and weight_type != "param":
										__lowerCAmelCase	   =""".""".join([key, weight_type]	)
					elif weight_type is not None and weight_type == "param":
										__lowerCAmelCase	   =""".""".join([key, hf_param_name]	)
					else:
										__lowerCAmelCase	   =key

					__lowerCAmelCase	   =value if """lm_head""" in full_key else value[0]


lowercase_	       =			{
    '''W_a''': '''linear_1.weight''',
    '''W_b''': '''linear_2.weight''',
    '''b_a''': '''linear_1.bias''',
    '''b_b''': '''linear_2.bias''',
    '''ln_W''': '''norm.weight''',
    '''ln_b''': '''norm.bias''',
}






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Optional[int]				,   __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    str=None				,   __lowerCamelCase	:    Any=None	)  ->						Tuple:
					__lowerCAmelCase	   =False
					for key, mapped_key in MAPPING.items():
										__lowerCAmelCase	   ="""wav2vec2.""" + 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	   =True
															if "*" in mapped_key:
																				__lowerCAmelCase	   =name.split(__lowerCamelCase	)[0].split("""."""	)[-2]
																				__lowerCAmelCase	   =mapped_key.replace("""*"""				,   __lowerCamelCase	)
															if "weight_g" in name:
																				__lowerCAmelCase	   ="""weight_g"""
															elif "weight_v" in name:
																				__lowerCAmelCase	   ="""weight_v"""
															elif "bias" in name:
																				__lowerCAmelCase	   ="""bias"""
															elif "weight" in name:
																				# TODO: don't match quantizer.weight_proj
																				__lowerCAmelCase	   ="""weight"""
															else:
																				__lowerCAmelCase	   =None
															if hf_dict is not None:
																				rename_dict(__lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase	)
															else:
																				set_recursively(__lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase	)
															return is_used
					return is_used






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    Tuple				,   __lowerCamelCase	:    Optional[int]	)  ->						List[str]:
					__lowerCAmelCase	   =[]
					__lowerCAmelCase	   =fairseq_model.state_dict()

					__lowerCAmelCase	   =hf_model.wavaveca.feature_extractor

					for name, value in fairseq_dict.items():
										__lowerCAmelCase	   =False
										if "conv_layers" in name:
															load_conv_layer(
															    __lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase				,   hf_model.config.feat_extract_norm == """group"""				,   )
															__lowerCAmelCase	   =True
										else:
															__lowerCAmelCase	   =load_wavaveca_layer(__lowerCamelCase				,   __lowerCamelCase				,   __lowerCamelCase	)
										if not is_used:
															unused_weights.append(__lowerCamelCase	)

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






def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[Any]				,   __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    str				,   __lowerCamelCase	:    Optional[Any]				,   __lowerCamelCase	:    Dict	)  ->						int:
					__lowerCAmelCase	   =full_name.split("""conv_layers."""	)[-1]
					__lowerCAmelCase	   =name.split("""."""	)
					__lowerCAmelCase	   =int(items[0]	)
					__lowerCAmelCase	   =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."""	)
															__lowerCAmelCase	   =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."""	)
															__lowerCAmelCase	   =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.conv_layers[layer_id].layer_norm.bias.data.shape} was found."""	)
															__lowerCAmelCase	   =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.conv_layers[layer_id].layer_norm.weight.data.shape} was found."""	)
															__lowerCAmelCase	   =value
															logger.info(f"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""	)
					else:
										unused_weights.append(__lowerCamelCase	)






@torch.no_grad()
def 					__lowerCAmelCase		(  __lowerCamelCase	:    Tuple				,   __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    Optional[Any]=None				,   __lowerCamelCase	:    Tuple=None				,   __lowerCamelCase	:    Union[str, Any]=True				,   __lowerCamelCase	:    Union[str, Any]=False	)  ->						int:
					if config_path is not None:
										__lowerCAmelCase	   =WavaVecaConfig.from_pretrained(__lowerCamelCase	)
					else:
										__lowerCAmelCase	   =WavaVecaConfig()

					if is_seq_class:
										__lowerCAmelCase	   =read_txt_into_dict(__lowerCamelCase	)
										__lowerCAmelCase	   =idalabel
										__lowerCAmelCase	   =WavaVecaForSequenceClassification(__lowerCamelCase	)
										__lowerCAmelCase	   =WavaVecaFeatureExtractor(
										    feature_size=1				,   sampling_rate=16000				,   padding_value=0				,   do_normalize=__lowerCamelCase				,   return_attention_mask=__lowerCamelCase				,   )
										feature_extractor.save_pretrained(__lowerCamelCase	)

					elif is_finetuned:
										if dict_path:
															__lowerCAmelCase	   =Dictionary.load(__lowerCamelCase	)

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

															# fairseq has the <pad> and <s> switched
															__lowerCAmelCase	   =0
															__lowerCAmelCase	   =1
															with open(__lowerCamelCase				,   """w"""				,   encoding="""utf-8"""	) as vocab_handle:
																				json.dump(__lowerCamelCase				,   __lowerCamelCase	)
															__lowerCAmelCase	   =WavaVecaCTCTokenizer(
															    __lowerCamelCase				,   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=__lowerCamelCase				,   )
															__lowerCAmelCase	   =True if config.feat_extract_norm == """layer""" else False
															__lowerCAmelCase	   =WavaVecaFeatureExtractor(
															    feature_size=1				,   sampling_rate=16000				,   padding_value=0				,   do_normalize=__lowerCamelCase				,   return_attention_mask=__lowerCamelCase				,   )
															__lowerCAmelCase	   =WavaVecaProcessor(feature_extractor=__lowerCamelCase				,   tokenizer=__lowerCamelCase	)
															processor.save_pretrained(__lowerCamelCase	)

										__lowerCAmelCase	   =WavaVecaForCTC(__lowerCamelCase	)
					else:
										__lowerCAmelCase	   =WavaVecaForPreTraining(__lowerCamelCase	)

					if is_finetuned or is_seq_class:
										__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase	   =fairseq.checkpoint_utils.load_model_ensemble_and_task(
										    [checkpoint_path]				,   arg_overrides={"""data""": """/""".join(dict_path.split("""/"""	)[:-1]	)}	)
					else:
										__lowerCAmelCase	   =argparse.Namespace(task="""audio_pretraining"""	)
										__lowerCAmelCase	   =fairseq.tasks.setup_task(__lowerCamelCase	)

										__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase	   =fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]				,   task=__lowerCamelCase	)

					__lowerCAmelCase	   =model[0].eval()

					recursively_load_weights(__lowerCamelCase				,   __lowerCamelCase				,   not is_finetuned	)

					hf_wavavec.save_pretrained(__lowerCamelCase	)


if __name__ == "__main__":
				lowercase_	       =			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'''
				)
				parser.add_argument(
				    '''--is_seq_class''',
				    action='''store_true''',
				    help='''Whether the model to convert is a fine-tuned sequence classification model or not''',
				)
				lowercase_	       =			parser.parse_args()

				lowercase_	       =			not args.not_finetuned and not args.is_seq_class
				convert_wavaveca_checkpoint(
				    args.checkpoint_path,
				    args.pytorch_dump_folder_path,
				    args.config_path,
				    args.dict_path,
				    is_finetuned,
				    args.is_seq_class,
				)









def 					__lowerCAmelCase		(  __lowerCamelCase	:    list				,   __lowerCamelCase	:    int				,   __lowerCamelCase	:    int = 0				,   __lowerCamelCase	:    int = 0	)  ->						int:
					__lowerCAmelCase	   =right or len(__lowerCamelCase	) - 1
					if left > right:
										return -1
					elif list_data[left] == key:
										return left
					elif list_data[right] == key:
										return right
					else:
										return search(__lowerCamelCase				,   __lowerCamelCase				,   left + 1				,   right - 1	)


if __name__ == "__main__":
				import doctest

				doctest.testmod()




