






"""simple docstring"""





import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    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',
    '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',
}
__lowerCamelCase							=  [
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]



def   a    (     __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    Union[str, Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    Dict      ,   __UpperCAmelCase      :    Union[str, Any]  )    ->							Tuple:
						for attribute in key.split("""."""  ):
												__magic_name__:	Tuple							=       getattr(__UpperCAmelCase      ,   __UpperCAmelCase  )

						if weight_type is not None:
												__magic_name__:	Dict							=       getattr(__UpperCAmelCase      ,   __UpperCAmelCase  ).shape
						else:
												__magic_name__:	Tuple							=       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":
												__magic_name__:	Union[str, Any]							=       value
						elif weight_type == "weight_g":
												__magic_name__:	str							=       value
						elif weight_type == "weight_v":
												__magic_name__:	Optional[Any]							=       value
						elif weight_type == "bias":
												__magic_name__:	Tuple							=       value
						else:
												__magic_name__:	Union[str, Any]							=       value

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



def   a    (     __UpperCAmelCase      :    int      ,   __UpperCAmelCase      :    Optional[Any]  )    ->							Union[str, Any]:
						__magic_name__:	Union[str, Any]							=       []
						__magic_name__:	Tuple							=       fairseq_model.state_dict()

						__magic_name__:	List[str]							=       hf_model.feature_extractor
						__magic_name__:	str							=       hf_model.adapter

						for name, value in fairseq_dict.items():
												__magic_name__:	str							=       False
												if "conv_layers" in name:
																		load_conv_layer(
																		    __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   hf_model.config.feat_extract_norm == """group"""      ,   )
																		__magic_name__:	List[Any]							=       True
												elif any(x in name for x in ["""adaptor""", """w2v_encoder.proj.""", """w2v_proj_ln."""]  ):
																		load_adapter(__UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase  )
																		__magic_name__:	Union[str, Any]							=       True
												else:
																		for key, mapped_key in MAPPING.items():
																								if key in name or key.split("""w2v_model."""  )[-1] == name.split("""."""  )[0]:
																														__magic_name__:	str							=       True
																														if "*" in mapped_key:
																																				__magic_name__:	Dict							=       name.split(__UpperCAmelCase  )[0].split("""."""  )[-2]
																																				__magic_name__:	Optional[int]							=       mapped_key.replace("""*"""      ,   __UpperCAmelCase  )
																														if "weight_g" in name:
																																				__magic_name__:	Optional[Any]							=       """weight_g"""
																														elif "weight_v" in name:
																																				__magic_name__:	Dict							=       """weight_v"""
																														elif "bias" in name:
																																				__magic_name__:	List[Any]							=       """bias"""
																														elif "weight" in name:
																																				__magic_name__:	int							=       """weight"""
																														else:
																																				__magic_name__:	Optional[Any]							=       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   a    (     __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    Any      ,   __UpperCAmelCase      :    Dict      ,   __UpperCAmelCase      :    List[str]  )    ->							Union[str, Any]:
						__magic_name__:	List[Any]							=       full_name.split("""conv_layers."""  )[-1]
						__magic_name__:	str							=       name.split("""."""  )
						__magic_name__:	List[Any]							=       int(items[0]  )
						__magic_name__:	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.'
																		)
																		__magic_name__:	Dict							=       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.'
																		)
																		__magic_name__:	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."
																		)
																		__magic_name__:	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.'
																		)
																		__magic_name__:	Tuple							=       value
																		logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'  )
						else:
												unused_weights.append(__UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    Optional[Any]      ,   __UpperCAmelCase      :    List[Any]      ,   __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    Optional[Any]  )    ->							str:
						__magic_name__:	Dict							=       full_name.split("""adaptor."""  )[-1]
						__magic_name__:	Dict							=       name.split("""."""  )

						if items[1].isdigit():
												__magic_name__:	Optional[int]							=       int(items[1]  )
						else:
												__magic_name__:	Any							=       None

						if "adaptor" not in full_name:
												if "proj_ln" in full_name:
																		# has to be layer norm
																		if "bias" in name:
																								assert (
																								    value.shape == adapter.proj_layer_norm.bias.data.shape
																								), f'{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'
																								__magic_name__:	str							=       value
																								logger.info(f'Adapter proj layer norm bias was initialized from {full_name}.'  )
																		if "weight" in name:
																								assert (
																								    value.shape == adapter.proj_layer_norm.weight.data.shape
																								), f'{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'
																								__magic_name__:	Tuple							=       value
												else:
																		# has to be projection layer
																		if "bias" in name:
																								assert (
																								    value.shape == adapter.proj.bias.data.shape
																								), f'{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'
																								__magic_name__:	Dict							=       value
																								logger.info(f'Adapter proj layer bias was initialized from {full_name}.'  )
																		if "weight" in name:
																								assert (
																								    value.shape == adapter.proj.weight.data.shape
																								), f'{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'
																								__magic_name__:	Optional[int]							=       value
																								logger.info(f'Adapter proj layer weight was initialized from {full_name}.'  )
						elif isinstance(__UpperCAmelCase      ,   __UpperCAmelCase  ):
												if "bias" in name:
																		assert (
																		    value.shape == adapter.layers[layer_id].conv.bias.data.shape
																		), f'{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'
																		__magic_name__:	int							=       value
																		logger.info(f'Adapter layer {layer_id} bias was initialized from {full_name}.'  )
												elif "weight" in name:
																		assert (
																		    value.shape == adapter.layers[layer_id].conv.weight.data.shape
																		), f'{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'
																		__magic_name__:	Optional[Any]							=       value
																		logger.info(f'Adapter layer {layer_id} bias was initialized from {full_name}.'  )
						else:
												unused_weights.append(__UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    List[str]  )    ->							Any:
						__magic_name__,			__magic_name__:	Tuple							=       emb.weight.shape
						__magic_name__:	Union[str, Any]							=       nn.Linear(__UpperCAmelCase      ,   __UpperCAmelCase      ,   bias=__UpperCAmelCase  )
						__magic_name__:	Optional[Any]							=       emb.weight.data
						return lin_layer



