
import datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


_lowerCamelCase        ="""\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
"""

_lowerCamelCase        ="""
This metric wrap the official scoring script for version 2 of the Stanford Question
Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.

SQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions
written adversarially by crowdworkers to look similar to answerable ones.
To do well on SQuAD2.0, systems must not only answer questions when possible, but also
determine when no answer is supported by the paragraph and abstain from answering.
"""

_lowerCamelCase        ="""
Computes SQuAD v2 scores (F1 and EM).
Args:
    predictions: List of triple for question-answers to score with the following elements:
        - the question-answer \'id\' field as given in the references (see below)
        - the text of the answer
        - the probability that the question has no answer
    references: List of question-answers dictionaries with the following key-values:
            - \'id\': id of the question-answer pair (see above),
            - \'answers\': a list of Dict {\'text\': text of the answer as a string}
    no_answer_threshold: float
        Probability threshold to decide that a question has no answer.
Returns:
    \'exact\': Exact match (the normalized answer exactly match the gold answer)
    \'f1\': The F-score of predicted tokens versus the gold answer
    \'total\': Number of score considered
    \'HasAns_exact\': Exact match (the normalized answer exactly match the gold answer)
    \'HasAns_f1\': The F-score of predicted tokens versus the gold answer
    \'HasAns_total\': Number of score considered
    \'NoAns_exact\': Exact match (the normalized answer exactly match the gold answer)
    \'NoAns_f1\': The F-score of predicted tokens versus the gold answer
    \'NoAns_total\': Number of score considered
    \'best_exact\': Best exact match (with varying threshold)
    \'best_exact_thresh\': No-answer probability threshold associated to the best exact match
    \'best_f1\': Best F1 (with varying threshold)
    \'best_f1_thresh\': No-answer probability threshold associated to the best F1
Examples:

    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\', \'no_answer_probability\': 0.}]
    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> squad_v2_metric = datasets.load_metric(\"squad_v2\")
    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact\': 100.0, \'f1\': 100.0, \'total\': 1, \'HasAns_exact\': 100.0, \'HasAns_f1\': 100.0, \'HasAns_total\': 1, \'best_exact\': 100.0, \'best_exact_thresh\': 0.0, \'best_f1\': 100.0, \'best_f1_thresh\': 0.0}
"""





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,       _KWARGS_DESCRIPTION)
class 							A__		(     datasets.Metric):
      def 						UpperCamelCase__							(				self				):
       return datasets.MetricInfo(
           description=_DESCRIPTION   ,  citation=_CITATION   ,  inputs_description=_KWARGS_DESCRIPTION   ,  features=datasets.Features(
               {
                   """predictions""": {
                       """id""": datasets.Value("""string"""				),
                       """prediction_text""": datasets.Value("""string"""				),
                       """no_answer_probability""": datasets.Value("""float32"""				),
                   },
                   """references""": {
                       """id""": datasets.Value("""string"""				),
                       """answers""": datasets.features.Sequence(
                           {"""text""": datasets.Value("""string"""				), """answer_start""": datasets.Value("""int32"""				)}				),
                   },
               }				)   ,  codebase_urls=["""https://rajpurkar.github.io/SQuAD-explorer/"""]   ,  reference_urls=["""https://rajpurkar.github.io/SQuAD-explorer/"""]   ,  )






      def 						UpperCamelCase__							(				self   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__=1.0				):
       lowerCamelCase		:     int          =				{p['id']: p['no_answer_probability'] for p in predictions}
       lowerCamelCase		:     Union[str, Any]          =				[{'paragraphs': [{'qas': references}]}]
       lowerCamelCase		:     Any          =				{p['id']: p['prediction_text'] for p in predictions}

       lowerCamelCase		:     int          =				make_qid_to_has_ans(__magic_name__				)  # maps qid to True/False
       lowerCamelCase		:     List[str]          =				[k for k, v in qid_to_has_ans.items() if v]
       lowerCamelCase		:     int          =				[k for k, v in qid_to_has_ans.items() if not v]

       lowerCamelCase		:     int          =				get_raw_scores(__magic_name__   ,  __magic_name__				)
       lowerCamelCase		:     Tuple          =				apply_no_ans_threshold(__magic_name__   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__				)
       lowerCamelCase		:     int          =				apply_no_ans_threshold(__magic_name__   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__				)
       lowerCamelCase		:     Optional[Any]          =				make_eval_dict(__magic_name__   ,  __magic_name__				)

       if has_ans_qids:
        lowerCamelCase		:     List[Any]          =				make_eval_dict(__magic_name__   ,  __magic_name__   ,  qid_list=__magic_name__				)
        merge_eval(__magic_name__   ,  __magic_name__   ,  """HasAns"""				)
       if no_ans_qids:
        lowerCamelCase		:     Any          =				make_eval_dict(__magic_name__   ,  __magic_name__   ,  qid_list=__magic_name__				)
        merge_eval(__magic_name__   ,  __magic_name__   ,  """NoAns"""				)
       find_all_best_thresh(__magic_name__   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__   ,  __magic_name__				)
       return dict(__magic_name__				)






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

