




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

_lowerCamelCase					:     List[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					:     List[str]					  =  [
    "lm_head",
]




def 							_UpperCAmelCase     (UpperCamelCase_       :       Tuple					,			UpperCamelCase_       :       Dict					,			UpperCamelCase_       :       Optional[int]					,			UpperCamelCase_       :       str					,			UpperCamelCase_       :       Optional[int]				):
				'''simple docstring'''
				for attribute in key.split("""."""				):
								_lowerCAmelCase			:			Dict        =					getattr(UpperCamelCase_					,			UpperCamelCase_				)

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

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




def 							_UpperCAmelCase     (UpperCamelCase_       :       Tuple					,			UpperCamelCase_       :       Any					,			UpperCamelCase_       :       List[str]				):
				'''simple docstring'''
				_lowerCAmelCase			:			str        =					[]
				_lowerCAmelCase			:			Union[str, Any]        =					fairseq_model.state_dict()

				if not is_headless:
								_lowerCAmelCase			:			Tuple        =					hf_model.dataavec_audio.feature_extractor
								_lowerCAmelCase			:			List[Any]        =					hf_model.dataavec_audio.encoder.pos_conv_embed

				else:
								_lowerCAmelCase			:			List[Any]        =					hf_model.feature_extractor
								_lowerCAmelCase			:			Optional[Any]        =					hf_model.encoder.pos_conv_embed

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

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




def 							_UpperCAmelCase     (UpperCamelCase_       :       Any					,			UpperCamelCase_       :       Dict				):
				'''simple docstring'''
				_lowerCAmelCase			:			Optional[Any]        =					path.split("""."""				)
				return reduce(UpperCamelCase_					,			UpperCamelCase_					,			UpperCamelCase_				)




def 							_UpperCAmelCase     (UpperCamelCase_       :       Union[str, Any]					,			UpperCamelCase_       :       Optional[Any]					,			UpperCamelCase_       :       List[Any]					,			UpperCamelCase_       :       Dict				):
				'''simple docstring'''
				_lowerCAmelCase			:			Union[str, Any]        =					access_by_string(UpperCamelCase_					,			UpperCamelCase_				)
				_lowerCAmelCase			:			int        =					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			:			Any        =					fsq_value
				logger.info(F"{full_name} was correctly initialized from {hf_weight_path}."				)




def 							_UpperCAmelCase     (UpperCamelCase_       :       Any					,			UpperCamelCase_       :       Dict					,			UpperCamelCase_       :       List[Any]					,			UpperCamelCase_       :       Optional[Any]				):
				'''simple docstring'''
				_lowerCAmelCase			:			Dict        =					full_name.split("""conv_layers."""				)[-1]
				_lowerCAmelCase			:			Tuple        =					name.split("""."""				)
				_lowerCAmelCase			:			List[str]        =					int(items[0]				)
				_lowerCAmelCase			:			Union[str, Any]        =					int(items[1]				)

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

				set_weights(UpperCamelCase_					,			UpperCamelCase_					,			UpperCamelCase_					,			F"conv_layers.{layer_id}.{layer_type}.{weight_type}"				)




def 							_UpperCAmelCase     (UpperCamelCase_       :       Dict					,			UpperCamelCase_       :       int					,			UpperCamelCase_       :       List[Any]					,			UpperCamelCase_       :       str				):
				'''simple docstring'''
				_lowerCAmelCase			:			Tuple        =					full_name.split("""pos_conv."""				)[-1]
				_lowerCAmelCase			:			List[str]        =					name.split("""."""				)
				_lowerCAmelCase			:			Tuple        =					int(items[0]				)
				_lowerCAmelCase			:			Dict        =					int(items[1]				)

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

				set_weights(UpperCamelCase_					,			UpperCamelCase_					,			UpperCamelCase_					,			F"layers.{layer_id}.{layer_type}.{weight_type}"				)




@torch.no_grad()
def 							_UpperCAmelCase     (UpperCamelCase_       :       Union[str, Any]					,			UpperCamelCase_       :       Optional[int]					,			UpperCamelCase_       :       Union[str, Any]=None					,			UpperCamelCase_       :       Union[str, Any]=None					,			UpperCamelCase_       :       Union[str, Any]=True				):
				'''simple docstring'''
				if config_path is not None:
								_lowerCAmelCase			:			List[str]        =					DataaVecAudioConfig.from_pretrained(UpperCamelCase_				)
				else:
								_lowerCAmelCase			:			Any        =					DataaVecAudioConfig()

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

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

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

				_lowerCAmelCase			:			Union[str, Any]        =					load_dataavec(UpperCamelCase_				)

				recursively_load_weights(UpperCamelCase_					,			UpperCamelCase_					,			not is_finetuned				)

				_lowerCAmelCase			:			Union[str, Any]        =					WavaVecaProcessor.from_pretrained("""facebook/wav2vec2-large-lv60"""				)

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

				_lowerCAmelCase			:			Union[str, Any]        =					processor(UpperCamelCase_					,			return_tensors="""pt"""					,			padding=UpperCamelCase_				)

				_lowerCAmelCase			:			Optional[int]        =					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			:			Union[str, Any]        =					model(source=UpperCamelCase_					,			padding_mask=(1 - attention_mask)					,			mask=UpperCamelCase_					,			features_only=UpperCamelCase_				)[
								    """encoder_out"""
								].transpose(0					,			1				)
								_lowerCAmelCase			:			Optional[int]        =					hf_wavavec(UpperCamelCase_					,			attention_mask=UpperCamelCase_				)["""logits"""]

								_lowerCAmelCase			:			Any        =					torch.argmax(UpperCamelCase_					,			dim=-1				)
								_lowerCAmelCase			:			int        =					processor.batch_decode(UpperCamelCase_				)

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

				print(our_output.shape					,			their_output.shape				)
				_lowerCAmelCase			:			List[str]        =					torch.max(torch.abs(our_output - their_output				)				).item()
				print(F"max_absolute_diff = {max_absolute_diff}"				)  # ~ 1e-7
				_lowerCAmelCase			:			Optional[int]        =					torch.allclose(UpperCamelCase_					,			UpperCamelCase_					,			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(UpperCamelCase_				)

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


if __name__ == "__main__":
	_lowerCamelCase					:     Tuple					  =  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
	)







import numpy as np  # noqa: E402

from ...configuration_utils import ConfigMixin, register_to_config
from ...schedulers.scheduling_utils import SchedulerMixin


try:
	import librosa  # noqa: E402

	_lowerCamelCase					:     Union[str, Any]					  =  True
	_lowerCamelCase					:     Optional[Any]					  =  ""
except Exception as e:
	_lowerCamelCase					:     int					  =  False
	_lowerCamelCase					:     List[str]					  =  (
	    F'''Cannot import librosa because {e}. Make sure to correctly install librosa to be able to install it.'''
	)


from PIL import Image  # noqa: E402
class 				__snake_case     (_a  ,		_a						):
				lowerCAmelCase__			=   "mel_config.json"



				@register_to_config
				def __init__(					self					:						Dict						,		_UpperCAmelCase					:						int = 256						,		_UpperCAmelCase					:						int = 256						,		_UpperCAmelCase					:						int = 2_2050						,		_UpperCAmelCase					:						int = 2048						,		_UpperCAmelCase					:						int = 512						,		_UpperCAmelCase					:						int = 80						,		_UpperCAmelCase					:						int = 32						,		)     ->			List[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			List[str]        =					hop_length
								_lowerCAmelCase			:			Optional[Any]        =					sample_rate
								_lowerCAmelCase			:			Optional[Any]        =					n_fft
								_lowerCAmelCase			:			Dict        =					top_db
								_lowerCAmelCase			:			List[Any]        =					n_iter
								self.set_resolution(_UpperCAmelCase						,		_UpperCAmelCase					)
								_lowerCAmelCase			:			Optional[int]        =					None

								if not _librosa_can_be_imported:
												raise ValueError(_import_error					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[int]						,		_UpperCAmelCase					:						int						,		_UpperCAmelCase					:						int					)     ->			List[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			Dict        =					x_res
								_lowerCAmelCase			:			Union[str, Any]        =					y_res
								_lowerCAmelCase			:			Any        =					self.y_res
								_lowerCAmelCase			:			Optional[Any]        =					self.x_res * self.hop_length - 1



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						str = None						,		_UpperCAmelCase					:						np.ndarray = None					)     ->			List[str]:





								'''simple docstring'''




								if audio_file is not None:
												_lowerCAmelCase    ,			_lowerCAmelCase			:			Dict        =					librosa.load(_UpperCAmelCase						,		mono=_UpperCAmelCase						,		sr=self.sr					)
								else:
												_lowerCAmelCase			:			Any        =					raw_audio

								# Pad with silence if necessary.
								if len(self.audio					) < self.x_res * self.hop_length:
												_lowerCAmelCase			:			Optional[Any]        =					np.concatenate([self.audio, np.zeros((self.x_res * self.hop_length - len(self.audio					),)					)]					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Any					)     ->			int:





								'''simple docstring'''




								return len(self.audio					) // self.slice_size



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[str]						,		_UpperCAmelCase					:						int = 0					)     ->			np.ndarray:





								'''simple docstring'''




								return self.audio[self.slice_size * slice : self.slice_size * (slice + 1)]



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[Any]					)     ->			int:





								'''simple docstring'''




								return self.sr



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						int					)     ->			Image.Image:





								'''simple docstring'''




								_lowerCAmelCase			:			Union[str, Any]        =					librosa.feature.melspectrogram(
								    y=self.get_audio_slice(_UpperCAmelCase					)						,		sr=self.sr						,		n_fft=self.n_fft						,		hop_length=self.hop_length						,		n_mels=self.n_mels					)
								_lowerCAmelCase			:			List[Any]        =					librosa.power_to_db(_UpperCAmelCase						,		ref=np.max						,		top_db=self.top_db					)
								_lowerCAmelCase			:			List[Any]        =					(((log_S + self.top_db) * 255 / self.top_db).clip(0						,		255					) + 0.5).astype(np.uinta					)
								_lowerCAmelCase			:			Optional[Any]        =					Image.fromarray(_UpperCAmelCase					)
								return image



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[int]						,		_UpperCAmelCase					:						Image.Image					)     ->			np.ndarray:





								'''simple docstring'''




								_lowerCAmelCase			:			Optional[int]        =					np.frombuffer(image.tobytes()						,		dtype="""uint8"""					).reshape((image.height, image.width)					)
								_lowerCAmelCase			:			Optional[int]        =					bytedata.astype("""float"""					) * self.top_db / 255 - self.top_db
								_lowerCAmelCase			:			List[Any]        =					librosa.db_to_power(_UpperCAmelCase					)
								_lowerCAmelCase			:			Optional[int]        =					librosa.feature.inverse.mel_to_audio(
								    _UpperCAmelCase						,		sr=self.sr						,		n_fft=self.n_fft						,		hop_length=self.hop_length						,		n_iter=self.n_iter					)
								return audio

