


'''simple docstring'''
import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConformerConfig,
    WavaVecaConformerForCTC,
    WavaVecaConformerForPreTraining,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
__snake_case											=					logging.get_logger(__name__)

__snake_case											=					{
    '''post_extract_proj''': '''feature_projection.projection''',
    '''encoder.pos_conv.0''': '''encoder.pos_conv_embed.conv''',
    '''self_attn.linear_k''': '''encoder.layers.*.self_attn.linear_k''',
    '''self_attn.linear_v''': '''encoder.layers.*.self_attn.linear_v''',
    '''self_attn.linear_q''': '''encoder.layers.*.self_attn.linear_q''',
    '''self_attn.pos_bias_u''': '''encoder.layers.*.self_attn.pos_bias_u''',
    '''self_attn.pos_bias_v''': '''encoder.layers.*.self_attn.pos_bias_v''',
    '''self_attn.linear_out''': '''encoder.layers.*.self_attn.linear_out''',
    '''self_attn.linear_pos''': '''encoder.layers.*.self_attn.linear_pos''',
    '''self_attn.rotary_emb''': '''encoder.embed_positions''',
    '''self_attn_layer_norm''': '''encoder.layers.*.self_attn_layer_norm''',
    '''conv_module.pointwise_conv1''': '''encoder.layers.*.conv_module.pointwise_conv1''',
    '''conv_module.pointwise_conv2''': '''encoder.layers.*.conv_module.pointwise_conv2''',
    '''conv_module.depthwise_conv''': '''encoder.layers.*.conv_module.depthwise_conv''',
    '''conv_module.batch_norm''': '''encoder.layers.*.conv_module.batch_norm''',
    '''conv_module.layer_norm''': '''encoder.layers.*.conv_module.layer_norm''',
    '''ffn1.w_1''': '''encoder.layers.*.ffn1.intermediate_dense''',
    '''ffn1.w_2''': '''encoder.layers.*.ffn1.output_dense''',
    '''ffn1.layer_norm''': '''encoder.layers.*.ffn1_layer_norm''',
    '''ffn2.w_1''': '''encoder.layers.*.ffn2.intermediate_dense''',
    '''ffn2.w_2''': '''encoder.layers.*.ffn2.output_dense''',
    '''ffn2.layer_norm''': '''encoder.layers.*.ffn2_layer_norm''',
    '''final_layer_norm''': '''encoder.layers.*.final_layer_norm''',
    '''encoder.layer_norm''': '''encoder.layer_norm''',
    '''w2v_model.layer_norm''': '''feature_projection.layer_norm''',
    '''quantizer.weight_proj''': '''quantizer.weight_proj''',
    '''quantizer.vars''': '''quantizer.codevectors''',
    '''project_q''': '''project_q''',
    '''final_proj''': '''project_hid''',
    '''w2v_encoder.proj''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}
__snake_case											=					[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]


def       a				(      __a			, __a			, __a			, __a			, __a       )  ->     Optional[int]:



						'''simple docstring'''






						for attribute in key.split('''.'''       ):
												UpperCamelCase__       :str							 =			getattr(__a			, __a       )

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

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


def       a				(      __a			, __a			, __a       )  ->     Optional[int]:



						'''simple docstring'''






						UpperCamelCase__       :Dict							 =			[]
						UpperCamelCase__       :Any							 =			fairseq_model.state_dict()

						UpperCamelCase__       :Dict							 =			hf_model.wavaveca_conformer.feature_extractor

						for name, value in fairseq_dict.items():
												UpperCamelCase__       :Tuple							 =			False
												if "conv_layers" in name:
																		load_conv_layer(
																		    __a			, __a			, __a			, __a			, hf_model.config.feat_extract_norm == '''group'''			, )
																		UpperCamelCase__       :List[str]							 =			True
												else:
																		for key, mapped_key in MAPPING.items():
																								UpperCamelCase__       :Any							 =			'''wav2vec2_conformer.''' + 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__       :List[Any]							 =			True
																														if "*" in mapped_key:
																																				UpperCamelCase__       :List[str]							 =			name.split(__a       )[0].split('''.'''       )[-2]
																																				UpperCamelCase__       :Union[str, Any]							 =			mapped_key.replace('''*'''			, __a       )
																														if "pos_bias_u" in name:
																																				UpperCamelCase__       :int							 =			None
																														elif "pos_bias_v" in name:
																																				UpperCamelCase__       :Union[str, Any]							 =			None
																														elif "weight_g" in name:
																																				UpperCamelCase__       :Dict							 =			'''weight_g'''
																														elif "weight_v" in name:
																																				UpperCamelCase__       :int							 =			'''weight_v'''
																														elif "bias" in name:
																																				UpperCamelCase__       :Union[str, Any]							 =			'''bias'''
																														elif "weight" in name:
																																				# TODO: don't match quantizer.weight_proj
																																				UpperCamelCase__       :Union[str, Any]							 =			'''weight'''
																														elif "running_mean" in name:
																																				UpperCamelCase__       :Union[str, Any]							 =			'''running_mean'''
																														elif "inv_freq" in name:
																																				UpperCamelCase__       :Optional[Any]							 =			'''inv_freq'''
																														elif "running_var" in name:
																																				UpperCamelCase__       :Tuple							 =			'''running_var'''
																														elif "num_batches_tracked" in name:
																																				UpperCamelCase__       :Dict							 =			'''num_batches_tracked'''
																														else:
																																				UpperCamelCase__       :str							 =			None
																														set_recursively(__a			, __a			, __a			, __a			, __a       )
																								continue
												if not is_used:
																		unused_weights.append(__a       )

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


def       a				(      __a			, __a			, __a			, __a			, __a       )  ->     str:



						'''simple docstring'''






						UpperCamelCase__       :Union[str, Any]							 =			full_name.split('''conv_layers.'''       )[-1]
						UpperCamelCase__       :Dict							 =			name.split('''.'''       )
						UpperCamelCase__       :Tuple							 =			int(items[0]       )
						UpperCamelCase__       :int							 =			int(items[1]       )

						if type_id == 0:
												if "bias" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''       )
																		UpperCamelCase__       :str							 =			value
																		logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''       )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''       )
																		UpperCamelCase__       :List[str]							 =			value
																		logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''       )
						elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
												if "bias" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'''       )
																		UpperCamelCase__       :int							 =			value
																		logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''       )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
																								raise ValueError(
																								    f'''{full_name} has size {value.shape}, but'''
																								    f''' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'''       )
																		UpperCamelCase__       :List[str]							 =			value
																		logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''       )
						else:
												unused_weights.append(__a       )


