




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

_lowerCamelCase							:						Tuple							     =  {
    '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							:						str							     =  [
    'lm_head',
]






def 		_lowercase		(     lowercase_						,							lowercase_						,							lowercase_						,							lowercase_						,							lowercase_  ):



		'''simple docstring'''







		for attribute in key.split("."  ):
				_snake_case     	=	getattr(lowercase_						,							lowercase_  )

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

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






def 		_lowercase		(     lowercase_						,							lowercase_						,							lowercase_  ):



		'''simple docstring'''







		_snake_case     	=	[]
		_snake_case     	=	fairseq_model.state_dict()

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

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

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

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






def 		_lowercase		(     lowercase_						,							lowercase_  ):



		'''simple docstring'''







		_snake_case     	=	path.split("."  )
		return reduce(lowercase_						,							lowercase_						,							lowercase_  )






def 		_lowercase		(     lowercase_						,							lowercase_						,							lowercase_						,							lowercase_  ):



		'''simple docstring'''







		_snake_case     	=	access_by_string(lowercase_						,							lowercase_  )
		_snake_case     	=	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     	=	fsq_value
		logger.info(f'''{full_name} was correctly initialized from {hf_weight_path}.'''  )






def 		_lowercase		(     lowercase_						,							lowercase_						,							lowercase_						,							lowercase_  ):



		'''simple docstring'''







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

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

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






def 		_lowercase		(     lowercase_						,							lowercase_						,							lowercase_						,							lowercase_  ):



		'''simple docstring'''







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

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

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






@torch.no_grad()
def 		_lowercase		(     lowercase_						,							lowercase_						,							lowercase_=None						,							lowercase_=None						,							lowercase_=True  ):



		'''simple docstring'''







		if config_path is not None:
				_snake_case     	=	DataaVecAudioConfig.from_pretrained(lowercase_  )
		else:
				_snake_case     	=	DataaVecAudioConfig()

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

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

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

		_snake_case     	=	load_dataavec(lowercase_  )

		recursively_load_weights(lowercase_						,							lowercase_						,							not is_finetuned  )

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

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

		_snake_case     	=	processor(lowercase_						,							return_tensors="pt"						,							padding=lowercase_  )

		_snake_case     	=	inputs.input_values
		_snake_case     	=	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     	=	model(source=lowercase_						,							padding_mask=(1 - attention_mask)						,							mask=lowercase_						,							features_only=lowercase_  )[
				    "encoder_out"
				].transpose(0						,							1  )
				_snake_case     	=	hf_wavavec(lowercase_						,							attention_mask=lowercase_  )["logits"]

				_snake_case     	=	torch.argmax(lowercase_						,							dim=-1  )
				_snake_case     	=	processor.batch_decode(lowercase_  )

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

		print(our_output.shape						,							their_output.shape  )
		_snake_case     	=	torch.max(torch.abs(our_output - their_output  )  ).item()
		print(f'''max_absolute_diff = {max_absolute_diff}'''  )  # ~ 1e-7
		_snake_case     	=	torch.allclose(lowercase_						,							lowercase_						,							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(lowercase_  )

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


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							:						Tuple							     =  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 unittest

import numpy as np

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

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
						import torch

if is_vision_available():
						from PIL import Image

						from transformers import CLIPImageProcessor





class 					lowerCamelCase_							(	unittest.TestCase							):




					'''simple docstring'''






					def __init__(   self,  __UpperCamelCase,  __UpperCamelCase=7,  __UpperCamelCase=3,  __UpperCamelCase=18,  __UpperCamelCase=30,  __UpperCamelCase=400,  __UpperCamelCase=True,  __UpperCamelCase=None,  __UpperCamelCase=True,  __UpperCamelCase=None,  __UpperCamelCase=True,  __UpperCamelCase=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73],  __UpperCamelCase=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11],  __UpperCamelCase=True,  ):
							_snake_case     	=	size if size is not None else {"shortest_edge": 20}
							_snake_case     	=	crop_size if crop_size is not None else {"height": 18, "width": 18}
							_snake_case     	=	parent
							_snake_case     	=	batch_size
							_snake_case     	=	num_channels
							_snake_case     	=	image_size
							_snake_case     	=	min_resolution
							_snake_case     	=	max_resolution
							_snake_case     	=	do_resize
							_snake_case     	=	size
							_snake_case     	=	do_center_crop
							_snake_case     	=	crop_size
							_snake_case     	=	do_normalize
							_snake_case     	=	image_mean
							_snake_case     	=	image_std
							_snake_case     	=	do_convert_rgb






					def     __SCREAMING_SNAKE_CASE     (   self			):
							return {
							    "do_resize": self.do_resize,
							    "size": self.size,
							    "do_center_crop": self.do_center_crop,
							    "crop_size": self.crop_size,
							    "do_normalize": self.do_normalize,
							    "image_mean": self.image_mean,
							    "image_std": self.image_std,
							    "do_convert_rgb": self.do_convert_rgb,
							}






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase=False,  __UpperCamelCase=False,  __UpperCamelCase=False			):
							assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

							if equal_resolution:
									_snake_case     	=	[]
									for i in range(self.batch_size			):
											image_inputs.append(
											    np.random.randint(
											        255,  size=(self.num_channels, self.max_resolution, self.max_resolution),  dtype=np.uinta			)			)
							else:
									_snake_case     	=	[]
									for i in range(self.batch_size			):
											_snake_case	,   _snake_case     	=	np.random.choice(np.arange(self.min_resolution,  self.max_resolution			),  2			)
											image_inputs.append(np.random.randint(255,  size=(self.num_channels, width, height),  dtype=np.uinta			)			)

							if not numpify and not torchify:
									# PIL expects the channel dimension as last dimension
									_snake_case     	=	[Image.fromarray(np.moveaxis(__UpperCamelCase,  0,  -1			)			) for x in image_inputs]

							if torchify:
									_snake_case     	=	[torch.from_numpy(__UpperCamelCase			) for x in image_inputs]

							return image_inputs





