





import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
A   : Optional[int]			=				logging.get_logger(__name__)

A   : Optional[int]			=				{
    '''post_extract_proj''': '''feature_projection.projection''',
    '''encoder.pos_conv.0''': '''encoder.pos_conv_embed.conv''',
    '''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''',
    '''quantizer.weight_proj''': '''quantizer.weight_proj''',
    '''quantizer.vars''': '''quantizer.codevectors''',
    '''project_q''': '''project_q''',
    '''final_proj''': '''project_hid''',
    '''w2v_encoder.proj''': '''ctc_proj''',
    '''mask_emb''': '''masked_spec_embed''',
}
A   : Optional[int]			=				[
    '''ctc_proj''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]







def        __lowerCamelCase (							__a       :List[str]   ,					__a       :Any   ,					__a       :Optional[Any]   ,					__a       :Tuple   ,					__a       :Tuple   ,					__a       :Optional[Any]   )		->    List[Any]:






		"""simple docstring"""

		for attribute in key.split("""."""   ):
				if is_finetuned:
						if attribute in ["quantizer", "project_q", "project_hid"]:
								# those layers are only relevant for pretraining and should be dropped
								return

						if attribute == "ctc_proj":
								# we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
								A__  				=       """lm_head"""

				A__  				=       getattr(__a   ,					__a   )

		if weight_type is not None:
				A__  				=       getattr(__a   ,					__a   ).shape
		else:
				A__  				=       hf_pointer.shape

		assert hf_shape == value.shape, (
		    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__  				=       value
		elif weight_type == "weight_g":
				A__  				=       value
		elif weight_type == "weight_v":
				A__  				=       value
		elif weight_type == "bias":
				A__  				=       value
		else:
				A__  				=       value

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







def        __lowerCamelCase (							__a       :str   ,					__a       :List[Any]   ,					__a       :Dict   )		->    Tuple:






		"""simple docstring"""

		A__  				=       []
		A__  				=       fairseq_model.state_dict()

		A__  				=       hf_model.unispeech.feature_extractor

		for name, value in fairseq_dict.items():
				A__  				=       False
				if "conv_layers" in name:
						load_conv_layer(
						    __a   ,					__a   ,					__a   ,					__a   ,					hf_model.config.feat_extract_norm == """group"""   ,					)
						A__  				=       True
				else:
						for key, mapped_key in MAPPING.items():
								A__  				=       """unispeech.""" + 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__  				=       True
										if "*" in mapped_key:
												A__  				=       name.split(__a   )[0].split("""."""   )[-2]
												A__  				=       mapped_key.replace("""*"""   ,					__a   )
										if "weight_g" in name:
												A__  				=       """weight_g"""
										elif "weight_v" in name:
												A__  				=       """weight_v"""
										elif "bias" in name:
												A__  				=       """bias"""
										elif "weight" in name:
												# TODO: don't match quantizer.weight_proj
												A__  				=       """weight"""
										else:
												A__  				=       None
										set_recursively(__a   ,					__a   ,					__a   ,					__a   ,					__a   ,					__a   )
								continue
				if not is_used:
						unused_weights.append(__a   )

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







def        __lowerCamelCase (							__a       :Union[str, Any]   ,					__a       :List[str]   ,					__a       :Union[str, Any]   ,					__a       :int   ,					__a       :Any   )		->    Tuple:






		"""simple docstring"""

		A__  				=       full_name.split("""conv_layers."""   )[-1]
		A__  				=       name.split("""."""   )
		A__  				=       int(items[0]   )
		A__  				=       int(items[1]   )

		if type_id == 0:
				if "bias" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
						    F'{full_name} has size {value.shape}, but'
						    F' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'
						)
						A__  				=       value
						logger.info(F'Feat extract conv layer {layer_id} was initialized from {full_name}.'   )
				elif "weight" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
						    F'{full_name} has size {value.shape}, but'
						    F' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'
						)
						A__  				=       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:
						assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
						    F'{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'
						    " found."
						)
						A__  				=       value
						logger.info(F'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'   )
				elif "weight" in name:
						assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
						    F'{full_name} has size {value.shape}, but'
						    F' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'
						)
						A__  				=       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        __lowerCamelCase (							__a       :Dict   ,					__a       :List[str]   ,					__a       :Optional[Any]=None   ,					__a       :Optional[Any]=None   ,					__a       :int=True   )		->    List[Any]:






		"""simple docstring"""

		if config_path is not None:
				A__  				=       UniSpeechConfig.from_pretrained(__a   )
		else:
				A__  				=       UniSpeechConfig()

		if is_finetuned:
				if dict_path:
						A__  				=       Dictionary.load_from_json(__a   )

						# important change bos & pad token id since CTC symbol is <pad> and
						# not <s> as in fairseq
						A__  				=       target_dict.pad_index
						A__  				=       target_dict.bos_index
						A__  				=       target_dict.eos_index
						A__  				=       len(target_dict.symbols   )
						A__  				=       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   )
						A__  				=       target_dict.indices

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

				A__  				=       UniSpeechForCTC(__a   )
		else:
				A__  				=       UniSpeechForPreTraining(__a   )

		if is_finetuned:
				A__				,      A__				,      A__  				=       fairseq.checkpoint_utils.load_model_ensemble_and_task(
				    [checkpoint_path]   ,					arg_overrides={"""data""": """/""".join(dict_path.split("""/"""   )[:-1]   ), """w2v_path""": checkpoint_path}   )
		else:
				A__				,      A__				,      A__  				=       fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]   )

		A__  				=       model[0].eval()

		recursively_load_weights(__a   ,					__a   ,					__a   )

		hf_unispeech.save_pretrained(__a   )


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










import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
					from transformers import WhisperFeatureExtractor

if is_torch_available():
					import torch

A   : Dict			=				random.Random()







def        __lowerCamelCase (							__a       :Dict   ,					__a       :str=1.0   ,					__a       :List[Any]=None   ,					__a       :List[str]=None   )		->    Optional[int]:






		"""simple docstring"""

		if rng is None:
				A__  				=       global_rng

		A__  				=       []
		for batch_idx in range(shape[0]   ):
				values.append([]   )
				for _ in range(shape[1]   ):
						values[-1].append(rng.random() * scale   )

		return values


@require_torch
@require_torchaudio
class   A      (unittest.TestCase   ):
			'''simple docstring'''




			def __init__(      self       :   Dict							,     __lowerCAmelCase       :   Any							,     __lowerCAmelCase       :   str=7							,     __lowerCAmelCase       :   List[Any]=4_00							,     __lowerCAmelCase       :   Optional[Any]=20_00							,     __lowerCAmelCase       :   Dict=10							,     __lowerCAmelCase       :   Union[str, Any]=1_60							,     __lowerCAmelCase       :   List[Any]=8							,     __lowerCAmelCase       :   List[Any]=0.0							,     __lowerCAmelCase       :   int=40_00							,     __lowerCAmelCase       :   Any=False							,     __lowerCAmelCase       :   List[str]=True							,     )     ->     Optional[int]:







					"""simple docstring"""






					A__  				=       parent
					A__  				=       batch_size
					A__  				=       min_seq_length
					A__  				=       max_seq_length
					A__  				=       (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
					A__  				=       padding_value
					A__  				=       sampling_rate
					A__  				=       return_attention_mask
					A__  				=       do_normalize
					A__  				=       feature_size
					A__  				=       chunk_length
					A__  				=       hop_length






			def        a_	(      self       :   Optional[Any]       )     ->     Optional[int]:







					"""simple docstring"""






					return {
					    "feature_size": self.feature_size,
					    "hop_length": self.hop_length,
					    "chunk_length": self.chunk_length,
					    "padding_value": self.padding_value,
					    "sampling_rate": self.sampling_rate,
					    "return_attention_mask": self.return_attention_mask,
					    "do_normalize": self.do_normalize,
					}







			def        a_	(      self       :   List[str]							,     __lowerCAmelCase       :   List[Any]=False							,     __lowerCAmelCase       :   int=False       )     ->     str:







					"""simple docstring"""






					def _flatten(__lowerCAmelCase       :   Optional[int]       ):
							return list(itertools.chain(*__lowerCAmelCase       )       )

					if equal_length:
							A__  				=       [floats_list((self.max_seq_length, self.feature_size)       ) for _ in range(self.batch_size       )]
					else:
							# make sure that inputs increase in size
							A__  				=       [
							    floats_list((x, self.feature_size)       )
							    for x in range(self.min_seq_length							,     self.max_seq_length							,     self.seq_length_diff       )
							]
					if numpify:
							A__  				=       [np.asarray(__lowerCAmelCase       ) for x in speech_inputs]
					return speech_inputs


@require_torch
@require_torchaudio
class   A      (SCREAMING_SNAKE_CASE       ,					unittest.TestCase   ):
			'''simple docstring'''




			__lowerCamelCase     :	int          =   WhisperFeatureExtractor if is_speech_available() else None






			def        a_	(      self       :   Any       )     ->     str:







					"""simple docstring"""






					A__  				=       WhisperFeatureExtractionTester(self       )






			def        a_	(      self       :   Dict       )     ->     Tuple:







					"""simple docstring"""






					A__  				=       self.feature_extraction_class(**self.feat_extract_dict       )

					with tempfile.TemporaryDirectory() as tmpdirname:
							A__  				=       feat_extract_first.save_pretrained(__lowerCAmelCase       )[0]
							check_json_file_has_correct_format(__lowerCAmelCase       )
							A__  				=       self.feature_extraction_class.from_pretrained(__lowerCAmelCase       )

					A__  				=       feat_extract_first.to_dict()
					A__  				=       feat_extract_second.to_dict()
					A__  				=       feat_extract_first.mel_filters
					A__  				=       feat_extract_second.mel_filters
					self.assertTrue(np.allclose(__lowerCAmelCase							,     __lowerCAmelCase       )       )
					self.assertEqual(__lowerCAmelCase							,     __lowerCAmelCase       )






			def        a_	(      self       :   Dict       )     ->     str:







					"""simple docstring"""






					A__  				=       self.feature_extraction_class(**self.feat_extract_dict       )

					with tempfile.TemporaryDirectory() as tmpdirname:
							A__  				=       os.path.join(__lowerCAmelCase							,     """feat_extract.json"""       )
							feat_extract_first.to_json_file(__lowerCAmelCase       )
							A__  				=       self.feature_extraction_class.from_json_file(__lowerCAmelCase       )

					A__  				=       feat_extract_first.to_dict()
					A__  				=       feat_extract_second.to_dict()
					A__  				=       feat_extract_first.mel_filters
					A__  				=       feat_extract_second.mel_filters
					self.assertTrue(np.allclose(__lowerCAmelCase							,     __lowerCAmelCase       )       )
					self.assertEqual(__lowerCAmelCase							,     __lowerCAmelCase       )






			def        a_	(      self       :   Union[str, Any]       )     ->     List[Any]:







					"""simple docstring"""






					A__  				=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
					# create three inputs of length 800, 1000, and 1200
					A__  				=       [floats_list((1, x)       )[0] for x in range(8_00							,     14_00							,     2_00       )]
					A__  				=       [np.asarray(__lowerCAmelCase       ) for speech_input in speech_inputs]

					# Test feature size
					A__  				=       feature_extractor(__lowerCAmelCase							,     padding="""max_length"""							,     return_tensors="""np"""       ).input_features
					self.assertTrue(input_features.ndim == 3       )
					self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames       )
					self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size       )

					# Test not batched input
					A__  				=       feature_extractor(speech_inputs[0]							,     return_tensors="""np"""       ).input_features
					A__  				=       feature_extractor(np_speech_inputs[0]							,     return_tensors="""np"""       ).input_features
					self.assertTrue(np.allclose(__lowerCAmelCase							,     __lowerCAmelCase							,     atol=1e-3       )       )

					# Test batched
					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""np"""       ).input_features
					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""np"""       ).input_features
					for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase							,     __lowerCAmelCase       ):
							self.assertTrue(np.allclose(__lowerCAmelCase							,     __lowerCAmelCase							,     atol=1e-3       )       )

					# Test 2-D numpy arrays are batched.
					A__  				=       [floats_list((1, x)       )[0] for x in (8_00, 8_00, 8_00)]
					A__  				=       np.asarray(__lowerCAmelCase       )
					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""np"""       ).input_features
					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""np"""       ).input_features
					for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase							,     __lowerCAmelCase       ):
							self.assertTrue(np.allclose(__lowerCAmelCase							,     __lowerCAmelCase							,     atol=1e-3       )       )

					# Test truncation required
					A__  				=       [floats_list((1, x)       )[0] for x in range(2_00							,     (feature_extractor.n_samples + 5_00)							,     2_00       )]
					A__  				=       [np.asarray(__lowerCAmelCase       ) for speech_input in speech_inputs]

					A__  				=       [x[: feature_extractor.n_samples] for x in speech_inputs]
					A__  				=       [np.asarray(__lowerCAmelCase       ) for speech_input in speech_inputs_truncated]

					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""np"""       ).input_features
					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""np"""       ).input_features
					for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase							,     __lowerCAmelCase       ):
							self.assertTrue(np.allclose(__lowerCAmelCase							,     __lowerCAmelCase							,     atol=1e-3       )       )






			def        a_	(      self       :   Dict       )     ->     Union[str, Any]:







					"""simple docstring"""






					import torch

					A__  				=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
					A__  				=       np.random.rand(1_00							,     32       ).astype(np.floataa       )
					A__  				=       np_speech_inputs.tolist()

					for inputs in [py_speech_inputs, np_speech_inputs]:
							A__  				=       feature_extractor.pad([{"""input_features""": inputs}]							,     return_tensors="""np"""       )
							self.assertTrue(np_processed.input_features.dtype == np.floataa       )
							A__  				=       feature_extractor.pad([{"""input_features""": inputs}]							,     return_tensors="""pt"""       )
							self.assertTrue(pt_processed.input_features.dtype == torch.floataa       )






			def        a_	(      self       :   List[Any]							,     __lowerCAmelCase       :   Any       )     ->     Dict:







					"""simple docstring"""






					A__  				=       load_dataset("""hf-internal-testing/librispeech_asr_dummy"""							,     """clean"""							,     split="""validation"""       )
					# automatic decoding with librispeech
					A__  				=       ds.sort("""id"""       ).select(range(__lowerCAmelCase       )       )[:num_samples]["""audio"""]

					return [x["array"] for x in speech_samples]






			def        a_	(      self       :   Dict       )     ->     Any:







					"""simple docstring"""






					A__  				=       torch.tensor(
					    [
					        0.1_1_9_3, -0.0_9_4_6, -0.1_0_9_8, -0.0_1_9_6, 0.0_2_2_5, -0.0_6_9_0, -0.1_7_3_6, 0.0_9_5_1,
					        0.0_9_7_1, -0.0_8_1_7, -0.0_7_0_2, 0.0_1_6_2, 0.0_2_6_0, 0.0_0_1_7, -0.0_1_9_2, -0.1_6_7_8,
					        0.0_7_0_9, -0.1_8_6_7, -0.0_6_5_5, -0.0_2_7_4, -0.0_2_3_4, -0.1_8_8_4, -0.0_5_1_6, -0.0_5_5_4,
					        -0.0_2_7_4, -0.1_4_2_5, -0.1_4_2_3, 0.0_8_3_7, 0.0_3_7_7, -0.0_8_5_4
					    ]       )
					# fmt: on

					A__  				=       self._load_datasamples(1       )
					A__  				=       WhisperFeatureExtractor()
					A__  				=       feature_extractor(__lowerCAmelCase							,     return_tensors="""pt"""       ).input_features
					self.assertEqual(input_features.shape							,     (1, 80, 30_00)       )
					self.assertTrue(torch.allclose(input_features[0, 0, :30]							,     __lowerCAmelCase							,     atol=1e-4       )       )







			def        a_	(      self       :   Dict       )     ->     Tuple:







					"""simple docstring"""






					A__  				=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
					A__  				=       self._load_datasamples(1       )[0]
					A__  				=       ((audio - audio.min()) / (audio.max() - audio.min())) * 6_55_35  # Rescale to [0, 65535] to show issue
					A__  				=       feat_extract.zero_mean_unit_var_norm([audio]							,     attention_mask=__lowerCAmelCase       )[0]

					self.assertTrue(np.all(np.mean(__lowerCAmelCase       ) < 1e-3       )       )
					self.assertTrue(np.all(np.abs(np.var(__lowerCAmelCase       ) - 1       ) < 1e-3       )       )