@torch.no_grad()
def       a				(      __a			, __a			, __a=None			, __a=None			, __a=True       )  ->     int:



						'''simple docstring'''






						if config_path is not None:
												UpperCamelCase__       :Dict							 =			WavaVecaConformerConfig.from_pretrained(__a			, hidden_act='''swish'''       )
						else:
												UpperCamelCase__       :int							 =			WavaVecaConformerConfig()

						if "rope" in checkpoint_path:
												UpperCamelCase__       :str							 =			'''rotary'''

						if is_finetuned:
												if dict_path:
																		UpperCamelCase__       :Dict							 =			Dictionary.load(__a       )

																		# important change bos & pad token id since CTC symbol is <pad> and
																		# not <s> as in fairseq
																		UpperCamelCase__       :List[Any]							 =			target_dict.pad_index
																		UpperCamelCase__       :Any							 =			target_dict.bos_index
																		UpperCamelCase__       :Dict							 =			target_dict.eos_index
																		UpperCamelCase__       :List[Any]							 =			len(target_dict.symbols       )
																		UpperCamelCase__       :List[str]							 =			os.path.join(__a			, '''vocab.json'''       )
																		if not os.path.isdir(__a       ):
																								logger.error('''--pytorch_dump_folder_path ({}) should be a directory'''.format(__a       )       )
																								return
																		os.makedirs(__a			, exist_ok=__a       )
																		UpperCamelCase__       :List[str]							 =			target_dict.indices

																		# fairseq has the <pad> and <s> switched
																		UpperCamelCase__       :List[Any]							 =			0
																		UpperCamelCase__       :List[Any]							 =			1
																		with open(__a			, '''w'''			, encoding='''utf-8'''       ) as vocab_handle:
																								json.dump(__a			, __a       )
																		UpperCamelCase__       :Union[str, Any]							 =			WavaVecaCTCTokenizer(
																		    __a			, unk_token=target_dict.unk_word			, pad_token=target_dict.pad_word			, bos_token=target_dict.bos_word			, eos_token=target_dict.eos_word			, word_delimiter_token='''|'''			, do_lower_case=__a			, )
																		UpperCamelCase__       :int							 =			True if config.feat_extract_norm == '''layer''' else False
																		UpperCamelCase__       :int							 =			WavaVecaFeatureExtractor(
																		    feature_size=1			, sampling_rate=16000			, padding_value=0			, do_normalize=__a			, return_attention_mask=__a			, )
																		UpperCamelCase__       :Optional[int]							 =			WavaVecaProcessor(feature_extractor=__a			, tokenizer=__a       )
																		processor.save_pretrained(__a       )

												UpperCamelCase__       :List[Any]							 =			WavaVecaConformerForCTC(__a       )
						else:
												UpperCamelCase__       :List[str]							 =			WavaVecaConformerForPreTraining(__a       )

						if is_finetuned:
												UpperCamelCase__       :int							 =			fairseq.checkpoint_utils.load_model_ensemble_and_task(
												    [checkpoint_path]			, arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''       )[:-1]       )}       )
						else:
												UpperCamelCase__       :int							 =			argparse.Namespace(task='''audio_pretraining'''       )
												UpperCamelCase__       :Any							 =			fairseq.tasks.setup_task(__a       )

												UpperCamelCase__       :List[Any]							 =			fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]			, task=__a       )

						UpperCamelCase__       :List[str]							 =			model[0].eval()

						recursively_load_weights(__a			, __a			, not is_finetuned       )

						hf_wavavec.save_pretrained(__a       )