_lowerCamelCase        ={
    """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        =[
    """lm_head""",
]





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 for attribute in key.split("""."""							):
  lowerCamelCase		:     str          =				getattr(lowerCamelCase,     lowerCamelCase							)

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

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





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     List[str]          =				[]
 lowerCamelCase		:     Union[str, Any]          =				fairseq_model.state_dict()

 if not is_headless:
  lowerCamelCase		:     Tuple          =				hf_model.dataavec_audio.feature_extractor
  lowerCamelCase		:     Tuple          =				hf_model.dataavec_audio.encoder.pos_conv_embed

 else:
  lowerCamelCase		:     Tuple          =				hf_model.feature_extractor
  lowerCamelCase		:     int          =				hf_model.encoder.pos_conv_embed

 for name, value in fairseq_dict.items():
  lowerCamelCase		:     Any          =				False
  if "conv_layers" in name:
   load_conv_layer(
       lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     )
   lowerCamelCase		:     str          =				True
  elif "pos_conv" in name:
   load_pos_conv_layer(
       lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     )
   lowerCamelCase		:     Union[str, Any]          =				True
  else:
   for key, mapped_key in MAPPING.items():
    if not is_headless:
     lowerCamelCase		:     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]:
     lowerCamelCase		:     Union[str, Any]          =				True
     if "*" in mapped_key:
      lowerCamelCase		:     Tuple          =				name.split(lowerCamelCase							)[0].split("""."""							)[-2]
      lowerCamelCase		:     Union[str, Any]          =				mapped_key.replace("""*""",     lowerCamelCase							)
     if "weight_g" in name:
      lowerCamelCase		:     Tuple          =				"""weight_g"""
     elif "weight_v" in name:
      lowerCamelCase		:     Tuple          =				"""weight_v"""
     elif "bias" in name:
      lowerCamelCase		:     List[str]          =				"""bias"""
     elif "weight" in name:
      # TODO: don't match quantizer.weight_proj
      lowerCamelCase		:     List[Any]          =				"""weight"""
     else:
      lowerCamelCase		:     Any          =				None
     set_recursively(lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							)
    continue
  if not is_used:
   unused_weights.append(lowerCamelCase							)

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





def    _a (						lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Union[str, Any]          =				path.split("""."""							)
 return reduce(lowerCamelCase,     lowerCamelCase,     lowerCamelCase							)





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Tuple          =				access_by_string(lowerCamelCase,     lowerCamelCase							)
 lowerCamelCase		:     List[Any]          =				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.'''							)
 lowerCamelCase		:     Dict          =				fsq_value
 logger.info(F'''{full_name} was correctly initialized from {hf_weight_path}.'''							)





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Optional[int]          =				full_name.split("""conv_layers."""							)[-1]
 lowerCamelCase		:     int          =				name.split("""."""							)
 lowerCamelCase		:     Any          =				int(items[0]							)
 lowerCamelCase		:     Union[str, Any]          =				int(items[1]							)

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

 set_weights(lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     F'''conv_layers.{layer_id}.{layer_type}.{weight_type}'''							)





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Dict          =				full_name.split("""pos_conv."""							)[-1]
 lowerCamelCase		:     Any          =				name.split("""."""							)
 lowerCamelCase		:     List[str]          =				int(items[0]							)
 lowerCamelCase		:     List[str]          =				int(items[1]							)

 lowerCamelCase		:     List[Any]          =				name.split("""."""							)[-1]
 if type_id != 0:
  unused_weights.append(lowerCamelCase							)
  return
 else:
  lowerCamelCase		:     Union[str, Any]          =				"""conv"""

 set_weights(lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     F'''layers.{layer_id}.{layer_type}.{weight_type}'''							)





@torch.no_grad()
def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase=None,     lowerCamelCase=None,     lowerCamelCase=True							):
 if config_path is not None:
  lowerCamelCase		:     Optional[Any]          =				DataaVecAudioConfig.from_pretrained(lowerCamelCase							)
 else:
  lowerCamelCase		:     Union[str, Any]          =				DataaVecAudioConfig()

 if not is_finetuned:
  # Modify final_proj layer name
  lowerCamelCase		:     Optional[Any]          =				DataaVecAudioModel(lowerCamelCase							)
  lowerCamelCase		:     List[Any]          =				os.path.dirname(lowerCamelCase							)

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

 def load_dataavec(lowerCamelCase							):
  lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase		:     int          =				fairseq.checkpoint_utils.load_model_ensemble_and_task([path]							)
  return model[0].eval()

 lowerCamelCase		:     Optional[int]          =				load_dataavec(lowerCamelCase							)

 recursively_load_weights(lowerCamelCase,     lowerCamelCase,     not is_finetuned							)

 lowerCamelCase		:     int          =				WavaVecaProcessor.from_pretrained("""facebook/wav2vec2-large-lv60"""							)

 lowerCamelCase		:     Union[str, Any]          =				load_dataset("""patrickvonplaten/librispeech_asr_dummy""",     """clean""",     split="""validation"""							)
 lowerCamelCase		:     Any          =				[x["""array"""] for x in ds[:4]["""audio"""]]

 lowerCamelCase		:     List[Any]          =				processor(lowerCamelCase,     return_tensors="""pt""",     padding=lowerCamelCase							)

 lowerCamelCase		:     Optional[Any]          =				inputs.input_values
 lowerCamelCase		:     int          =				inputs.attention_mask
 #    input_values = inputs.input_values[:, :-1]
 #    attention_mask = inputs.attention_mask[:, :-1]

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

  lowerCamelCase		:     Any          =				torch.argmax(lowerCamelCase,     dim=-1							)
  lowerCamelCase		:     List[str]          =				processor.batch_decode(lowerCamelCase							)

  print(F'''Expected Output: {ds[:4]["text"]}, Pred: {output_string}'''							)
 else:
  lowerCamelCase		:     Union[str, Any]          =				model(source=lowerCamelCase,     padding_mask=(1 - attention_mask),     mask=lowerCamelCase,     features_only=lowerCamelCase							)[
      """layer_results"""
  ][-1][0].transpose(0,     1							)
  lowerCamelCase		:     Optional[Any]          =				hf_wavavec(lowerCamelCase,     attention_mask=lowerCamelCase							)["""last_hidden_state"""]

 print(our_output.shape,     their_output.shape							)
 lowerCamelCase		:     str          =				torch.max(torch.abs(our_output - their_output							)							).item()
 print(F'''max_absolute_diff = {max_absolute_diff}'''							)  # ~ 1e-7
 lowerCamelCase		:     int          =				torch.allclose(lowerCamelCase,     lowerCamelCase,     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(lowerCamelCase							)

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


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_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        =parser.parse_args()
      convert_wavaveca_checkpoint(
          args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
      )




