





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

A   : List[Any]             =			{
    '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',
}
A   : int             =			[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]





def       lowercase_      (		lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__						)   ->int:
						for attribute in key.split('.'						):
												_snake_case:   Optional[Any]  					=		getattr(lowercase__				,  lowercase__						)

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

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





def       lowercase_      (		lowercase__				,  lowercase__						)   ->Any:
						_snake_case:   int  					=		[]
						_snake_case:   int  					=		fairseq_model.state_dict()

						_snake_case:   Any  					=		hf_model.feature_extractor
						_snake_case:   Any  					=		hf_model.adapter

						for name, value in fairseq_dict.items():
												_snake_case:   Tuple  					=		False
												if "conv_layers" in name:
																		load_conv_layer(
																		    lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  hf_model.config.feat_extract_norm == 'group'				,  )
																		_snake_case:   Optional[Any]  					=		True
												elif any(x in name for x in ['adaptor', 'w2v_encoder.proj.', 'w2v_proj_ln.']						):
																		load_adapter(lowercase__				,  lowercase__				,  lowercase__				,  lowercase__						)
																		_snake_case:   List[Any]  					=		True
												else:
																		for key, mapped_key in MAPPING.items():
																								if key in name or key.split('w2v_model.'						)[-1] == name.split('.'						)[0]:
																														_snake_case:   List[str]  					=		True
																														if "*" in mapped_key:
																																				_snake_case:   Tuple  					=		name.split(lowercase__						)[0].split('.'						)[-2]
																																				_snake_case:   List[str]  					=		mapped_key.replace('*'				,  lowercase__						)
																														if "weight_g" in name:
																																				_snake_case:   Optional[int]  					=		'weight_g'
																														elif "weight_v" in name:
																																				_snake_case:   Optional[Any]  					=		'weight_v'
																														elif "bias" in name:
																																				_snake_case:   Optional[Any]  					=		'bias'
																														elif "weight" in name:
																																				_snake_case:   Tuple  					=		'weight'
																														else:
																																				_snake_case:   List[Any]  					=		None
																														set_recursively(lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__						)
																								continue
												if not is_used:
																		unused_weights.append(lowercase__						)

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