if __name__ == "__main__":
		__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'''
		)
		__snake_case											=					parser.parse_args()
		convert_wavaveca_conformer_checkpoint(
		    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
		)



'''simple docstring'''
# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
		import jax
		import jaxlib

__snake_case											=					get_logger()

__snake_case											=					None
class        lowercase     (							TensorFormatter[Mapping, 'jax.Array', Mapping] ):

		"""simple docstring"""

		def __init__(		self   ,       UpperCamelCase_=None   ,       UpperCamelCase_=None   ,       **UpperCamelCase_							):

							'''simple docstring'''


							super().__init__(features=UpperCamelCase_							)
							import jax
							from jaxlib.xla_client import Device

							if isinstance(UpperCamelCase_   ,       UpperCamelCase_							):
												raise ValueError(
												    F'''Expected {device} to be a `str` not {type(UpperCamelCase_							)}, as `jaxlib.xla_extension.Device` '''
												    '''is not serializable neither with `pickle` nor with `dill`. Instead you can surround '''
												    '''the device with `str()` to get its string identifier that will be internally mapped '''
												    '''to the actual `jaxlib.xla_extension.Device`.'''							)
							UpperCamelCase__       :Tuple							 =			device if isinstance(UpperCamelCase_   ,       UpperCamelCase_							) else str(jax.devices()[0]							)
							# using global variable since `jaxlib.xla_extension.Device` is not serializable neither
							# with `pickle` nor with `dill`, so we need to use a global variable instead
							global DEVICE_MAPPING
							if DEVICE_MAPPING is None:
												UpperCamelCase__       :Optional[Any]							 =			self._map_devices_to_str()
							if self.device not in list(DEVICE_MAPPING.keys()							):
												logger.warning(
												    F'''Device with string identifier {self.device} not listed among the available '''
												    F'''devices: {list(DEVICE_MAPPING.keys()							)}, so falling back to the default '''
												    F'''device: {str(jax.devices()[0]							)}.'''							)
												UpperCamelCase__       :Optional[int]							 =			str(jax.devices()[0]							)
							UpperCamelCase__       :Tuple							 =			jnp_array_kwargs




		@staticmethod
		def    lowerCAmelCase__    (		):

							'''simple docstring'''


							import jax

							return {str(UpperCamelCase_							): device for device in jax.devices()}




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							import jax
							import jax.numpy as jnp

							if isinstance(UpperCamelCase_   ,       UpperCamelCase_							) and column:
												if all(
												    isinstance(UpperCamelCase_   ,       jax.Array							) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column							):
																	return jnp.stack(UpperCamelCase_   ,       axis=0							)
							return column




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							import jax
							import jax.numpy as jnp

							if isinstance(UpperCamelCase_   ,       (str, bytes, type(UpperCamelCase_							))							):
												return value
							elif isinstance(UpperCamelCase_   ,       (np.character, np.ndarray)							) and np.issubdtype(value.dtype   ,       np.character							):
												return value.tolist()

							UpperCamelCase__       :Optional[int]							 =			{}

							if isinstance(UpperCamelCase_   ,       (np.number, np.ndarray)							) and np.issubdtype(value.dtype   ,       np.integer							):
												# the default int precision depends on the jax config
												# see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
												if jax.config.jax_enable_xaa:
																	UpperCamelCase__       :List[str]							 =			{'''dtype''': jnp.intaa}
												else:
																	UpperCamelCase__       :Union[str, Any]							 =			{'''dtype''': jnp.intaa}
							elif isinstance(UpperCamelCase_   ,       (np.number, np.ndarray)							) and np.issubdtype(value.dtype   ,       np.floating							):
												UpperCamelCase__       :Optional[Any]							 =			{'''dtype''': jnp.floataa}
							elif config.PIL_AVAILABLE and "PIL" in sys.modules:
												import PIL.Image

												if isinstance(UpperCamelCase_   ,       PIL.Image.Image							):
																	UpperCamelCase__       :str							 =			np.asarray(UpperCamelCase_							)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
							global DEVICE_MAPPING
							if DEVICE_MAPPING is None:
												UpperCamelCase__       :Dict							 =			self._map_devices_to_str()

							with jax.default_device(DEVICE_MAPPING[self.device]							):
												# calling jnp.array on a np.ndarray does copy the data
												# see https://github.com/google/jax/issues/4486
												return jnp.array(UpperCamelCase_   ,       **{**default_dtype, **self.jnp_array_kwargs}							)




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							import jax

							# support for torch, tf, jax etc.
							if config.TORCH_AVAILABLE and "torch" in sys.modules:
												import torch

												if isinstance(UpperCamelCase_   ,       torch.Tensor							):
																	return self._tensorize(data_struct.detach().cpu().numpy()[()]							)
							if hasattr(UpperCamelCase_   ,       '''__array__'''							) and not isinstance(UpperCamelCase_   ,       jax.Array							):
												UpperCamelCase__       :int							 =			data_struct.__array__()
							# support for nested types like struct of list of struct
							if isinstance(UpperCamelCase_   ,       np.ndarray							):
												if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
																	return self._consolidate([self.recursive_tensorize(UpperCamelCase_							) for substruct in data_struct]							)
							elif isinstance(UpperCamelCase_   ,       (list, tuple)							):
												return self._consolidate([self.recursive_tensorize(UpperCamelCase_							) for substruct in data_struct]							)
							return self._tensorize(UpperCamelCase_							)




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							return map_nested(self._recursive_tensorize   ,       UpperCamelCase_   ,       map_list=UpperCamelCase_							)




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			self.numpy_arrow_extractor().extract_row(UpperCamelCase_							)
							UpperCamelCase__       :List[Any]							 =			self.python_features_decoder.decode_row(UpperCamelCase_							)
							return self.recursive_tensorize(UpperCamelCase_							)




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			self.numpy_arrow_extractor().extract_column(UpperCamelCase_							)
							UpperCamelCase__       :Tuple							 =			self.python_features_decoder.decode_column(UpperCamelCase_   ,       pa_table.column_names[0]							)
							UpperCamelCase__       :Dict							 =			self.recursive_tensorize(UpperCamelCase_							)
							UpperCamelCase__       :str							 =			self._consolidate(UpperCamelCase_							)
							return column




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :int							 =			self.numpy_arrow_extractor().extract_batch(UpperCamelCase_							)
							UpperCamelCase__       :Union[str, Any]							 =			self.python_features_decoder.decode_batch(UpperCamelCase_							)
							UpperCamelCase__       :List[str]							 =			self.recursive_tensorize(UpperCamelCase_							)
							for column_name in batch:
												UpperCamelCase__       :Optional[int]							 =			self._consolidate(batch[column_name]							)
							return batch