




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

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






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Union[str, Any]  ,	__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		List[Any]  ,	__lowerCAmelCase			:		Union[str, Any]  ,	__lowerCAmelCase			:		Tuple						)    ->       Union[str, Any]:
 for attribute in key.split("."						):
  _UpperCamelCase      :	Optional[Any]					      =  getattr(__lowerCAmelCase  ,	__lowerCAmelCase						)

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

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






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		str  ,	__lowerCAmelCase			:		Dict  ,	__lowerCAmelCase			:		Dict						)    ->       Tuple:
 _UpperCamelCase      :	int					      =  []
 _UpperCamelCase      :	int					      =  fairseq_model.state_dict()

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

 else:
  _UpperCamelCase      :	str					      =  hf_model.feature_extractor
  _UpperCamelCase      :	Optional[int]					      =  hf_model.encoder.pos_conv_embed

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






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		List[str]  ,	__lowerCAmelCase			:		str  ,	__lowerCAmelCase			:		List[str]  ,	__lowerCAmelCase			:		Dict						)    ->       Tuple:
 _UpperCamelCase      :	str					      =  access_by_string(__lowerCAmelCase  ,	__lowerCAmelCase						)
 _UpperCamelCase      :	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."						)
 _UpperCamelCase      :	List[str]					      =  fsq_value
 logger.info(f"{full_name} was correctly initialized from {hf_weight_path}."						)






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		List[Any]  ,	__lowerCAmelCase			:		Optional[int]  ,	__lowerCAmelCase			:		Tuple  ,	__lowerCAmelCase			:		str						)    ->       List[Any]:
 _UpperCamelCase      :	Tuple					      =  full_name.split("conv_layers."						)[-1]
 _UpperCamelCase      :	str					      =  name.split("."						)
 _UpperCamelCase      :	Union[str, Any]					      =  int(items[0]						)
 _UpperCamelCase      :	List[Any]					      =  int(items[1]						)

 _UpperCamelCase      :	Tuple					      =  name.split("."						)[-1]
 if type_id == 0:
  _UpperCamelCase      :	str					      =  "conv"
 elif type_id == 2:
  _UpperCamelCase      :	Optional[Any]					      =  "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			:		Dict  ,	__lowerCAmelCase			:		List[str]  ,	__lowerCAmelCase			:		Tuple  ,	__lowerCAmelCase			:		Dict						)    ->       Union[str, Any]:
 _UpperCamelCase      :	int					      =  full_name.split("pos_conv."						)[-1]
 _UpperCamelCase      :	Union[str, Any]					      =  name.split("."						)
 _UpperCamelCase      :	Optional[Any]					      =  int(items[0]						)
 _UpperCamelCase      :	List[str]					      =  int(items[1]						)

 _UpperCamelCase      :	Optional[Any]					      =  name.split("."						)[-1]
 if type_id != 0:
  unused_weights.append(__lowerCAmelCase						)
  return
 else:
  _UpperCamelCase      :	str					      =  "conv"

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