@require_torch
@require_vision
class 					lowerCamelCase_							(	UpperCAmelCase_,	unittest.TestCase							):




					'''simple docstring'''






					lowercase__   =			CLIPImageProcessor if is_vision_available() else None






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	CLIPImageProcessingTester(self			)






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							return self.image_processor_tester.prepare_image_processor_dict()






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.image_processing_class(**self.image_processor_dict			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_resize"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "size"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_center_crop"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "center_crop"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_normalize"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "image_mean"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "image_std"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_convert_rgb"			)			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.image_processing_class.from_dict(self.image_processor_dict			)
							self.assertEqual(image_processor.size,  {"shortest_edge": 20}			)
							self.assertEqual(image_processor.crop_size,  {"height": 18, "width": 18}			)

							_snake_case     	=	self.image_processing_class.from_dict(self.image_processor_dict,  size=42,  crop_size=84			)
							self.assertEqual(image_processor.size,  {"shortest_edge": 42}			)
							self.assertEqual(image_processor.crop_size,  {"height": 84, "width": 84}			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							pass






					def     __SCREAMING_SNAKE_CASE     (   self			):
							# Initialize image_processing
							_snake_case     	=	self.image_processing_class(**self.image_processor_dict			)
							# create random PIL images
							_snake_case     	=	self.image_processor_tester.prepare_inputs(equal_resolution=__UpperCamelCase			)
							for image in image_inputs:
									self.assertIsInstance(__UpperCamelCase,  Image.Image			)

							# Test not batched input
							_snake_case     	=	image_processing(image_inputs[0],  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        1,
							        self.image_processor_tester.num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )

							# Test batched
							_snake_case     	=	image_processing(__UpperCamelCase,  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )






					def     __SCREAMING_SNAKE_CASE     (   self			):
							# Initialize image_processing
							_snake_case     	=	self.image_processing_class(**self.image_processor_dict			)
							# create random numpy tensors
							_snake_case     	=	self.image_processor_tester.prepare_inputs(equal_resolution=__UpperCamelCase,  numpify=__UpperCamelCase			)
							for image in image_inputs:
									self.assertIsInstance(__UpperCamelCase,  np.ndarray			)

							# Test not batched input
							_snake_case     	=	image_processing(image_inputs[0],  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        1,
							        self.image_processor_tester.num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )

							# Test batched
							_snake_case     	=	image_processing(__UpperCamelCase,  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )






					def     __SCREAMING_SNAKE_CASE     (   self			):
							# Initialize image_processing
							_snake_case     	=	self.image_processing_class(**self.image_processor_dict			)
							# create random PyTorch tensors
							_snake_case     	=	self.image_processor_tester.prepare_inputs(equal_resolution=__UpperCamelCase,  torchify=__UpperCamelCase			)
							for image in image_inputs:
									self.assertIsInstance(__UpperCamelCase,  torch.Tensor			)

							# Test not batched input
							_snake_case     	=	image_processing(image_inputs[0],  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        1,
							        self.image_processor_tester.num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )

							# Test batched
							_snake_case     	=	image_processing(__UpperCamelCase,  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        self.image_processor_tester.batch_size,
							        self.image_processor_tester.num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )





@require_torch
@require_vision
class 					lowerCamelCase_							(	UpperCAmelCase_,	unittest.TestCase							):




					'''simple docstring'''






					lowercase__   =			CLIPImageProcessor if is_vision_available() else None






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	CLIPImageProcessingTester(self,  num_channels=4			)
							_snake_case     	=	3






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							return self.image_processor_tester.prepare_image_processor_dict()






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.image_processing_class(**self.image_processor_dict			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_resize"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "size"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_center_crop"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "center_crop"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_normalize"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "image_mean"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "image_std"			)			)
							self.assertTrue(hasattr(__UpperCamelCase,  "do_convert_rgb"			)			)






					def     __SCREAMING_SNAKE_CASE     (   self			):
							pass






					def     __SCREAMING_SNAKE_CASE     (   self			):
							# Initialize image_processing
							_snake_case     	=	self.image_processing_class(**self.image_processor_dict			)
							# create random PIL images
							_snake_case     	=	self.image_processor_tester.prepare_inputs(equal_resolution=__UpperCamelCase			)
							for image in image_inputs:
									self.assertIsInstance(__UpperCamelCase,  Image.Image			)

							# Test not batched input
							_snake_case     	=	image_processing(image_inputs[0],  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        1,
							        self.expected_encoded_image_num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )

							# Test batched
							_snake_case     	=	image_processing(__UpperCamelCase,  return_tensors="pt"			).pixel_values
							self.assertEqual(
							    encoded_images.shape,  (
							        self.image_processor_tester.batch_size,
							        self.expected_encoded_image_num_channels,
							        self.image_processor_tester.crop_size["height"],
							        self.image_processor_tester.crop_size["width"],
							    ),  )