@torch.no_grad()
def   a    (     __UpperCAmelCase      :    Tuple      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    Dict      ,   __UpperCAmelCase      :    Dict      ,   __UpperCAmelCase      :    Union[str, Any]      ,   __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    Optional[Any]      ,   __UpperCAmelCase      :    Tuple      ,   __UpperCAmelCase      :    str      ,   __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    Dict      ,   )    ->							Optional[int]:
						__magic_name__:	Optional[int]							=       WavaVecaConfig.from_pretrained(
						    __UpperCAmelCase      ,   add_adapter=__UpperCAmelCase      ,   adapter_stride=__UpperCAmelCase      ,   adapter_kernel_size=__UpperCAmelCase      ,   use_auth_token=__UpperCAmelCase      ,   output_hidden_size=__UpperCAmelCase      ,   )
						__magic_name__:	Optional[int]							=       MBartConfig.from_pretrained(__UpperCAmelCase  )

						# load model
						__magic_name__,			__magic_name__,			__magic_name__:	str							=       fairseq.checkpoint_utils.load_model_ensemble_and_task(
						    [checkpoint_path]      ,   arg_overrides={
						        """config_yaml""": config_yaml_path,
						        """data""": """/""".join(dict_path.split("""/"""  )[:-1]  ),
						        """w2v_path""": checkpoint_path,
						        """load_pretrained_decoder_from""": None,
						    }      ,   )
						__magic_name__:	Tuple							=       model[0].eval()

						# load feature extractor
						__magic_name__:	Tuple							=       WavaVecaFeatureExtractor.from_pretrained(__UpperCAmelCase      ,   use_auth_token=__UpperCAmelCase  )

						# set weights for wav2vec2 encoder
						__magic_name__:	int							=       WavaVecaModel(__UpperCAmelCase  )

						recursively_load_weights_wavaveca(model.encoder      ,   __UpperCAmelCase  )

						# load decoder weights
						__magic_name__:	Optional[Any]							=       MBartForCausalLM(__UpperCAmelCase  )
						__magic_name__,			__magic_name__:	Dict							=       hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()      ,   strict=__UpperCAmelCase  )
						logger.warning(f'The following keys are missing when loading the decoder weights: {missing_keys}'  )
						logger.warning(f'The following keys are unexpected when loading the decoder weights: {unexpected_keys}'  )

						__magic_name__:	List[str]							=       SpeechEncoderDecoderModel(encoder=__UpperCAmelCase      ,   decoder=__UpperCAmelCase  )
						__magic_name__:	Optional[int]							=       False

						__magic_name__:	Union[str, Any]							=       MBartaaTokenizer(__UpperCAmelCase  )
						tokenizer.save_pretrained(__UpperCAmelCase  )

						__magic_name__:	List[Any]							=       hf_wavavec.config.to_dict()
						__magic_name__:	Any							=       tokenizer.pad_token_id
						__magic_name__:	Union[str, Any]							=       tokenizer.bos_token_id
						__magic_name__:	Optional[int]							=       tokenizer.eos_token_id
						__magic_name__:	int							=       """mbart50"""
						__magic_name__:	int							=       """wav2vec2"""

						__magic_name__:	str							=       tokenizer.eos_token_id
						__magic_name__:	Tuple							=       2_5_0_0_0_4
						__magic_name__:	Optional[int]							=       tokenizer.eos_token_id

						__magic_name__:	List[Any]							=       SpeechEncoderDecoderConfig.from_dict(__UpperCAmelCase  )

						hf_wavavec.save_pretrained(__UpperCAmelCase  )
						feature_extractor.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_yaml_path', default=None, type=str, help='Path to yaml file of fine-tuned model')
							parser.add_argument(
							    '--encoder_config_path',
							    default='facebook/wav2vec2-xls-r-1b',
							    type=str,
							    help='Path to hf encoder wav2vec2 checkpoint config',
							)
							parser.add_argument(
							    '--decoder_config_path',
							    default='facebook/mbart-large-50-one-to-many-mmt',
							    type=str,
							    help='Path to hf decoder checkpoint config',
							)
							parser.add_argument('--add_adapter', default=True, type=bool, help='whethere to add model adapter layers')
							parser.add_argument('--adapter_stride', default=2, type=int, help='stride of adapter layers')
							parser.add_argument('--adapter_kernel_size', default=3, type=int, help='kernel size of adapter layers')
							parser.add_argument('--encoder_output_dim', default=10_24, type=int, help='encoder output dim')
							parser.add_argument('--start_token_id', default=25_00_04, type=int, help='`decoder_start_token_id` of model config')

							__lowerCamelCase							=  parser.parse_args()
							convert_wavaveca_checkpoint(
							    args.checkpoint_path,
							    args.pytorch_dump_folder_path,
							    args.dict_path,
							    args.config_yaml_path,
							    encoder_config_path=args.encoder_config_path,
							    decoder_config_path=args.decoder_config_path,
							    add_adapter=args.add_adapter,
							    adapter_kernel_size=args.adapter_kernel_size,
							    adapter_stride=args.adapter_stride,
							    decoder_start_token_id=args.start_token_id,
							    encoder_output_dim=args.encoder_output_dim,
							)











