




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

_lowerCAmelCase    : Optional[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    : Optional[Any]		= [
    'lm_head',
]



def     lowerCAmelCase		(		_lowerCAmelCase				: List[Any]    ,				_lowerCAmelCase				: List[str]    ,				_lowerCAmelCase				: Optional[int]    ,				_lowerCAmelCase				: Union[str, Any]    ,				_lowerCAmelCase				: Tuple						):




 """simple docstring"""







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

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

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



def     lowerCAmelCase		(		_lowerCAmelCase				: List[Any]    ,				_lowerCAmelCase				: Union[str, Any]    ,				_lowerCAmelCase				: Union[str, Any]						):




 """simple docstring"""







 UpperCAmelCase__							=    []
 UpperCAmelCase__							=    fairseq_model.state_dict()

 if not is_headless:
  UpperCAmelCase__							=    hf_model.dataavec_audio.feature_extractor
  UpperCAmelCase__							=    hf_model.dataavec_audio.encoder.pos_conv_embed

 else:
  UpperCAmelCase__							=    hf_model.feature_extractor
  UpperCAmelCase__							=    hf_model.encoder.pos_conv_embed

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




 """simple docstring"""







 UpperCAmelCase__							=    path.split("."						)
 return reduce(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase						)



def     lowerCAmelCase		(		_lowerCAmelCase				: Any    ,				_lowerCAmelCase				: Tuple    ,				_lowerCAmelCase				: List[Any]    ,				_lowerCAmelCase				: Any						):




 """simple docstring"""







 UpperCAmelCase__							=    access_by_string(_lowerCAmelCase    ,				_lowerCAmelCase						)
 UpperCAmelCase__							=    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__							=    fsq_value
 logger.info(F'''{full_name} was correctly initialized from {hf_weight_path}.'''						)



def     lowerCAmelCase		(		_lowerCAmelCase				: Any    ,				_lowerCAmelCase				: str    ,				_lowerCAmelCase				: Any    ,				_lowerCAmelCase				: List[str]						):




 """simple docstring"""







 UpperCAmelCase__							=    full_name.split("conv_layers."						)[-1]
 UpperCAmelCase__							=    name.split("."						)
 UpperCAmelCase__							=    int(items[0]						)
 UpperCAmelCase__							=    int(items[1]						)

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

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



def     lowerCAmelCase		(		_lowerCAmelCase				: int    ,				_lowerCAmelCase				: Dict    ,				_lowerCAmelCase				: List[Any]    ,				_lowerCAmelCase				: Union[str, Any]						):




 """simple docstring"""







 UpperCAmelCase__							=    full_name.split("pos_conv."						)[-1]
 UpperCAmelCase__							=    name.split("."						)
 UpperCAmelCase__							=    int(items[0]						)
 UpperCAmelCase__							=    int(items[1]						)

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

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



@torch.no_grad()
def     lowerCAmelCase		(		_lowerCAmelCase				: Dict    ,				_lowerCAmelCase				: Union[str, Any]    ,				_lowerCAmelCase				: Dict=None    ,				_lowerCAmelCase				: Optional[int]=None    ,				_lowerCAmelCase				: Optional[int]=True						):




 """simple docstring"""







 if config_path is not None:
  UpperCAmelCase__							=    DataaVecAudioConfig.from_pretrained(_lowerCAmelCase						)
 else:
  UpperCAmelCase__							=    DataaVecAudioConfig()

 if not is_finetuned:
  # Modify final_proj layer name
  UpperCAmelCase__							=    DataaVecAudioModel(_lowerCAmelCase						)
  UpperCAmelCase__							=    os.path.dirname(_lowerCAmelCase						)

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

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

 UpperCAmelCase__							=    load_dataavec(_lowerCAmelCase						)

 recursively_load_weights(_lowerCAmelCase    ,				_lowerCAmelCase    ,				not is_finetuned						)

 UpperCAmelCase__							=    WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"						)

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

 UpperCAmelCase__							=    processor(_lowerCAmelCase    ,				return_tensors="pt"    ,				padding=_lowerCAmelCase						)

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

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

  UpperCAmelCase__							=    torch.argmax(_lowerCAmelCase    ,				dim=-1						)
  UpperCAmelCase__							=    processor.batch_decode(_lowerCAmelCase						)

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

 print(our_output.shape    ,				their_output.shape						)
 UpperCAmelCase__							=    torch.max(torch.abs(our_output - their_output						)						).item()
 print(F'''max_absolute_diff = {max_absolute_diff}'''						)  # ~ 1e-7
 UpperCAmelCase__							=    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    : Union[str, 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    : 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 ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCAmelCase					:						Any         =     logging.get_logger(__name__)

_UpperCAmelCase					:						Dict         =     {
    'ZinengTang/tvlt-base': 'https://huggingface.co/ZinengTang/tvlt-base/blob/main/config.json',
}


class 		lowerCAmelCase_				(  UpperCamelCase ):
     __lowerCamelCase       :    Any 	=     '''tvlt'''





     def __init__(	self							, _UpperCAmelCase=224							, _UpperCAmelCase=2048							, _UpperCAmelCase=128							, _UpperCAmelCase=[16, 16]							, _UpperCAmelCase=[16, 16]							, _UpperCAmelCase=3							, _UpperCAmelCase=1							, _UpperCAmelCase=8							, _UpperCAmelCase=768							, _UpperCAmelCase=12							, _UpperCAmelCase=12							, _UpperCAmelCase=3072							, _UpperCAmelCase="gelu"							, _UpperCAmelCase=0.0							, _UpperCAmelCase=0.0							, _UpperCAmelCase=0.02							, _UpperCAmelCase=1E-6							, _UpperCAmelCase=True							, _UpperCAmelCase=False							, _UpperCAmelCase=16							, _UpperCAmelCase=512							, _UpperCAmelCase=8							, _UpperCAmelCase=2048							, _UpperCAmelCase=0.75							, _UpperCAmelCase=0.15							, _UpperCAmelCase="frame-level"							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase="classification"							, **_UpperCAmelCase							, )		->				List[Any]:
        super().__init__(**_UpperCAmelCase  )

        if audio_mask_type not in ("frame-level", "patch_level"):
           raise ValueError(
               """audio_mask_type must be one of two acceptable strategies - {'frame_level', 'patch-level') """
               F'''got {audio_mask_type}'''  )

        lowerCamelCase									=image_size
        lowerCamelCase									=spectrogram_length
        lowerCamelCase									=frequency_length
        lowerCamelCase									=image_patch_size
        lowerCamelCase									=audio_patch_size
        lowerCamelCase									=num_image_channels
        lowerCamelCase									=num_audio_channels
        lowerCamelCase									=num_frames

        lowerCamelCase									=hidden_size
        lowerCamelCase									=num_hidden_layers
        lowerCamelCase									=num_attention_heads
        lowerCamelCase									=intermediate_size
        lowerCamelCase									=hidden_act
        lowerCamelCase									=hidden_dropout_prob
        lowerCamelCase									=attention_probs_dropout_prob
        lowerCamelCase									=initializer_range
        lowerCamelCase									=layer_norm_eps
        lowerCamelCase									=qkv_bias
        lowerCamelCase									=use_mean_pooling

        lowerCamelCase									=decoder_num_attention_heads
        lowerCamelCase									=decoder_hidden_size
        lowerCamelCase									=decoder_num_hidden_layers
        lowerCamelCase									=decoder_intermediate_size
        lowerCamelCase									=pixel_mask_ratio
        lowerCamelCase									=audio_mask_ratio
        lowerCamelCase									=audio_mask_type

        lowerCamelCase									=task_matching
        lowerCamelCase									=task_mae
        lowerCamelCase									=loss_type