



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

__UpperCAmelCase						  =     {
    '''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''',
}
__UpperCAmelCase						  =     [
    '''lm_head''',
]





def 						__lowerCamelCase       (      __magic_name__	:  Union[str, Any]     , __magic_name__	:  Union[str, Any]     , __magic_name__	:  List[Any]     , __magic_name__	:  str     , __magic_name__	:  Dict						):
							for attribute in key.split("."						):
														a__:       int											=getattr(__magic_name__     , __magic_name__						)

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

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





def 						__lowerCamelCase       (      __magic_name__	:  List[str]     , __magic_name__	:  Tuple     , __magic_name__	:  Optional[Any]						):
							a__:       Optional[int]											=[]
							a__:       Any											=fairseq_model.state_dict()

							if not is_headless:
														a__:       Any											=hf_model.dataavec_audio.feature_extractor
														a__:       Tuple											=hf_model.dataavec_audio.encoder.pos_conv_embed

							else:
														a__:       int											=hf_model.feature_extractor
														a__:       Any											=hf_model.encoder.pos_conv_embed

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

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





def 						__lowerCamelCase       (      __magic_name__	:  str     , __magic_name__	:  str						):
							a__:       Optional[int]											=path.split("."						)
							return reduce(__magic_name__     , __magic_name__     , __magic_name__						)





def 						__lowerCamelCase       (      __magic_name__	:  int     , __magic_name__	:  str     , __magic_name__	:  Dict     , __magic_name__	:  int						):
							a__:       Optional[int]											=access_by_string(__magic_name__     , __magic_name__						)
							a__:       str											=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."						)
							a__:       Optional[Any]											=fsq_value
							logger.info(F"{full_name} was correctly initialized from {hf_weight_path}."						)





def 						__lowerCamelCase       (      __magic_name__	:  List[str]     , __magic_name__	:  Optional[int]     , __magic_name__	:  Optional[Any]     , __magic_name__	:  Optional[Any]						):
							a__:       Optional[Any]											=full_name.split("conv_layers."						)[-1]
							a__:       List[str]											=name.split("."						)
							a__:       Tuple											=int(items[0]						)
							a__:       Optional[Any]											=int(items[1]						)

							a__:       Optional[Any]											=name.split("."						)[-1]
							if type_id == 0:
														a__:       Union[str, Any]											="conv"
							elif type_id == 2:
														a__:       Optional[Any]											="layer_norm"
							else:
														unused_weights.append(__magic_name__						)
														return

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





def 						__lowerCamelCase       (      __magic_name__	:  Optional[Any]     , __magic_name__	:  str     , __magic_name__	:  List[Any]     , __magic_name__	:  str						):
							a__:       str											=full_name.split("pos_conv."						)[-1]
							a__:       Union[str, Any]											=name.split("."						)
							a__:       Dict											=int(items[0]						)
							a__:       Any											=int(items[1]						)

							a__:       List[str]											=name.split("."						)[-1]
							if type_id != 0:
														unused_weights.append(__magic_name__						)
														return
							else:
														a__:       List[Any]											="conv"

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





@torch.no_grad()
def 						__lowerCamelCase       (      __magic_name__	:  Any     , __magic_name__	:  Tuple     , __magic_name__	:  int=None     , __magic_name__	:  Dict=None     , __magic_name__	:  Any=True						):
							if config_path is not None:
														a__:       Any											=DataaVecAudioConfig.from_pretrained(__magic_name__						)
							else:
														a__:       Optional[int]											=DataaVecAudioConfig()

							if not is_finetuned:
														# Modify final_proj layer name
														a__:       Union[str, Any]											=DataaVecAudioModel(__magic_name__						)
														a__:       Union[str, Any]											=os.path.dirname(__magic_name__						)

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

							def load_dataavec(__magic_name__	:  List[str]						):
														a__							,			a__							,			a__:       str											=fairseq.checkpoint_utils.load_model_ensemble_and_task([path]						)
														return model[0].eval()

							a__:       Dict											=load_dataavec(__magic_name__						)

							recursively_load_weights(__magic_name__     , __magic_name__     , not is_finetuned						)

							a__:       List[Any]											=WavaVecaProcessor.from_pretrained("facebook/wav2vec2-large-lv60"						)

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

							a__:       Any											=processor(__magic_name__     , return_tensors="pt"     , padding=__magic_name__						)

							a__:       int											=inputs.input_values
							a__:       Optional[int]											=inputs.attention_mask
							#    input_values = inputs.input_values[:, :-1]
							#    attention_mask = inputs.attention_mask[:, :-1]

							hf_wavavec.eval()
							model.eval()
							if is_finetuned:
														a__:       Dict											=model(source=__magic_name__     , padding_mask=(1 - attention_mask)     , mask=__magic_name__     , features_only=__magic_name__						)[
														    "encoder_out"
														].transpose(0     , 1						)
														a__:       Optional[Any]											=hf_wavavec(__magic_name__     , attention_mask=__magic_name__						)["logits"]

														a__:       Union[str, Any]											=torch.argmax(__magic_name__     , dim=-1						)
														a__:       Tuple											=processor.batch_decode(__magic_name__						)

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

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

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


if __name__ == "__main__":
		__UpperCAmelCase						  =     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'''
		)
		__UpperCAmelCase						  =     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 json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


__UpperCAmelCase						  =     10_00_00_00_00_00
__UpperCAmelCase						  =     1_00_00

__UpperCAmelCase  ,							__UpperCAmelCase						  =     os.path.split(__file__)
__UpperCAmelCase						  =     os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))





def 						__lowerCamelCase       (      __magic_name__	:  int     , __magic_name__	:  int						):
							a__:       Any											=pa.Table.from_pydict({"col": [0] * chunk_size}						)
							a__:       Optional[Any]											=pa.concat_tables([table] * (num_examples // chunk_size)						)
							return datasets.Dataset(__magic_name__     , fingerprint="table_100B"						)


@dataclass
class 	lowerCamelCase__ :
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42


						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Tuple											=np.random.default_rng(self.seed			)
													a__:       Tuple											=rng.integers(low=self.low      ,       high=self.high      ,       size=self.size			).tolist()


						def __iter__( self						:  Any			):
													return iter(self._sampled_values			)







						def __len__( self						:  Any			):
													return self.size





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Union[str, Any]											=dataset[0]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Optional[Any]											=dataset[-1]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       List[str]											=dataset[range(len(__magic_name__						) // 2     , len(__magic_name__						) // 2 + 1_024						)]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Optional[int]											=dataset[RandIter(0     , len(__magic_name__						)     , 1_024     , seed=42						)]





def 						__lowerCamelCase       (      ):
							a__:       Dict											={"num examples": SPEED_TEST_N_EXAMPLES}
							a__:       Union[str, Any]											=(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
							print("generating dataset"						)
							a__:       Dict											=generate_100B_dataset(num_examples=__magic_name__     , chunk_size=__magic_name__						)
							print("Functions"						)
							for func in functions:
														print(func.__name__						)
														a__:       Dict											=func(__magic_name__						)

							with open(__magic_name__     , "wb"						) as f:
														f.write(json.dumps(__magic_name__						).encode("utf-8"						)						)


if __name__ == "__main__":  # useful to run the profiler
		benchmark_table_100B()





