






"""simple docstring"""



import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import WavaVecaProcessor, logging
from transformers.models.dataavec.configuration_dataavec_audio import DataaVecAudioConfig

# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
from transformers.models.dataavec.dataavec_audio import DataaVecAudioModel as Dummy  # noqa: F401
from transformers.models.dataavec.modeling_dataavec_audio import DataaVecAudioForCTC, DataaVecAudioModel


logging.set_verbosity_info()
_lowerCAmelCase     :     List[Any]            =	logging.get_logger(__name__)

_lowerCAmelCase     :     Any            =	{
    '''post_extract_proj''': '''feature_projection.projection''',
    '''models.0.layer_norm''': '''feature_projection.layer_norm''',
    '''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''',
    '''w2v_encoder.proj''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}
_lowerCAmelCase     :     int            =	[
    '''lm_head''',
]


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Optional[int]       , snake_case    :      List[str]       , snake_case    :      Optional[int]       , snake_case    :      Tuple       , snake_case    :      Tuple							)->						List[str]:






     '''simple docstring'''



     for attribute in key.split("."							):
          UpperCAmelCase__		:  Any				=				getattr(UpperCAmelCase__       , UpperCAmelCase__							)

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

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


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Union[str, Any]       , snake_case    :      Optional[Any]       , snake_case    :      int							)->						str:






     '''simple docstring'''



     UpperCAmelCase__		:  Tuple				=				[]
     UpperCAmelCase__		:  Optional[int]				=				fairseq_model.state_dict()

     if not is_headless:
          UpperCAmelCase__		:  Tuple				=				hf_model.dataavec_audio.feature_extractor
          UpperCAmelCase__		:  Union[str, Any]				=				hf_model.dataavec_audio.encoder.pos_conv_embed

     else:
          UpperCAmelCase__		:  List[str]				=				hf_model.feature_extractor
          UpperCAmelCase__		:  Dict				=				hf_model.encoder.pos_conv_embed

     for name, value in fairseq_dict.items():
          UpperCAmelCase__		:  Optional[Any]				=				False
          if "conv_layers" in name:
               load_conv_layer(
                   UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__       , )
               UpperCAmelCase__		:  Union[str, Any]				=				True
          elif "pos_conv" in name:
               load_pos_conv_layer(
                   UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__       , )
               UpperCAmelCase__		:  int				=				True
          else:
               for key, mapped_key in MAPPING.items():
                    if not is_headless:
                         UpperCAmelCase__		:  List[Any]				=				"data2vec_audio." + 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]:
                         UpperCAmelCase__		:  List[str]				=				True
                         if "*" in mapped_key:
                              UpperCAmelCase__		:  Dict				=				name.split(UpperCAmelCase__							)[0].split("."							)[-2]
                              UpperCAmelCase__		:  Optional[int]				=				mapped_key.replace("*"       , UpperCAmelCase__							)
                         if "weight_g" in name:
                              UpperCAmelCase__		:  str				=				"weight_g"
                         elif "weight_v" in name:
                              UpperCAmelCase__		:  Optional[Any]				=				"weight_v"
                         elif "bias" in name:
                              UpperCAmelCase__		:  str				=				"bias"
                         elif "weight" in name:
                              # TODO: don't match quantizer.weight_proj
                              UpperCAmelCase__		:  int				=				"weight"
                         else:
                              UpperCAmelCase__		:  Dict				=				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      SCREAMING_SNAKE_CASE__    (   snake_case    :      Any       , snake_case    :      List[str]							)->						List[Any]:






     '''simple docstring'''



     UpperCAmelCase__		:  Optional[int]				=				path.split("."							)
     return reduce(UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__							)


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Optional[int]       , snake_case    :      List[Any]       , snake_case    :      List[Any]       , snake_case    :      List[str]							)->						Dict:






     '''simple docstring'''



     UpperCAmelCase__		:  str				=				access_by_string(UpperCAmelCase__       , UpperCAmelCase__							)
     UpperCAmelCase__		:  Optional[int]				=				hf_weight.data

     if fsq_value.shape != hf_value.shape:
          raise ValueError(f'{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found.'							)
     UpperCAmelCase__		:  Optional[Any]				=				fsq_value
     logger.info(f'{full_name} was correctly initialized from {hf_weight_path}.'							)


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Optional[Any]       , snake_case    :      Optional[int]       , snake_case    :      Dict       , snake_case    :      int							)->						Any:






     '''simple docstring'''



     UpperCAmelCase__		:  List[Any]				=				full_name.split("conv_layers."							)[-1]
     UpperCAmelCase__		:  Dict				=				name.split("."							)
     UpperCAmelCase__		:  Union[str, Any]				=				int(items[0]							)
     UpperCAmelCase__		:  List[str]				=				int(items[1]							)

     UpperCAmelCase__		:  Any				=				name.split("."							)[-1]
     if type_id == 0:
          UpperCAmelCase__		:  Optional[int]				=				"conv"
     elif type_id == 2:
          UpperCAmelCase__		:  Optional[int]				=				"layer_norm"
     else:
          unused_weights.append(UpperCAmelCase__							)
          return

     set_weights(UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__       , f'conv_layers.{layer_id}.{layer_type}.{weight_type}'							)


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      str       , snake_case    :      Tuple       , snake_case    :      int       , snake_case    :      Optional[int]							)->						Optional[int]:






     '''simple docstring'''



     UpperCAmelCase__		:  int				=				full_name.split("pos_conv."							)[-1]
     UpperCAmelCase__		:  Union[str, Any]				=				name.split("."							)
     UpperCAmelCase__		:  Union[str, Any]				=				int(items[0]							)
     UpperCAmelCase__		:  str				=				int(items[1]							)

     UpperCAmelCase__		:  Any				=				name.split("."							)[-1]
     if type_id != 0:
          unused_weights.append(UpperCAmelCase__							)
          return
     else:
          UpperCAmelCase__		:  Tuple				=				"conv"

     set_weights(UpperCAmelCase__       , UpperCAmelCase__       , UpperCAmelCase__       , f'layers.{layer_id}.{layer_type}.{weight_type}'							)


@torch.no_grad()
def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Any       , snake_case    :      Optional[Any]       , snake_case    :      Optional[Any]=None       , snake_case    :      List[Any]=None       , snake_case    :      Optional[int]=True							)->						Dict:






     '''simple docstring'''



     if config_path is not None:
          UpperCAmelCase__		:  Tuple				=				DataaVecAudioConfig.from_pretrained(UpperCAmelCase__							)
     else:
          UpperCAmelCase__		:  Optional[int]				=				DataaVecAudioConfig()

     if not is_finetuned:
          # Modify final_proj layer name
          UpperCAmelCase__		:  Dict				=				DataaVecAudioModel(UpperCAmelCase__							)
          UpperCAmelCase__		:  Optional[int]				=				os.path.dirname(UpperCAmelCase__							)

          UpperCAmelCase__		:  List[Any]				=				torch.load(UpperCAmelCase__							)
          UpperCAmelCase__		:  Optional[Any]				=				state_dict["model"].pop("final_proj.0.weight"							)
          UpperCAmelCase__		:  Optional[int]				=				state_dict["model"].pop("final_proj.0.bias"							)
          UpperCAmelCase__		:  Optional[int]				=				os.path.join(UpperCAmelCase__       , "converted.pt"							)
          torch.save(UpperCAmelCase__       , UpperCAmelCase__							)
     else:
          UpperCAmelCase__		:  str				=				DataaVecAudioForCTC(UpperCAmelCase__							)
          UpperCAmelCase__		:  Optional[int]				=				checkpoint_path

     def load_dataavec(snake_case    :      Optional[Any]							):
          UpperCAmelCase__				,							UpperCAmelCase__				,							UpperCAmelCase__		:  List[Any]				=				fairseq.checkpoint_utils.load_model_ensemble_and_task([path]							)
          return model[0].eval()

     UpperCAmelCase__		:  Any				=				load_dataavec(UpperCAmelCase__							)

     recursively_load_weights(UpperCAmelCase__       , UpperCAmelCase__       , not is_finetuned							)

     UpperCAmelCase__		:  List[Any]				=				WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"							)

     UpperCAmelCase__		:  Optional[Any]				=				load_dataset("patrickvonplaten/librispeech_asr_dummy"       , "clean"       , split="validation"							)
     UpperCAmelCase__		:  int				=				[x["array"] for x in ds[:4]["audio"]]

     UpperCAmelCase__		:  List[str]				=				processor(UpperCAmelCase__       , return_tensors="pt"       , padding=UpperCAmelCase__							)

     UpperCAmelCase__		:  Optional[int]				=				inputs.input_values
     UpperCAmelCase__		:  Tuple				=				inputs.attention_mask
     #    input_values = inputs.input_values[:, :-1]
     #    attention_mask = inputs.attention_mask[:, :-1]

     hf_wavavec.eval()
     model.eval()
     if is_finetuned:
          UpperCAmelCase__		:  List[str]				=				model(source=UpperCAmelCase__       , padding_mask=(1 - attention_mask)       , mask=UpperCAmelCase__       , features_only=UpperCAmelCase__							)[
              "encoder_out"
          ].transpose(0       , 1							)
          UpperCAmelCase__		:  str				=				hf_wavavec(UpperCAmelCase__       , attention_mask=UpperCAmelCase__							)["logits"]

          UpperCAmelCase__		:  str				=				torch.argmax(UpperCAmelCase__       , dim=-1							)
          UpperCAmelCase__		:  int				=				processor.batch_decode(UpperCAmelCase__							)

          print(f'Expected Output: {ds[:4]["text"]}, Pred: {output_string}'							)
     else:
          UpperCAmelCase__		:  Dict				=				model(source=UpperCAmelCase__       , padding_mask=(1 - attention_mask)       , mask=UpperCAmelCase__       , features_only=UpperCAmelCase__							)[
              "layer_results"
          ][-1][0].transpose(0       , 1							)
          UpperCAmelCase__		:  Tuple				=				hf_wavavec(UpperCAmelCase__       , attention_mask=UpperCAmelCase__							)["last_hidden_state"]

     print(our_output.shape       , their_output.shape							)
     UpperCAmelCase__		:  int				=				torch.max(torch.abs(our_output - their_output							)							).item()
     print(f'max_absolute_diff = {max_absolute_diff}'							)  # ~ 1e-7
     UpperCAmelCase__		:  int				=				torch.allclose(UpperCAmelCase__       , UpperCAmelCase__       , atol=1E-3							)
     print("Do both models output the same tensors?"       , "🔥" if success else "💩"							)
     if not success:
          raise Exception("Something went wRoNg"							)

     hf_wavavec.save_pretrained(UpperCAmelCase__							)

     if is_finetuned:
          processor.save_pretrained(UpperCAmelCase__							)
     else:
          processor.feature_extractor.save_pretrained(UpperCAmelCase__							)


if __name__ == "__main__":
 _lowerCAmelCase     :     Any            =	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"""
 )
 _lowerCAmelCase     :     Union[str, Any]            =	parser.parse_args()
 convert_wavaveca_checkpoint(
     args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
 )



from sklearn.metrics import mean_absolute_error

import datasets


_lowerCamelCase    :      Any	 =						'''\
@article{scikit-learn,
 title={Scikit-learn: Machine Learning in {P}ython},
 author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
 journal={Journal of Machine Learning Research},
 volume={12},
 pages={2825--2830},
 year={2011}
}
'''

_lowerCamelCase    :      str	 =						'''\
Mean Absolute Error (MAE) is the mean of the magnitude of difference between the predicted and actual
values.
'''


_lowerCamelCase    :      List[str]	 =						'''
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {"raw_values", "uniform_average"} or array-like of shape (n_outputs,), default="uniform_average"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

                 "raw_values" : Returns a full set of errors in case of multioutput input.

                 "uniform_average" : Errors of all outputs are averaged with uniform weight.

Returns:
    mae : mean absolute error.
        If multioutput is "raw_values", then mean absolute error is returned for each output separately. If multioutput is "uniform_average" or an ndarray of weights, then the weighted average of all output errors is returned.
        MAE output is non-negative floating point. The best value is 0.0.
Examples:

    >>> mae_metric = datasets.load_metric("mae")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mae_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mae\': 0.5}

    If you\'re using multi-dimensional lists, then set the config as follows :

    >>> mae_metric = datasets.load_metric("mae", "multilist")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mae_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mae\': 0.75}
    >>> results = mae_metric.compute(predictions=predictions, references=references, multioutput=\'raw_values\')
    >>> print(results)
    {\'mae\': array([0.5, 1. ])}
'''


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




							def 		__snake_case(  self      :    int					)    -> str:



												'''simple docstring'''

												return datasets.MetricInfo(
												    description=_DESCRIPTION		,							citation=_CITATION		,							inputs_description=_KWARGS_DESCRIPTION		,							features=datasets.Features(self._get_feature_types()					)		,							reference_urls=[
												        "https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_absolute_error.html"
												    ]		,							)




							def 		__snake_case(  self      :    int					)    -> Tuple:



												'''simple docstring'''

												if self.config_name == "multilist":
																	return {
																	    "predictions": datasets.Sequence(datasets.Value("float"					)					),
																	    "references": datasets.Sequence(datasets.Value("float"					)					),
																	}
												else:
																	return {
																	    "predictions": datasets.Value("float"					),
																	    "references": datasets.Value("float"					),
																	}





							def 		__snake_case(  self      :    Tuple		,							_UpperCamelCase      :    List[str]		,							_UpperCamelCase      :    List[Any]		,							_UpperCamelCase      :    str=None		,							_UpperCamelCase      :    Any="uniform_average"					)    -> int:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE				=				mean_absolute_error(_UpperCamelCase		,							_UpperCamelCase		,							sample_weight=_UpperCamelCase		,							multioutput=_UpperCamelCase					)

												return {"mae": mae_score}