"""simple docstring"""


import heapq as hq
import math
from collections.abc import Iterator




class 						__UpperCamelCase     :
					def __init__(							self							:	Union[str, Any]     , UpperCAmelCase							:	List[Any]   )	->       List[str]:
											lowerCAmelCase   :Dict						=			str(id_   )
											lowerCAmelCase   :str						=			None
											lowerCAmelCase   :Any						=			None
											lowerCAmelCase   :List[str]						=			[]
											lowerCAmelCase   :List[Any]						=			{}  # {vertex:distance}
					def __lt__(							self							:	Union[str, Any]     , UpperCAmelCase							:	Any   )	->       str:
											return self.key < other.key
					def __repr__(							self							:	Union[str, Any]   )	->       Dict:
											return self.id
					def        UpperCAmelCase__ (							self							:	Any     , UpperCAmelCase							:	Tuple   )	->       List[str]:
											self.neighbors.append(UpperCAmelCase   )
					def        UpperCAmelCase__ (							self							:	str     , UpperCAmelCase							:	Optional[Any]     , UpperCAmelCase							:	List[Any]   )	->       Optional[int]:
											lowerCAmelCase   :str						=			weight


def        UpperCAmelCase      (					a__     ,      a__     ,      a__     ,      a__ ):





						'''simple docstring'''






						graph[a - 1].add_neighbor(graph[b - 1] )
						graph[b - 1].add_neighbor(graph[a - 1] )
						# add the edges:
						graph[a - 1].add_edge(graph[b - 1]     ,      a__ )
						graph[b - 1].add_edge(graph[a - 1]     ,      a__ )


def        UpperCAmelCase      (					a__     ,      a__ ):





						'''simple docstring'''






						lowerCAmelCase   :List[str]						=			[]
						for u in graph:
												lowerCAmelCase   :Tuple						=			math.inf
												lowerCAmelCase   :Any						=			None
						lowerCAmelCase   :Dict						=			0
						lowerCAmelCase   :int						=			graph[:]
						while q:
												lowerCAmelCase   :str						=			min(a__ )
												q.remove(a__ )
												for v in u.neighbors:
																		if (v in q) and (u.edges[v.id] < v.key):
																								lowerCAmelCase   :List[str]						=			u
																								lowerCAmelCase   :List[Any]						=			u.edges[v.id]
						for i in range(1     ,      len(a__ ) ):
												a.append((int(graph[i].id ) + 1, int(graph[i].pi.id ) + 1) )
						return a


def        UpperCAmelCase      (					a__     ,      a__ ):





						'''simple docstring'''






						for u in graph:
												lowerCAmelCase   :List[str]						=			math.inf
												lowerCAmelCase   :str						=			None
						lowerCAmelCase   :Optional[int]						=			0

						lowerCAmelCase   :Dict						=			list(a__ )
						hq.heapify(a__ )

						while h:
												lowerCAmelCase   :Union[str, Any]						=			hq.heappop(a__ )
												for v in u.neighbors:
																		if (v in h) and (u.edges[v.id] < v.key):
																								lowerCAmelCase   :int						=			u
																								lowerCAmelCase   :List[Any]						=			u.edges[v.id]
																								hq.heapify(a__ )

						for i in range(1     ,      len(a__ ) ):
												yield (int(graph[i].id ) + 1, int(graph[i].pi.id ) + 1)


def        UpperCAmelCase      (					):





						'''simple docstring'''








if __name__ == "__main__":
		import doctest

		doctest.testmod()