


"""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()
A	:  Dict            =     logging.get_logger(__name__)

A	:  List[str]            =     {
    "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",
}
A	:  Any            =     [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''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
			elif weight_type == "running_mean":
						__lowerCAmelCase							       =     value
			elif weight_type == "running_var":
						__lowerCAmelCase							       =     value
			elif weight_type == "num_batches_tracked":
						__lowerCAmelCase							       =     value
			elif weight_type == "inv_freq":
						__lowerCAmelCase							       =     value
			else:
						__lowerCAmelCase							       =     value

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


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			__lowerCAmelCase							       =     []
			__lowerCAmelCase							       =     fairseq_model.state_dict()

			__lowerCAmelCase							       =     hf_model.wavaveca_conformer.feature_extractor

			for name, value in fairseq_dict.items():
						__lowerCAmelCase							       =     False
						if "conv_layers" in name:
									load_conv_layer(
									    _UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					hf_model.config.feat_extract_norm == "group"       ,					)
									__lowerCAmelCase							       =     True
						else:
									for key, mapped_key in MAPPING.items():
												__lowerCAmelCase							       =     "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]:
															__lowerCAmelCase							       =     True
															if "*" in mapped_key:
																		__lowerCAmelCase							       =     name.split(_UpperCamelCase )[0].split("." )[-2]
																		__lowerCAmelCase							       =     mapped_key.replace("*"       ,					_UpperCamelCase )
															if "pos_bias_u" in name:
																		__lowerCAmelCase							       =     None
															elif "pos_bias_v" in name:
																		__lowerCAmelCase							       =     None
															elif "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"
															elif "running_mean" in name:
																		__lowerCAmelCase							       =     "running_mean"
															elif "inv_freq" in name:
																		__lowerCAmelCase							       =     "inv_freq"
															elif "running_var" in name:
																		__lowerCAmelCase							       =     "running_var"
															elif "num_batches_tracked" in name:
																		__lowerCAmelCase							       =     "num_batches_tracked"
															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 _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			__lowerCAmelCase							       =     full_name.split("conv_layers." )[-1]
			__lowerCAmelCase							       =     name.split("." )
			__lowerCAmelCase							       =     int(items[0] )
			__lowerCAmelCase							       =     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." )
									__lowerCAmelCase							       =     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." )
									__lowerCAmelCase							       =     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." )
									__lowerCAmelCase							       =     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." )
									__lowerCAmelCase							       =     value
									logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}." )
			else:
						unused_weights.append(_UpperCamelCase )


@torch.no_grad()
def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase=None       ,					_UpperCamelCase=None       ,					_UpperCamelCase=True ):

			'''simple docstring'''




			if config_path is not None:
						__lowerCAmelCase							       =     WavaVecaConformerConfig.from_pretrained(_UpperCamelCase       ,					hidden_act="swish" )
			else:
						__lowerCAmelCase							       =     WavaVecaConformerConfig()

			if "rope" in checkpoint_path:
						__lowerCAmelCase							       =     "rotary"

			if is_finetuned:
						if dict_path:
									__lowerCAmelCase							       =     Dictionary.load(_UpperCamelCase )

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

									# fairseq has the <pad> and <s> switched
									__lowerCAmelCase							       =     0
									__lowerCAmelCase							       =     1
									with open(_UpperCamelCase       ,					"w"       ,					encoding="utf-8" ) as vocab_handle:
												json.dump(_UpperCamelCase       ,					_UpperCamelCase )
									__lowerCAmelCase							       =     WavaVecaCTCTokenizer(
									    _UpperCamelCase       ,					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=_UpperCamelCase       ,					)
									__lowerCAmelCase							       =     True if config.feat_extract_norm == "layer" else False
									__lowerCAmelCase							       =     WavaVecaFeatureExtractor(
									    feature_size=1       ,					sampling_rate=1_6000       ,					padding_value=0       ,					do_normalize=_UpperCamelCase       ,					return_attention_mask=_UpperCamelCase       ,					)
									__lowerCAmelCase							       =     WavaVecaProcessor(feature_extractor=_UpperCamelCase       ,					tokenizer=_UpperCamelCase )
									processor.save_pretrained(_UpperCamelCase )

						__lowerCAmelCase							       =     WavaVecaConformerForCTC(_UpperCamelCase )
			else:
						__lowerCAmelCase							       =     WavaVecaConformerForPreTraining(_UpperCamelCase )

			if is_finetuned:
						__lowerCAmelCase       ,      __lowerCAmelCase       ,      __lowerCAmelCase							       =     fairseq.checkpoint_utils.load_model_ensemble_and_task(
						    [checkpoint_path]       ,					arg_overrides={"data": "/".join(dict_path.split("/" )[:-1] )} )
			else:
						__lowerCAmelCase							       =     argparse.Namespace(task="audio_pretraining" )
						__lowerCAmelCase							       =     fairseq.tasks.setup_task(_UpperCamelCase )

						__lowerCAmelCase       ,      __lowerCAmelCase       ,      __lowerCAmelCase							       =     fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]       ,					task=_UpperCamelCase )

			__lowerCAmelCase							       =     model[0].eval()

			recursively_load_weights(_UpperCamelCase       ,					_UpperCamelCase       ,					not is_finetuned )

			hf_wavavec.save_pretrained(_UpperCamelCase )


if __name__ == "__main__":
	A	:  str            =     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"
	)
	A	:  str            =     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"""


