






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

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







def 		_UpperCamelCase      (			UpperCamelCase_				:    List[Any]			,				UpperCamelCase_				:    Optional[int]			,				UpperCamelCase_				:    Tuple			,				UpperCamelCase_				:    List[str]			,				UpperCamelCase_				:    Any						)     ->      List[Any]:

						"""simple docstring"""



						for attribute in key.split('.'						):
												lowerCAmelCase__			  =  getattr(UpperCamelCase_			,				UpperCamelCase_						)

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

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







def 		_UpperCamelCase      (			UpperCamelCase_				:    Dict			,				UpperCamelCase_				:    Optional[int]			,				UpperCamelCase_				:    Dict						)     ->      Dict:

						"""simple docstring"""



						lowerCAmelCase__			  =  []
						lowerCAmelCase__			  =  fairseq_model.state_dict()

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

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

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

						"""simple docstring"""



						lowerCAmelCase__			  =  path.split('.'						)
						return reduce(UpperCamelCase_			,				UpperCamelCase_			,				UpperCamelCase_						)







def 		_UpperCamelCase      (			UpperCamelCase_				:    str			,				UpperCamelCase_				:    List[Any]			,				UpperCamelCase_				:    Dict			,				UpperCamelCase_				:    Optional[Any]						)     ->      List[Any]:

						"""simple docstring"""



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







def 		_UpperCamelCase      (			UpperCamelCase_				:    Union[str, Any]			,				UpperCamelCase_				:    Union[str, Any]			,				UpperCamelCase_				:    int			,				UpperCamelCase_				:    Union[str, Any]						)     ->      str:

						"""simple docstring"""



						lowerCAmelCase__			  =  full_name.split('conv_layers.'						)[-1]
						lowerCAmelCase__			  =  name.split('.'						)
						lowerCAmelCase__			  =  int(items[0]						)
						lowerCAmelCase__			  =  int(items[1]						)

						lowerCAmelCase__			  =  name.split('.'						)[-1]
						if type_id == 0:
												lowerCAmelCase__			  =  'conv'
						elif type_id == 2:
												lowerCAmelCase__			  =  '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_				:    Any			,				UpperCamelCase_				:    Union[str, Any]			,				UpperCamelCase_				:    List[Any]			,				UpperCamelCase_				:    int						)     ->      Dict:

						"""simple docstring"""



						lowerCAmelCase__			  =  full_name.split('pos_conv.'						)[-1]
						lowerCAmelCase__			  =  name.split('.'						)
						lowerCAmelCase__			  =  int(items[0]						)
						lowerCAmelCase__			  =  int(items[1]						)

						lowerCAmelCase__			  =  name.split('.'						)[-1]
						if type_id != 0:
												unused_weights.append(UpperCamelCase_						)
												return
						else:
												lowerCAmelCase__			  =  'conv'

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







