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

__lowerCamelCase     :   int							= {
    """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     :   Tuple							= [
    """lm_head""",
]


def     UpperCamelCase      (					__lowerCamelCase  :				Optional[Any]       ,      __lowerCamelCase  :				Optional[Any]       ,      __lowerCamelCase  :				Tuple       ,      __lowerCamelCase  :				int       ,      __lowerCamelCase  :				int					):
       for attribute in key.split("."					):
              snake_case :     Tuple           =   getattr(__lowerCamelCase       ,      __lowerCamelCase					)

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

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


def     UpperCamelCase      (					__lowerCamelCase  :				int       ,      __lowerCamelCase  :				List[Any]       ,      __lowerCamelCase  :				Any					):
       snake_case :     Union[str, Any]           =   []
       snake_case :     Optional[Any]           =   fairseq_model.state_dict()

       if not is_headless:
              snake_case :     Optional[Any]           =   hf_model.dataavec_audio.feature_extractor
              snake_case :     int           =   hf_model.dataavec_audio.encoder.pos_conv_embed

       else:
              snake_case :     List[str]           =   hf_model.feature_extractor
              snake_case :     int           =   hf_model.encoder.pos_conv_embed

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


def     UpperCamelCase      (					__lowerCamelCase  :				Tuple       ,      __lowerCamelCase  :				Optional[Any]       ,      __lowerCamelCase  :				List[Any]       ,      __lowerCamelCase  :				Optional[int]					):
       snake_case :     Any           =   access_by_string(__lowerCamelCase       ,      __lowerCamelCase					)
       snake_case :     Union[str, 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."""					)
       snake_case :     Tuple           =   fsq_value
       logger.info(f"""{full_name} was correctly initialized from {hf_weight_path}."""					)


def     UpperCamelCase      (					__lowerCamelCase  :				Union[str, Any]       ,      __lowerCamelCase  :				Any       ,      __lowerCamelCase  :				str       ,      __lowerCamelCase  :				str					):
       snake_case :     Optional[int]           =   full_name.split("conv_layers."					)[-1]
       snake_case :     List[str]           =   name.split("."					)
       snake_case :     Dict           =   int(items[0]					)
       snake_case :     str           =   int(items[1]					)

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

       set_weights(__lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase       ,      f"""conv_layers.{layer_id}.{layer_type}.{weight_type}"""					)


def     UpperCamelCase      (					__lowerCamelCase  :				Optional[Any]       ,      __lowerCamelCase  :				List[str]       ,      __lowerCamelCase  :				Any       ,      __lowerCamelCase  :				str					):
       snake_case :     Tuple           =   full_name.split("pos_conv."					)[-1]
       snake_case :     str           =   name.split("."					)
       snake_case :     Any           =   int(items[0]					)
       snake_case :     int           =   int(items[1]					)

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

       set_weights(__lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase       ,      f"""layers.{layer_id}.{layer_type}.{weight_type}"""					)


@torch.no_grad()
def     UpperCamelCase      (					__lowerCamelCase  :				Union[str, Any]       ,      __lowerCamelCase  :				Any       ,      __lowerCamelCase  :				Dict=None       ,      __lowerCamelCase  :				Any=None       ,      __lowerCamelCase  :				str=True					):
       if config_path is not None:
              snake_case :     Optional[int]           =   DataaVecAudioConfig.from_pretrained(__lowerCamelCase					)
       else:
              snake_case :     Optional[int]           =   DataaVecAudioConfig()

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

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

       def load_dataavec(__lowerCamelCase  :				Union[str, Any]					):
              snake_case :     Any           =   fairseq.checkpoint_utils.load_model_ensemble_and_task([path]					)
              return model[0].eval()

       snake_case :     Any           =   load_dataavec(__lowerCamelCase					)

       recursively_load_weights(__lowerCamelCase       ,      __lowerCamelCase       ,      not is_finetuned					)

       snake_case :     Tuple           =   WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"					)

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

       snake_case :     Dict           =   processor(__lowerCamelCase       ,      return_tensors="pt"       ,      padding=__lowerCamelCase					)

       snake_case :     Optional[Any]           =   inputs.input_values
       snake_case :     Dict           =   inputs.attention_mask
       #    input_values = inputs.input_values[:, :-1]
       #    attention_mask = inputs.attention_mask[:, :-1]

       hf_wavavec.eval()
       model.eval()
       if is_finetuned:
              snake_case :     Union[str, Any]           =   model(source=__lowerCamelCase       ,      padding_mask=(1 - attention_mask)       ,      mask=__lowerCamelCase       ,      features_only=__lowerCamelCase					)[
                  "encoder_out"
              ].transpose(0       ,      1					)
              snake_case :     Dict           =   hf_wavavec(__lowerCamelCase       ,      attention_mask=__lowerCamelCase					)["logits"]

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

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

       print(our_output.shape       ,      their_output.shape					)
       snake_case :     Tuple           =   torch.max(torch.abs(our_output - their_output					)					).item()
       print(f"""max_absolute_diff = {max_absolute_diff}"""					)  # ~ 1e-7
       snake_case :     Dict           =   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     :   Optional[int]							= 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     :   Optional[int]							= parser.parse_args()
   convert_wavaveca_checkpoint(
       args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
   )







def     UpperCamelCase      (					__lowerCamelCase  :				int = 998001					):
       snake_case :     Optional[Any]           =   0
       for i in range(999       ,      99       ,      -1					):  # 3 digit numbers range from 999 down to 100
              for j in range(999       ,      99       ,      -1					):
                     snake_case :     Tuple           =   str(i * j					)
                     if product_string == product_string[::-1] and i * j < n:
                            snake_case :     str           =   max(__lowerCamelCase       ,      i * j					)
       return answer


if __name__ == "__main__":
   print(F'{solution() = }')