def       lowercase_      (		lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__						)   ->Any:
						_snake_case:   int  					=		full_name.split('conv_layers.'						)[-1]
						_snake_case:   Tuple  					=		name.split('.'						)
						_snake_case:   str  					=		int(items[0]						)
						_snake_case:   Union[str, Any]  					=		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.'''
																		)
																		_snake_case:   int  					=		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.'''
																		)
																		_snake_case:   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:
																		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."
																		)
																		_snake_case:   Optional[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.'''
																		)
																		_snake_case:   Tuple  					=		value
																		logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''						)
						else:
												unused_weights.append(lowercase__						)





def       lowercase_      (		lowercase__				,  lowercase__				,  lowercase__				,  lowercase__						)   ->List[str]:
						_snake_case:   str  					=		full_name.split('adaptor.'						)[-1]
						_snake_case:   int  					=		name.split('.'						)

						if items[1].isdigit():
												_snake_case:   Dict  					=		int(items[1]						)
						else:
												_snake_case:   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.'''
																								_snake_case:   Optional[int]  					=		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.'''
																								_snake_case:   List[Any]  					=		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.'''
																								_snake_case:   int  					=		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.'''
																								_snake_case:   int  					=		value
																								logger.info(F'''Adapter proj layer weight was initialized from {full_name}.'''						)
						elif isinstance(lowercase__				,  lowercase__						):
												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.'''
																		_snake_case:   str  					=		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.'''
																		_snake_case:   Any  					=		value
																		logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''						)
						else:
												unused_weights.append(lowercase__						)





def       lowercase_      (		lowercase__						)   ->Optional[Any]:
						_snake_case					,						_snake_case:   Optional[int]  					=		emb.weight.shape
						_snake_case:   Tuple  					=		nn.Linear(lowercase__				,  lowercase__				,  bias=lowercase__						)
						_snake_case:   int  					=		emb.weight.data
						return lin_layer





@torch.no_grad()
def       lowercase_      (		lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  lowercase__				,  )   ->Dict:
						_snake_case:   str  					=		WavaVecaConfig.from_pretrained(
						    lowercase__				,  add_adapter=lowercase__				,  adapter_stride=lowercase__				,  adapter_kernel_size=lowercase__				,  use_auth_token=lowercase__				,  output_hidden_size=lowercase__				,  )
						_snake_case:   Any  					=		MBartConfig.from_pretrained(lowercase__						)

						# load model
						_snake_case					,						_snake_case					,						_snake_case:   Tuple  					=		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,
						    }				,  )
						_snake_case:   Union[str, Any]  					=		model[0].eval()

						# load feature extractor
						_snake_case:   Optional[int]  					=		WavaVecaFeatureExtractor.from_pretrained(lowercase__				,  use_auth_token=lowercase__						)

						# set weights for wav2vec2 encoder
						_snake_case:   Any  					=		WavaVecaModel(lowercase__						)

						recursively_load_weights_wavaveca(model.encoder				,  lowercase__						)

						# load decoder weights
						_snake_case:   Dict  					=		MBartForCausalLM(lowercase__						)
						_snake_case					,						_snake_case:   Optional[int]  					=		hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()				,  strict=lowercase__						)
						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}'''						)

						_snake_case:   Tuple  					=		SpeechEncoderDecoderModel(encoder=lowercase__				,  decoder=lowercase__						)
						_snake_case:   Optional[int]  					=		False

						_snake_case:   Dict  					=		MBartaaTokenizer(lowercase__						)
						tokenizer.save_pretrained(lowercase__						)

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

						_snake_case:   Union[str, Any]  					=		tokenizer.eos_token_id
						_snake_case:   Any  					=		25_0004
						_snake_case:   List[str]  					=		tokenizer.eos_token_id

						_snake_case:   Optional[Any]  					=		SpeechEncoderDecoderConfig.from_dict(lowercase__						)

						hf_wavavec.save_pretrained(lowercase__						)
						feature_extractor.save_pretrained(lowercase__						)


if __name__ == "__main__":
		A   : str             =			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=1_024, type=int, help='encoder output dim')
		parser.add_argument('--start_token_id', default=250_004, type=int, help='`decoder_start_token_id` of model config')

		A   : Tuple             =			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'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


A   : Optional[Any]             =			{'configuration_deit': ['DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DeiTConfig', 'DeiTOnnxConfig']}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A   : str             =			['DeiTFeatureExtractor']
		A   : Optional[int]             =			['DeiTImageProcessor']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A   : Optional[Any]             =			[
		    'DEIT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'DeiTForImageClassification',
		    'DeiTForImageClassificationWithTeacher',
		    'DeiTForMaskedImageModeling',
		    'DeiTModel',
		    'DeiTPreTrainedModel',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A   : int             =			[
		    'TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'TFDeiTForImageClassification',
		    'TFDeiTForImageClassificationWithTeacher',
		    'TFDeiTForMaskedImageModeling',
		    'TFDeiTModel',
		    'TFDeiTPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_deit import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, DeiTConfig, DeiTOnnxConfig

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_deit import DeiTFeatureExtractor
				from .image_processing_deit import DeiTImageProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_deit import (
				    DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DeiTForImageClassification,
				    DeiTForImageClassificationWithTeacher,
				    DeiTForMaskedImageModeling,
				    DeiTModel,
				    DeiTPreTrainedModel,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_deit import (
				    TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFDeiTForImageClassification,
				    TFDeiTForImageClassificationWithTeacher,
				    TFDeiTForMaskedImageModeling,
				    TFDeiTModel,
				    TFDeiTPreTrainedModel,
				)


else:
		import sys

		A   : Union[str, Any]             =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





