




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()
__lowerCamelCase   :						Optional[Any]				=			logging.get_logger(__name__)

__lowerCamelCase   :						str				=			{
    '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',
}
__lowerCamelCase   :						Optional[int]				=			[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'projector',
    'classifier',
]







def 		lowercase__   (		__A:	Tuple				):





					'''simple docstring'''






					__magic_name__							:	List[Any]    							=							{}
					with open(_lowerCamelCase							,'''r'''				) as file:
										for line_number, line in enumerate(_lowerCamelCase				):
															__magic_name__							:	Tuple    							=							line.strip()
															if line:
																				__magic_name__							:	List[Any]    							=							line.split()
																				__magic_name__							:	Optional[int]    							=							line_number
																				__magic_name__							:	str    							=							words[0]
																				__magic_name__							:	int    							=							value
					return result







def 		lowercase__   (		__A:	int							,__A:	Tuple							,__A:	Tuple							,__A:	Optional[int]							,__A:	Any				):





					'''simple docstring'''






					for attribute in key.split('''.'''				):
										__magic_name__							:	int    							=							getattr(_lowerCamelCase							,_lowerCamelCase				)

					__magic_name__							:	str    							=							None
					for param_key in PARAM_MAPPING.keys():
										if full_name.endswith(_lowerCamelCase				):
															__magic_name__							:	str    							=							PARAM_MAPPING[full_name.split('''.'''				)[-1]]
															__magic_name__							:	Optional[Any]    							=							"param"

					if weight_type is not None and weight_type != "param":
										__magic_name__							:	int    							=							getattr(_lowerCamelCase							,_lowerCamelCase				).shape
					elif weight_type is not None and weight_type == "param":
										__magic_name__							:	Any    							=							hf_pointer
										for attribute in hf_param_name.split('''.'''				):
															__magic_name__							:	Union[str, Any]    							=							getattr(_lowerCamelCase							,_lowerCamelCase				)
										__magic_name__							:	Optional[int]    							=							shape_pointer.shape

										# let's reduce dimension
										__magic_name__							:	Dict    							=							value[0]
					else:
										__magic_name__							:	Optional[Any]    							=							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":
										__magic_name__							:	Optional[Any]    							=							value
					elif weight_type == "weight_g":
										__magic_name__							:	Union[str, Any]    							=							value
					elif weight_type == "weight_v":
										__magic_name__							:	Optional[Any]    							=							value
					elif weight_type == "bias":
										__magic_name__							:	List[str]    							=							value
					elif weight_type == "param":
										for attribute in hf_param_name.split('''.'''				):
															__magic_name__							:	List[str]    							=							getattr(_lowerCamelCase							,_lowerCamelCase				)
										__magic_name__							:	Optional[Any]    							=							value
					else:
										__magic_name__							:	int    							=							value

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







def 		lowercase__   (		__A:	Any							,__A:	str							,__A:	Tuple							,__A:	str							,__A:	str				):





					'''simple docstring'''






					__magic_name__							:	Tuple    							=							None
					for param_key in PARAM_MAPPING.keys():
										if full_name.endswith(_lowerCamelCase				):
															__magic_name__							:	str    							=							PARAM_MAPPING[full_name.split('''.'''				)[-1]]
															__magic_name__							:	List[str]    							=							"param"

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

					__magic_name__							:	int    							=							value if "lm_head" in full_key else value[0]


