


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()
UpperCAmelCase      =   logging.get_logger(__name__)

UpperCAmelCase      =   {
    """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""",
}
UpperCAmelCase      =   [
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Any:




						"""simple docstring"""

						for attribute in key.split('''.''' ):
												snake_case_	=				getattr(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

						if weight_type is not None:
												snake_case_	=				getattr(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE ).shape
						else:
												snake_case_	=				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_	=				value
						elif weight_type == "weight_g":
												snake_case_	=				value
						elif weight_type == "weight_v":
												snake_case_	=				value
						elif weight_type == "bias":
												snake_case_	=				value
						else:
												snake_case_	=				value

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Optional[int]:




						"""simple docstring"""

						snake_case_	=				[]
						snake_case_	=				fairseq_model.state_dict()

						snake_case_	=				hf_model.feature_extractor
						snake_case_	=				hf_model.adapter

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

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  List[Any]:




						"""simple docstring"""

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  str:




						"""simple docstring"""

						snake_case_	=				full_name.split('''adaptor.''' )[-1]
						snake_case_	=				name.split('''.''' )

						if items[1].isdigit():
												snake_case_	=				int(items[1] )
						else:
												snake_case_	=				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_	=				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_	=				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_	=				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_	=				value
																								logger.info(f'''Adapter proj layer weight was initialized from {full_name}.''' )
						elif isinstance(SCREAMING_SNAKE_CASE , SCREAMING_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.'''
																		snake_case_	=				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_	=				value
																		logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
						else:
												unused_weights.append(SCREAMING_SNAKE_CASE )






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  Any:




						"""simple docstring"""

						snake_case_	,   snake_case_	=				emb.weight.shape
						snake_case_	=				nn.Linear(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , bias=SCREAMING_SNAKE_CASE )
						snake_case_	=				emb.weight.data
						return lin_layer






@torch.no_grad()
def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , )->  Optional[int]:




						"""simple docstring"""

						snake_case_	=				WavaVecaConfig.from_pretrained(
						    SCREAMING_SNAKE_CASE , add_adapter=SCREAMING_SNAKE_CASE , adapter_stride=SCREAMING_SNAKE_CASE , adapter_kernel_size=SCREAMING_SNAKE_CASE , use_auth_token=SCREAMING_SNAKE_CASE , output_hidden_size=SCREAMING_SNAKE_CASE , )
						snake_case_	=				MBartConfig.from_pretrained(SCREAMING_SNAKE_CASE )

						# load model
						snake_case_	,   snake_case_	,   snake_case_	=				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_	=				model[0].eval()

						# load feature extractor
						snake_case_	=				WavaVecaFeatureExtractor.from_pretrained(SCREAMING_SNAKE_CASE , use_auth_token=SCREAMING_SNAKE_CASE )

						# set weights for wav2vec2 encoder
						snake_case_	=				WavaVecaModel(SCREAMING_SNAKE_CASE )

						recursively_load_weights_wavaveca(model.encoder , SCREAMING_SNAKE_CASE )

						# load decoder weights
						snake_case_	=				MBartForCausalLM(SCREAMING_SNAKE_CASE )
						snake_case_	,   snake_case_	=				hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() , strict=SCREAMING_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}''' )

						snake_case_	=				SpeechEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE , decoder=SCREAMING_SNAKE_CASE )
						snake_case_	=				False

						snake_case_	=				MBartaaTokenizer(SCREAMING_SNAKE_CASE )
						tokenizer.save_pretrained(SCREAMING_SNAKE_CASE )

						snake_case_	=				hf_wavavec.config.to_dict()
						snake_case_	=				tokenizer.pad_token_id
						snake_case_	=				tokenizer.bos_token_id
						snake_case_	=				tokenizer.eos_token_id
						snake_case_	=				'''mbart50'''
						snake_case_	=				'''wav2vec2'''

						snake_case_	=				tokenizer.eos_token_id
						snake_case_	=				25_0004
						snake_case_	=				tokenizer.eos_token_id

						snake_case_	=				SpeechEncoderDecoderConfig.from_dict(SCREAMING_SNAKE_CASE )

						hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE )
						feature_extractor.save_pretrained(SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
			UpperCAmelCase      =   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=1024, type=int, help="""encoder output dim""")
			parser.add_argument("""--start_token_id""", default=25_0004, type=int, help="""`decoder_start_token_id` of model config""")

			UpperCAmelCase      =   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,
			)



import numpy as np

import datasets


UpperCAmelCase      =   """
Compute the Mahalanobis Distance

Mahalonobis distance is the distance between a point and a distribution.
And not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.
It was introduced by Prof. P. C. Mahalanobis in 1936
and has been used in various statistical applications ever since
[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]
"""

UpperCAmelCase      =   """\
@article{de2000mahalanobis,
  title={The mahalanobis distance},
  author={De Maesschalck, Roy and Jouan-Rimbaud, Delphine and Massart, D{\'e}sir{\'e} L},
  journal={Chemometrics and intelligent laboratory systems},
  volume={50},
  number={1},
  pages={1--18},
  year={2000},
  publisher={Elsevier}
}
"""

UpperCAmelCase      =   """
Args:
    X: List of datapoints to be compared with the `reference_distribution`.
    reference_distribution: List of datapoints from the reference distribution we want to compare to.
Returns:
    mahalanobis: The Mahalonobis distance for each datapoint in `X`.
Examples:

    >>> mahalanobis_metric = datasets.load_metric(\"mahalanobis\")
    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])
    >>> print(results)
    {'mahalanobis': array([0.5])}
"""







@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION ,						_KWARGS_DESCRIPTION	)
class lowerCAmelCase_						(			datasets.Metric	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       return datasets.MetricInfo(
           description=_DESCRIPTION		,				citation=_CITATION		,				inputs_description=_KWARGS_DESCRIPTION		,				features=datasets.Features(
               {
                   '''X''': datasets.Sequence(datasets.Value('''float'''		,				id='''sequence''' )		,				id='''X''' ),
               } )		,				)





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       # convert to numpy arrays
       snake_case_	=				np.array(_UpperCAmelCase )
       snake_case_	=				np.array(_UpperCAmelCase )

       # Assert that arrays are 2D
       if len(X.shape ) != 2:
             raise ValueError('''Expected `X` to be a 2D vector''' )
       if len(reference_distribution.shape ) != 2:
             raise ValueError('''Expected `reference_distribution` to be a 2D vector''' )
       if reference_distribution.shape[0] < 2:
             raise ValueError(
                 '''Expected `reference_distribution` to be a 2D vector with more than one element in the first dimension''' )

       # Get mahalanobis distance for each prediction
       snake_case_	=				X - np.mean(_UpperCAmelCase )
       snake_case_	=				np.cov(reference_distribution.T )
       try:
             snake_case_	=				np.linalg.inv(_UpperCAmelCase )
       except np.linalg.LinAlgError:
             snake_case_	=				np.linalg.pinv(_UpperCAmelCase )
       snake_case_	=				np.dot(_UpperCAmelCase		,				_UpperCAmelCase )
       snake_case_	=				np.dot(_UpperCAmelCase		,				X_minus_mu.T ).diagonal()

       return {"mahalanobis": mahal_dist}