



'''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							:int				    =					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''',
    '''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							:Optional[int]				    =					[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			for attribute in key.split("""."""    ):
						A_	:						List[Any]	=     getattr(__snake_case						,      __snake_case    )

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

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



def 	a			( lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						List[str]	=     []
			A_	:						str	=     fairseq_model.state_dict()

			A_	:						Optional[int]	=     hf_model.feature_extractor
			A_	:						Optional[Any]	=     hf_model.adapter

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

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



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						str	=     full_name.split("""conv_layers."""    )[-1]
			A_	:						int	=     name.split("""."""    )
			A_	:						Optional[int]	=     int(items[0]    )
			A_	:						Optional[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.'
									)
									A_	:						List[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.'
									)
									A_	:						Optional[int]	=     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."
									)
									A_	:						Dict	=     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.'
									)
									A_	:						int	=     value
									logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'    )
			else:
						unused_weights.append(__snake_case    )



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						str	=     full_name.split("""adaptor."""    )[-1]
			A_	:						Tuple	=     name.split("""."""    )

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



def 	a			( lowerCamelCase__    ):
			'''simple docstring'''





			A_,    A_	:						Tuple	=     emb.weight.shape
			A_	:						str	=     nn.Linear(__snake_case						,      __snake_case						,      bias=__snake_case    )
			A_	:						int	=     emb.weight.data
			return lin_layer



@torch.no_grad()
def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      ):
			'''simple docstring'''





			A_	:						Tuple	=     WavaVecaConfig.from_pretrained(
			    __snake_case						,      add_adapter=__snake_case						,      adapter_stride=__snake_case						,      adapter_kernel_size=__snake_case						,      use_auth_token=__snake_case						,      output_hidden_size=__snake_case						,      )
			A_	:						Union[str, Any]	=     MBartConfig.from_pretrained(__snake_case    )

			# load model
			A_,    A_,    A_	:						List[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,
			    }						,      )
			A_	:						str	=     model[0].eval()

			# load feature extractor
			A_	:						Tuple	=     WavaVecaFeatureExtractor.from_pretrained(__snake_case						,      use_auth_token=__snake_case    )

			# set weights for wav2vec2 encoder
			A_	:						List[str]	=     WavaVecaModel(__snake_case    )

			recursively_load_weights_wavaveca(model.encoder						,      __snake_case    )

			# load decoder weights
			A_	:						Any	=     MBartForCausalLM(__snake_case    )
			A_,    A_	:						Any	=     hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()						,      strict=__snake_case    )
			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}'    )

			A_	:						Any	=     SpeechEncoderDecoderModel(encoder=__snake_case						,      decoder=__snake_case    )
			A_	:						Any	=     False

			A_	:						str	=     MBartaaTokenizer(__snake_case    )
			tokenizer.save_pretrained(__snake_case    )

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

			A_	:						Optional[int]	=     tokenizer.eos_token_id
			A_	:						List[Any]	=     25_00_04
			A_	:						Optional[Any]	=     tokenizer.eos_token_id

			A_	:						Union[str, Any]	=     SpeechEncoderDecoderConfig.from_dict(__snake_case    )

			hf_wavavec.save_pretrained(__snake_case    )
			feature_extractor.save_pretrained(__snake_case    )


if __name__ == "__main__":
						lowerCamelCase							:Tuple				    =					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_0_2_4, type=int, help='''encoder output dim''')
						parser.add_argument('''--start_token_id''', default=2_5_0_0_0_4, type=int, help='''`decoder_start_token_id` of model config''')

						lowerCamelCase							:List[str]				    =					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 List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


_a														=		logging.get_logger(__name__)






class     _UpperCAmelCase(    lowerCamelCase	):



		def 						UpperCAmelCase							(					self				,		__a)       ->       Tuple:






							'''simple docstring'''
							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [label.strip() for label in labels.split(''',''') if label.strip()]
							return labels




		def __call__(					self				,		__a				,		__a				,		__a)       ->       Optional[int]:






							'''simple docstring'''
							if len(__a) == 0 or len(__a) == 0:
												raise ValueError('''You must include at least one label and at least one sequence.''')
							if hypothesis_template.format(labels[0]) == hypothesis_template:
												raise ValueError(
												    (
												        '''The provided hypothesis_template "{}" was not able to be formatted with the target labels. '''
												        '''Make sure the passed template includes formatting syntax such as {{}} where the label should go.'''
												    ).format(__a))

							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [sequences]

							_UpperCamelCase							 =  []
							for sequence in sequences:
												sequence_pairs.extend([[sequence, hypothesis_template.format(__a)] for label in labels])

							return sequence_pairs, sequences






@add_end_docstrings(lowerCamelCase	)
class     _UpperCAmelCase(    lowerCamelCase	):



		def __init__(					self				,		__a=ZeroShotClassificationArgumentHandler()				,		*__a				,		**__a)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  args_parser
							super().__init__(*__a				,		**__a)
							if self.entailment_id == -1:
												logger.warning(
												    '''Failed to determine \'entailment\' label id from the label2id mapping in the model config. Setting to '''
												    '''-1. Define a descriptive label2id mapping in the model config to ensure correct outputs.''')



		@property
		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							for label, ind in self.model.config.labelaid.items():
												if label.lower().startswith('''entail'''):
																	return ind
							return -1



		def 						UpperCAmelCase							(					self				,		__a				,		__a=True				,		__a=True				,		__a=TruncationStrategy.ONLY_FIRST				,		**__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  self.framework
							if self.tokenizer.pad_token is None:
												# Override for tokenizers not supporting padding
												logger.error(
												    '''Tokenizer was not supporting padding necessary for zero-shot, attempting to use '''
												    ''' `pad_token=eos_token`''')
												_UpperCamelCase							 =  self.tokenizer.eos_token
							try:
												_UpperCamelCase							 =  self.tokenizer(
												    __a				,		add_special_tokens=__a				,		return_tensors=__a				,		padding=__a				,		truncation=__a				,		)
							except Exception as e:
												if "too short" in str(__a):
																	# tokenizers might yell that we want to truncate
																	# to a value that is not even reached by the input.
																	# In that case we don't want to truncate.
																	# It seems there's not a really better way to catch that
																	# exception.

																	_UpperCamelCase							 =  self.tokenizer(
																	    __a				,		add_special_tokens=__a				,		return_tensors=__a				,		padding=__a				,		truncation=TruncationStrategy.DO_NOT_TRUNCATE				,		)
												else:
																	raise e

							return inputs



		def 						UpperCAmelCase							(					self				,		**__a)       ->       Any:






							'''simple docstring'''
							if kwargs.get('''multi_class'''				,		__a) is not None:
												_UpperCamelCase							 =  kwargs['''multi_class''']
												logger.warning(
												    '''The `multi_class` argument has been deprecated and renamed to `multi_label`. '''
												    '''`multi_class` will be removed in a future version of Transformers.''')
							_UpperCamelCase							 =  {}
							if "candidate_labels" in kwargs:
												_UpperCamelCase							 =  self._args_parser._parse_labels(kwargs['''candidate_labels'''])
							if "hypothesis_template" in kwargs:
												_UpperCamelCase							 =  kwargs['''hypothesis_template''']

							_UpperCamelCase							 =  {}
							if "multi_label" in kwargs:
												_UpperCamelCase							 =  kwargs['''multi_label''']
							return preprocess_params, {}, postprocess_params



		def __call__(					self				,		__a				,		*__a				,		**__a				,		)       ->       int:






							'''simple docstring'''
							if len(__a) == 0:
												pass
							elif len(__a) == 1 and "candidate_labels" not in kwargs:
												_UpperCamelCase							 =  args[0]
							else:
												raise ValueError(F'''Unable to understand extra arguments {args}''')

							return super().__call__(__a				,		**__a)



		def 						UpperCAmelCase							(					self				,		__a				,		__a=None				,		__a="This example is {}.")       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self._args_parser(__a				,		__a				,		__a)

							for i, (candidate_label, sequence_pair) in enumerate(zip(__a				,		__a)):
												_UpperCamelCase							 =  self._parse_and_tokenize([sequence_pair])

												yield {
												    "candidate_label": candidate_label,
												    "sequence": sequences[0],
												    "is_last": i == len(__a) - 1,
												    **model_input,
												}



		def 						UpperCAmelCase							(					self				,		__a)       ->       List[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  inputs['''candidate_label''']
							_UpperCamelCase							 =  inputs['''sequence''']
							_UpperCamelCase							 =  {k: inputs[k] for k in self.tokenizer.model_input_names}
							_UpperCamelCase							 =  self.model(**__a)

							_UpperCamelCase							 =  {
							    '''candidate_label''': candidate_label,
							    '''sequence''': sequence,
							    '''is_last''': inputs['''is_last'''],
							    **outputs,
							}
							return model_outputs




		def 						UpperCAmelCase							(					self				,		__a				,		__a=False)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  [outputs['''candidate_label'''] for outputs in model_outputs]
							_UpperCamelCase							 =  [outputs['''sequence'''] for outputs in model_outputs]
							_UpperCamelCase							 =  np.concatenate([output['''logits'''].numpy() for output in model_outputs])
							_UpperCamelCase							 =  logits.shape[0]
							_UpperCamelCase							 =  len(__a)
							_UpperCamelCase							 =  N // n
							_UpperCamelCase							 =  logits.reshape((num_sequences, n, -1))

							if multi_label or len(__a) == 1:
												# softmax over the entailment vs. contradiction dim for each label independently
												_UpperCamelCase							 =  self.entailment_id
												_UpperCamelCase							 =  -1 if entailment_id == 0 else 0
												_UpperCamelCase							 =  reshaped_outputs[..., [contradiction_id, entailment_id]]
												_UpperCamelCase							 =  np.exp(__a) / np.exp(__a).sum(-1				,		keepdims=__a)
												_UpperCamelCase							 =  scores[..., 1]
							else:
												# softmax the "entailment" logits over all candidate labels
												_UpperCamelCase							 =  reshaped_outputs[..., self.entailment_id]
												_UpperCamelCase							 =  np.exp(__a) / np.exp(__a).sum(-1				,		keepdims=__a)

							_UpperCamelCase							 =  list(reversed(scores[0].argsort()))
							return {
							    "sequence": sequences[0],
							    "labels": [candidate_labels[i] for i in top_inds],
							    "scores": scores[0, top_inds].tolist(),
							}