@torch.no_grad()
def 		_UpperCamelCase      (			UpperCamelCase_				:    Tuple			,				UpperCamelCase_				:    str			,				UpperCamelCase_				:    Tuple=None			,				UpperCamelCase_				:    Tuple=None			,				UpperCamelCase_				:    Optional[int]=True						)     ->      Optional[int]:

						"""simple docstring"""



						if config_path is not None:
												lowerCAmelCase__			  =  DataaVecAudioConfig.from_pretrained(UpperCamelCase_						)
						else:
												lowerCAmelCase__			  =  DataaVecAudioConfig()

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

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

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

						lowerCAmelCase__			  =  load_dataavec(UpperCamelCase_						)

						recursively_load_weights(UpperCamelCase_			,				UpperCamelCase_			,				not is_finetuned						)

						lowerCAmelCase__			  =  WavaVecaProcessor.from_pretrained('facebook/wav2vec2-large-lv60'						)

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

						lowerCAmelCase__			  =  processor(UpperCamelCase_			,				return_tensors='pt'			,				padding=UpperCamelCase_						)

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

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

												lowerCAmelCase__			  =  torch.argmax(UpperCamelCase_			,				dim=-1						)
												lowerCAmelCase__			  =  processor.batch_decode(UpperCamelCase_						)

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

						print(our_output.shape			,				their_output.shape						)
						lowerCAmelCase__			  =  torch.max(torch.abs(our_output - their_output						)						).item()
						print(F"max_absolute_diff = {max_absolute_diff}"						)  # ~ 1e-7
						lowerCAmelCase__			  =  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__":
					__snake_case			:			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"""
					)
					__snake_case			:			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
					)













import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json







def 		_UpperCamelCase      (			UpperCamelCase_				:    Tuple						)     ->      str:

						"""simple docstring"""



						lowerCAmelCase__			  =  filter(lambda UpperCamelCase_						: p.requires_grad			,				model.parameters()						)
						lowerCAmelCase__			  =  sum([np.prod(p.size()						) for p in model_parameters]						)
						return params


__snake_case			:			Dict			   =   logging.getLogger(__name__)


class 			__SCREAMING_SNAKE_CASE    (      pl.Callback):






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase      ):






												"""simple docstring"""
												lowerCAmelCase__			  =  {F"lr_group_{i}": param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups      )}
												pl_module.logger.log_metrics(_UpperCamelCase      )






						@rank_zero_only
						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase=True      ):






												"""simple docstring"""
												logger.info(F"***** {type_path} results at step {trainer.global_step:05d} *****"      )
												lowerCAmelCase__			  =  trainer.callback_metrics
												trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']}      )
												# Log results
												lowerCAmelCase__			  =  Path(pl_module.hparams.output_dir      )
												if type_path == "test":
																		lowerCAmelCase__			  =  od / 'test_results.txt'
																		lowerCAmelCase__			  =  od / 'test_generations.txt'
												else:
																		# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
																		# If people want this it will be easy enough to add back.
																		lowerCAmelCase__			  =  od / F"{type_path}_results/{trainer.global_step:05d}.txt"
																		lowerCAmelCase__			  =  od / F"{type_path}_generations/{trainer.global_step:05d}.txt"
																		results_file.parent.mkdir(exist_ok=_UpperCamelCase      )
																		generations_file.parent.mkdir(exist_ok=_UpperCamelCase      )
												with open(_UpperCamelCase   ,  'a+'      ) as writer:
																		for key in sorted(_UpperCamelCase      ):
																								if key in ["log", "progress_bar", "preds"]:
																														continue
																								lowerCAmelCase__			  =  metrics[key]
																								if isinstance(_UpperCamelCase   ,  torch.Tensor      ):
																														lowerCAmelCase__			  =  val.item()
																								lowerCAmelCase__			  =  F"{key}: {val:.6f}\n"
																								writer.write(_UpperCamelCase      )

												if not save_generations:
																		return

												if "preds" in metrics:
																		lowerCAmelCase__			  =  '\n'.join(metrics['preds']      )
																		generations_file.open('w+'      ).write(_UpperCamelCase      )






						@rank_zero_only
						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase      ):






												"""simple docstring"""
												try:
																		lowerCAmelCase__			  =  pl_module.model.model.num_parameters()
												except AttributeError:
																		lowerCAmelCase__			  =  pl_module.model.num_parameters()

												lowerCAmelCase__			  =  count_trainable_parameters(_UpperCamelCase      )
												# mp stands for million parameters
												trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1E6, 'grad_mp': n_trainable_pars / 1E6}      )






						@rank_zero_only
						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase      ):






												"""simple docstring"""
												save_json(pl_module.metrics   ,  pl_module.metrics_save_path      )
												return self._write_logs(_UpperCamelCase   ,  _UpperCamelCase   ,  'test'      )







						@rank_zero_only
						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase      ):






												"""simple docstring"""
												save_json(pl_module.metrics   ,  pl_module.metrics_save_path      )
												# Uncommenting this will save val generations
												# return self._write_logs(trainer, pl_module, "valid")







def 		_UpperCamelCase      (			UpperCamelCase_				:    Union[str, Any]			,				UpperCamelCase_				:    Union[str, Any]			,				UpperCamelCase_				:    List[Any]=1			,				UpperCamelCase_				:    List[str]=False						)     ->      Optional[Any]:

						"""simple docstring"""



						if metric == "rouge2":
												lowerCAmelCase__			  =  '{val_avg_rouge2:.4f}-{step_count}'
						elif metric == "bleu":
												lowerCAmelCase__			  =  '{val_avg_bleu:.4f}-{step_count}'
						elif metric == "loss":
												lowerCAmelCase__			  =  '{val_avg_loss:.4f}-{step_count}'
						else:
												raise NotImplementedError(
												    F"seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"
												    ' this function.'						)

						lowerCAmelCase__			  =  ModelCheckpoint(
						    dirpath=UpperCamelCase_			,				filename=UpperCamelCase_			,				monitor=F"val_{metric}"			,				mode='min' if 'loss' in metric else 'max'			,				save_top_k=UpperCamelCase_			,				)
						return checkpoint_callback







def 		_UpperCamelCase      (			UpperCamelCase_				:    Union[str, Any]			,				UpperCamelCase_				:    int						)     ->      Dict:

						"""simple docstring"""



						return EarlyStopping(
						    monitor=F"val_{metric}"			,				mode='min' if 'loss' in metric else 'max'			,				patience=UpperCamelCase_			,				verbose=UpperCamelCase_			,				)