__lowerCamelCase   :						int				=			{
    '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 		lowercase__   (		__A:	Optional[int]							,__A:	Optional[int]							,__A:	Any=None							,__A:	List[str]=None				):





					'''simple docstring'''






					__magic_name__							:	Optional[Any]    							=							False
					for key, mapped_key in MAPPING.items():
										__magic_name__							:	int    							=							"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]:
															__magic_name__							:	Dict    							=							True
															if "*" in mapped_key:
																				__magic_name__							:	Tuple    							=							name.split(_lowerCamelCase				)[0].split('''.'''				)[-2]
																				__magic_name__							:	List[str]    							=							mapped_key.replace('''*'''							,_lowerCamelCase				)
															if "weight_g" in name:
																				__magic_name__							:	Any    							=							"weight_g"
															elif "weight_v" in name:
																				__magic_name__							:	Union[str, Any]    							=							"weight_v"
															elif "bias" in name:
																				__magic_name__							:	Any    							=							"bias"
															elif "weight" in name:
																				# TODO: don't match quantizer.weight_proj
																				__magic_name__							:	int    							=							"weight"
															else:
																				__magic_name__							:	str    							=							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 		lowercase__   (		__A:	str							,__A:	List[str]							,__A:	Any				):





					'''simple docstring'''






					__magic_name__							:	str    							=							[]
					__magic_name__							:	Union[str, Any]    							=							fairseq_model.state_dict()

					__magic_name__							:	Optional[int]    							=							hf_model.wavaveca.feature_extractor

					for name, value in fairseq_dict.items():
										__magic_name__							:	List[str]    							=							False
										if "conv_layers" in name:
															load_conv_layer(
															    _lowerCamelCase							,_lowerCamelCase							,_lowerCamelCase							,_lowerCamelCase							,hf_model.config.feat_extract_norm == '''group'''							,)
															__magic_name__							:	int    							=							True
										else:
															__magic_name__							:	List[str]    							=							load_wavaveca_layer(_lowerCamelCase							,_lowerCamelCase							,_lowerCamelCase				)
										if not is_used:
															unused_weights.append(_lowerCamelCase				)

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







def 		lowercase__   (		__A:	str							,__A:	int							,__A:	Any							,__A:	List[Any]							,__A:	Optional[int]				):





					'''simple docstring'''






					__magic_name__							:	Union[str, Any]    							=							full_name.split('''conv_layers.'''				)[-1]
					__magic_name__							:	List[str]    							=							name.split('''.'''				)
					__magic_name__							:	Dict    							=							int(items[0]				)
					__magic_name__							:	Tuple    							=							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.'''				)
															__magic_name__							:	Optional[int]    							=							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.'''				)
															__magic_name__							:	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:
															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.'''				)
															__magic_name__							:	List[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.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'''				)
															__magic_name__							:	str    							=							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 		lowercase__   (		__A:	List[Any]							,__A:	Optional[Any]							,__A:	str=None							,__A:	str=None							,__A:	Tuple=True							,__A:	Optional[int]=False				):





					'''simple docstring'''






					if config_path is not None:
										__magic_name__							:	Union[str, Any]    							=							WavaVecaConfig.from_pretrained(_lowerCamelCase				)
					else:
										__magic_name__							:	Optional[Any]    							=							WavaVecaConfig()

					if is_seq_class:
										__magic_name__							:	List[str]    							=							read_txt_into_dict(_lowerCamelCase				)
										__magic_name__							:	int    							=							idalabel
										__magic_name__							:	Optional[Any]    							=							WavaVecaForSequenceClassification(_lowerCamelCase				)
										__magic_name__							:	Dict    							=							WavaVecaFeatureExtractor(
										    feature_size=1							,sampling_rate=1_6_0_0_0							,padding_value=0							,do_normalize=_lowerCamelCase							,return_attention_mask=_lowerCamelCase							,)
										feature_extractor.save_pretrained(_lowerCamelCase				)

					elif is_finetuned:
										if dict_path:
															__magic_name__							:	Optional[Any]    							=							Dictionary.load(_lowerCamelCase				)

															# important change bos & pad token id since CTC symbol is <pad> and
															# not <s> as in fairseq
															__magic_name__							:	List[Any]    							=							target_dict.pad_index
															__magic_name__							:	Union[str, Any]    							=							target_dict.bos_index
															__magic_name__							:	Optional[int]    							=							target_dict.eos_index
															__magic_name__							:	Tuple    							=							len(target_dict.symbols				)
															__magic_name__							:	Optional[int]    							=							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				)
															__magic_name__							:	Any    							=							target_dict.indices

															# fairseq has the <pad> and <s> switched
															__magic_name__							:	Optional[Any]    							=							0
															__magic_name__							:	Dict    							=							1
															with open(_lowerCamelCase							,'''w'''							,encoding='''utf-8'''				) as vocab_handle:
																				json.dump(_lowerCamelCase							,_lowerCamelCase				)
															__magic_name__							:	Union[str, Any]    							=							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							,)
															__magic_name__							:	List[Any]    							=							True if config.feat_extract_norm == "layer" else False
															__magic_name__							:	Optional[Any]    							=							WavaVecaFeatureExtractor(
															    feature_size=1							,sampling_rate=1_6_0_0_0							,padding_value=0							,do_normalize=_lowerCamelCase							,return_attention_mask=_lowerCamelCase							,)
															__magic_name__							:	Tuple    							=							WavaVecaProcessor(feature_extractor=_lowerCamelCase							,tokenizer=_lowerCamelCase				)
															processor.save_pretrained(_lowerCamelCase				)

										__magic_name__							:	int    							=							WavaVecaForCTC(_lowerCamelCase				)
					else:
										__magic_name__							:	Optional[int]    							=							WavaVecaForPreTraining(_lowerCamelCase				)

					if is_finetuned or is_seq_class:
										__magic_name__							:	Dict    							=							fairseq.checkpoint_utils.load_model_ensemble_and_task(
										    [checkpoint_path]							,arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''				)[:-1]				)}				)
					else:
										__magic_name__							:	Any    							=							argparse.Namespace(task='''audio_pretraining'''				)
										__magic_name__							:	str    							=							fairseq.tasks.setup_task(_lowerCamelCase				)

										__magic_name__							:	List[str]    							=							fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]							,task=_lowerCamelCase				)

					__magic_name__							:	Any    							=							model[0].eval()

					recursively_load_weights(_lowerCamelCase							,_lowerCamelCase							,not is_finetuned				)

					hf_wavavec.save_pretrained(_lowerCamelCase				)


if __name__ == "__main__":
	__lowerCamelCase   :						Dict				=			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''',
	)
	__lowerCamelCase   :						Any				=			parser.parse_args()

	__lowerCamelCase   :						int				=			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,
	)









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase   :						List[Any]				=			{
    '''configuration_speech_to_text''': ['''SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Speech2TextConfig'''],
    '''processing_speech_to_text''': ['''Speech2TextProcessor'''],
}

try:
	if not is_sentencepiece_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			['''Speech2TextTokenizer''']

try:
	if not is_speech_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Optional[Any]				=			['''Speech2TextFeatureExtractor''']

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			[
	    '''TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFSpeech2TextForConditionalGeneration''',
	    '''TFSpeech2TextModel''',
	    '''TFSpeech2TextPreTrainedModel''',
	]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			[
	    '''SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''Speech2TextForConditionalGeneration''',
	    '''Speech2TextModel''',
	    '''Speech2TextPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
	from .processing_speech_to_text import SpeechaTextProcessor

	try:
		if not is_sentencepiece_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_speech_to_text import SpeechaTextTokenizer

	try:
		if not is_speech_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_speech_to_text import (
		    TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFSpeechaTextForConditionalGeneration,
		    TFSpeechaTextModel,
		    TFSpeechaTextPreTrainedModel,
		)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_speech_to_text import (
		    SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    SpeechaTextForConditionalGeneration,
		    SpeechaTextModel,
		    SpeechaTextPreTrainedModel,
		)

else:
	import sys

	__lowerCamelCase   :						List[str]				=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