@torch.no_grad()
def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Dict  ,	__lowerCAmelCase			:		List[Any]  ,	__lowerCAmelCase			:		Optional[int]=None  ,	__lowerCAmelCase			:		Optional[Any]=None  ,	__lowerCAmelCase			:		List[str]=True						)    ->       Any:
 if config_path is not None:
  _UpperCamelCase      :	Union[str, Any]					      =  DataaVecAudioConfig.from_pretrained(__lowerCAmelCase						)
 else:
  _UpperCamelCase      :	Optional[int]					      =  DataaVecAudioConfig()

 if not is_finetuned:
  # Modify final_proj layer name
  _UpperCamelCase      :	Any					      =  DataaVecAudioModel(__lowerCAmelCase						)
  _UpperCamelCase      :	Any					      =  os.path.dirname(__lowerCAmelCase						)

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

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

 _UpperCamelCase      :	Optional[int]					      =  load_dataavec(__lowerCAmelCase						)

 recursively_load_weights(__lowerCAmelCase  ,	__lowerCAmelCase  ,	not is_finetuned						)

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

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

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

 _UpperCamelCase      :	List[Any]					      =  inputs.input_values
 _UpperCamelCase      :	Union[str, Any]					      =  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[Any]					      =  model(source=__lowerCAmelCase  ,	padding_mask=(1 - attention_mask)  ,	mask=__lowerCAmelCase  ,	features_only=__lowerCAmelCase						)[
      "encoder_out"
  ].transpose(0  ,	1						)
  _UpperCamelCase      :	str					      =  hf_wavavec(__lowerCAmelCase  ,	attention_mask=__lowerCAmelCase						)["logits"]

  _UpperCamelCase      :	Dict					      =  torch.argmax(__lowerCAmelCase  ,	dim=-1						)
  _UpperCamelCase      :	List[Any]					      =  processor.batch_decode(__lowerCAmelCase						)

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

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









"""simple docstring"""
import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available


if is_torch_available():
    import numpy as np
    import torch

if is_vision_available():
    from PIL import Image






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Any  ,	__lowerCAmelCase			:		Dict=False  ,	__lowerCAmelCase			:		Optional[Any]=False  ,	__lowerCAmelCase			:		int=False						)    ->       List[str]:
 assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

 _UpperCamelCase      :	Optional[int]					      =  []
 for i in range(image_processor_tester.batch_size						):
  if equal_resolution:
   _UpperCamelCase      :	List[str]					      =  image_processor_tester.max_resolution
  else:
   # To avoid getting image width/height 0
   _UpperCamelCase      :	int					      =  image_processor_tester.min_resolution
   if getattr(__lowerCAmelCase  ,	"size_divisor"  ,	__lowerCAmelCase						):
    # If `size_divisor` is defined, the image needs to have width/size >= `size_divisor`
    _UpperCamelCase      :	Optional[int]					      =  max(image_processor_tester.size_divisor  ,	__lowerCAmelCase						)
   _UpperCamelCase      ,	_UpperCamelCase      :	Optional[Any]					      =  np.random.choice(np.arange(__lowerCAmelCase  ,	image_processor_tester.max_resolution						)  ,	2						)
  image_inputs.append(
      np.random.randint(255  ,	size=(image_processor_tester.num_channels, width, height)  ,	dtype=np.uinta						)						)

 if not numpify and not torchify:
  # PIL expects the channel dimension as last dimension
  _UpperCamelCase      :	Optional[Any]					      =  [Image.fromarray(np.moveaxis(__lowerCAmelCase  ,	0  ,	-1						)						) for image in image_inputs]

 if torchify:
  _UpperCamelCase      :	str					      =  [torch.from_numpy(__lowerCAmelCase						) for image in image_inputs]

 return image_inputs






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		Tuple=10  ,	__lowerCAmelCase			:		str=10  ,	__lowerCAmelCase			:		Any=False  ,	__lowerCAmelCase			:		List[Any]=False						)    ->       Optional[int]:
 _UpperCamelCase      :	Optional[int]					      =  []
 for i in range(image_processor_tester.num_frames						):
  video.append(np.random.randint(255  ,	size=(image_processor_tester.num_channels, width, height)  ,	dtype=np.uinta						)						)

 if not numpify and not torchify:
  # PIL expects the channel dimension as last dimension
  _UpperCamelCase      :	Optional[int]					      =  [Image.fromarray(np.moveaxis(__lowerCAmelCase  ,	0  ,	-1						)						) for frame in video]

 if torchify:
  _UpperCamelCase      :	int					      =  [torch.from_numpy(__lowerCAmelCase						) for frame in video]

 return video






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		List[Any]  ,	__lowerCAmelCase			:		Union[str, Any]=False  ,	__lowerCAmelCase			:		Any=False  ,	__lowerCAmelCase			:		Dict=False						)    ->       Any:
 assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

 _UpperCamelCase      :	Optional[Any]					      =  []
 for i in range(image_processor_tester.batch_size						):
  if equal_resolution:
   _UpperCamelCase      :	List[Any]					      =  image_processor_tester.max_resolution
  else:
   _UpperCamelCase      ,	_UpperCamelCase      :	Optional[int]					      =  np.random.choice(
       np.arange(image_processor_tester.min_resolution  ,	image_processor_tester.max_resolution						)  ,	2						)
   _UpperCamelCase      :	List[Any]					      =  prepare_video(
       image_processor_tester=__lowerCAmelCase  ,	width=__lowerCAmelCase  ,	height=__lowerCAmelCase  ,	numpify=__lowerCAmelCase  ,	torchify=__lowerCAmelCase  ,	)
  video_inputs.append(__lowerCAmelCase						)

 return video_inputs