import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
class     _UpperCamelCase    (  lowerCAmelCase__						):
							'''simple docstring'''







							__UpperCAmelCase		:       List[str]							 ="""Speech2TextFeatureExtractor"""
							__UpperCAmelCase		:       Union[str, Any]							 ="""Speech2TextTokenizer"""

							def __init__(	self       ,  __a       ,  __a      ):
										super().__init__(__a       ,  __a      )
										__lowerCAmelCase							       =     self.feature_extractor
										__lowerCAmelCase							       =     False

							def __call__(	self       ,  *__a       ,  **__a      ):
										# For backward compatibility
										if self._in_target_context_manager:
													return self.current_processor(*__a       ,  **__a      )

										if "raw_speech" in kwargs:
													warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead."      )
													__lowerCAmelCase							       =     kwargs.pop("raw_speech"      )
										else:
													__lowerCAmelCase							       =     kwargs.pop("audio"       ,  __a      )
										__lowerCAmelCase							       =     kwargs.pop("sampling_rate"       ,  __a      )
										__lowerCAmelCase							       =     kwargs.pop("text"       ,  __a      )
										if len(__a      ) > 0:
													__lowerCAmelCase							       =     args[0]
													__lowerCAmelCase							       =     args[1:]

										if audio is None and text is None:
													raise ValueError("You need to specify either an `audio` or `text` input to process."      )

										if audio is not None:
													__lowerCAmelCase							       =     self.feature_extractor(__a       ,  *__a       ,  sampling_rate=__a       ,  **__a      )
										if text is not None:
													__lowerCAmelCase							       =     self.tokenizer(__a       ,  **__a      )

										if text is None:
													return inputs
										elif audio is None:
													return encodings
										else:
													__lowerCAmelCase							       =     encodings["input_ids"]
													return inputs

							def      snake_case							(	self       ,  *__a       ,  **__a      ):
										return self.tokenizer.batch_decode(*__a       ,  **__a      )

							def      snake_case							(	self       ,  *__a       ,  **__a      ):
										return self.tokenizer.decode(*__a       ,  **__a      )



							@contextmanager
							def      snake_case							(	self      ):
										warnings.warn(
										    "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
										    "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
										    "your audio inputs, or in a separate call."      )
										__lowerCAmelCase							       =     True
										__lowerCAmelCase							       =     self.tokenizer
										yield
										__lowerCAmelCase							       =     self.feature_extractor
										__lowerCAmelCase							       =     False