class  _SCREAMING_SNAKE_CASE       :



 '''simple docstring'''





 __UpperCAmelCase		   =							None






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	str					      =  self.image_processing_class(**self.image_processor_dict   )
  _UpperCamelCase      :	Optional[int]					      =  json.loads(image_processor.to_json_string()   )
  for key, value in self.image_processor_dict.items():
   self.assertEqual(obj[key]					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Union[str, Any]					      =  self.image_processing_class(**self.image_processor_dict   )

  with tempfile.TemporaryDirectory() as tmpdirname:
   _UpperCamelCase      :	Tuple					      =  os.path.join(lowerCAmelCase__					,     "image_processor.json"   )
   image_processor_first.to_json_file(lowerCAmelCase__   )
   _UpperCamelCase      :	Tuple					      =  self.image_processing_class.from_json_file(lowerCAmelCase__   )

  self.assertEqual(image_processor_second.to_dict()					,     image_processor_first.to_dict()   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  self.image_processing_class(**self.image_processor_dict   )

  with tempfile.TemporaryDirectory() as tmpdirname:
   _UpperCamelCase      :	Optional[int]					      =  image_processor_first.save_pretrained(lowerCAmelCase__   )[0]
   check_json_file_has_correct_format(lowerCAmelCase__   )
   _UpperCamelCase      :	int					      =  self.image_processing_class.from_pretrained(lowerCAmelCase__   )

  self.assertEqual(image_processor_second.to_dict()					,     image_processor_first.to_dict()   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Tuple					      =  self.image_processing_class()
  self.assertIsNotNone(lowerCAmelCase__   )






 @require_torch
 @require_vision
 def lowercase_      (self   ):

  '''simple docstring'''







  if self.test_cast_dtype is not None:
   # Initialize image_processor
   _UpperCamelCase      :	Optional[int]					      =  self.image_processing_class(**self.image_processor_dict   )

   # create random PyTorch tensors
   _UpperCamelCase      :	str					      =  prepare_image_inputs(self.image_processor_tester					,     equal_resolution=lowerCAmelCase__					,     torchify=lowerCAmelCase__   )

   _UpperCamelCase      :	List[Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   )
   # for layoutLM compatiblity
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )

   _UpperCamelCase      :	List[Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to(torch.floataa   )
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )

   _UpperCamelCase      :	str					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to("cpu"					,     torch.bfloataa   )
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.bfloataa   )

   with self.assertRaises(lowerCAmelCase__   ):
    _UpperCamelCase      :	Optional[int]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to(torch.bfloataa					,     "cpu"   )

   # Try with text + image feature
   _UpperCamelCase      :	Union[str, Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   )
   encoding.update({"input_ids": torch.LongTensor([[1, 2, 3], [4, 5, 6]]   )}   )
   _UpperCamelCase      :	Optional[int]					      =  encoding.to(torch.floataa   )

   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )
   self.assertEqual(encoding.input_ids.dtype					,     torch.long   )



