




"""simple docstring"""







import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
__A   :  Dict					=      logging.get_logger(__name__)

__A   :  Union[str, Any]					=      {
    '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',
    'encoder.layer_norm_for_extract': 'layer_norm_for_extract',
    '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',
    'label_embs_concat': 'label_embeddings_concat',
    'mask_emb': 'masked_spec_embed',
    'spk_proj': 'speaker_proj',
}
__A   :  List[Any]					=      [
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'label_embeddings_concat',
    'speaker_proj',
    'layer_norm_for_extract',
]

def 						snake_case__       (					_lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase				)			->Dict:


		for attribute in key.split("."				):
				__lowercase				:		Optional[Any]												=				getattr(_lowerCamelCase,  _lowerCamelCase				)

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

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

def 						snake_case__       (					_lowerCamelCase,  _lowerCamelCase				)			->Tuple:
		__lowercase				:		List[Any]												=				[]
		__lowercase				:		str												=				fairseq_model.state_dict()

		__lowercase				:		Any												=				hf_model.unispeech_sat.feature_extractor

		for name, value in fairseq_dict.items():
				__lowercase				:		Any												=				False
				if "conv_layers" in name:
						load_conv_layer(
						    _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  hf_model.config.feat_extract_norm == "group",  )
						__lowercase				:		int												=				True
				else:
						for key, mapped_key in MAPPING.items():
								__lowercase				:		Optional[Any]												=				"unispeech_sat." + 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]:
										if "layer_norm_for_extract" in name and (".".join(name.split("."				)[:-1]				) != key):
												# special case since naming is very similar
												continue
										__lowercase				:		List[str]												=				True
										if "*" in mapped_key:
												__lowercase				:		Tuple												=				name.split(_lowerCamelCase				)[0].split("."				)[-2]
												__lowercase				:		Optional[Any]												=				mapped_key.replace("*",  _lowerCamelCase				)
										if "weight_g" in name:
												__lowercase				:		Any												=				"weight_g"
										elif "weight_v" in name:
												__lowercase				:		int												=				"weight_v"
										elif "bias" in name:
												__lowercase				:		Optional[Any]												=				"bias"
										elif "weight" in name:
												# TODO: don't match quantizer.weight_proj
												__lowercase				:		Any												=				"weight"
										else:
												__lowercase				:		Dict												=				None
										set_recursively(_lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase				)
								continue
				if not is_used:
						unused_weights.append(_lowerCamelCase				)

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

def 						snake_case__       (					_lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase				)			->Union[str, Any]:
		__lowercase				:		Any												=				full_name.split("conv_layers."				)[-1]
		__lowercase				:		int												=				name.split("."				)
		__lowercase				:		Optional[Any]												=				int(items[0]				)
		__lowercase				:		List[Any]												=				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.'				)
						__lowercase				:		int												=				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.'				)
						__lowercase				:		int												=				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[layer_id].layer_norm.bias.data.shape} was found.'				)
						__lowercase				:		Any												=				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[layer_id].layer_norm.weight.data.shape} was found.'				)
						__lowercase				:		int												=				value
						logger.info(F'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'				)
		else:
				unused_weights.append(_lowerCamelCase				)

@torch.no_grad()
def 						snake_case__       (					_lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase=None,  _lowerCamelCase=None,  _lowerCamelCase=True				)			->Optional[int]:


		if config_path is not None:
				__lowercase				:		Tuple												=				UniSpeechSatConfig.from_pretrained(_lowerCamelCase				)
		else:
				__lowercase				:		Union[str, Any]												=				UniSpeechSatConfig()

		__lowercase				:		List[str]												=				""

		if is_finetuned:
				__lowercase				:		str												=				UniSpeechSatForCTC(_lowerCamelCase				)
		else:
				__lowercase				:		Optional[Any]												=				UniSpeechSatForPreTraining(_lowerCamelCase				)

		__lowercase				:		Optional[int]												=				fairseq.checkpoint_utils.load_model_ensemble_and_task(
		    [checkpoint_path],  arg_overrides={"data": "/".join(dict_path.split("/"				)[:-1]				)}				)
		__lowercase				:		Any												=				model[0].eval()

		recursively_load_weights(_lowerCamelCase,  _lowerCamelCase				)

		hf_wavavec.save_pretrained(_lowerCamelCase				)


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   :  Dict					=      parser.parse_args()
						convert_unispeech_sat_checkpoint(
						    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
						)






"""simple docstring"""







import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate

def 						snake_case__       (					_lowerCamelCase				)			->Dict:






		"""simple docstring"""


		return x + 2





class       lowerCAmelCase__       (			unittest.TestCase    ):


				"""simple docstring"""





				def       snake_case    (     self					:	List[str]						):
						__lowercase				:		List[Any]												=				"x = 3"
						__lowercase				:		Optional[int]												=				{}
						__lowercase				:		List[Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result == 3
						self.assertDictEqual(lowercase__    ,   {"x": 3}						)

						__lowercase				:		Optional[int]												=				"x = y"
						__lowercase				:		Union[str, Any]												=				{"y": 5}
						__lowercase				:		str												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 5, "y": 5}						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Dict												=				"y = add_two(x)"
						__lowercase				:		List[Any]												=				{"x": 3}
						__lowercase				:		Optional[int]												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 5}						)

						# Won't work without the tool
						with CaptureStdout() as out:
								__lowercase				:		Any												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result is None
						assert "tried to execute add_two" in out.out



				def       snake_case    (     self					:	Any						):
						__lowercase				:		Optional[Any]												=				"x = 3"
						__lowercase				:		List[Any]												=				{}
						__lowercase				:		Dict												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result == 3
						self.assertDictEqual(lowercase__    ,   {"x": 3}						)



				def       snake_case    (     self					:	str						):
						__lowercase				:		Optional[int]												=				"test_dict = {'x': x, 'y': add_two(x)}"
						__lowercase				:		List[str]												=				{"x": 3}
						__lowercase				:		Tuple												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 5}						)
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_dict": {"x": 3, "y": 5}}						)



				def       snake_case    (     self					:	int						):
						__lowercase				:		Tuple												=				"x = 3\ny = 5"
						__lowercase				:		Optional[Any]												=				{}
						__lowercase				:		List[Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 5}						)



				def       snake_case    (     self					:	Any						):
						__lowercase				:		str												=				"text = f'This is x: {x}.'"
						__lowercase				:		Any												=				{"x": 3}
						__lowercase				:		List[Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == "This is x: 3."
						self.assertDictEqual(lowercase__    ,   {"x": 3, "text": "This is x: 3."}						)



				def       snake_case    (     self					:	str						):
						__lowercase				:		Optional[Any]												=				"if x <= 3:\n    y = 2\nelse:\n    y = 5"
						__lowercase				:		int												=				{"x": 3}
						__lowercase				:		Optional[int]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 2
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 2}						)

						__lowercase				:		Optional[int]												=				{"x": 8}
						__lowercase				:		Union[str, Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 8, "y": 5}						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Optional[Any]												=				"test_list = [x, add_two(x)]"
						__lowercase				:		str												=				{"x": 3}
						__lowercase				:		Any												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						self.assertListEqual(lowercase__    ,   [3, 5]						)
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_list": [3, 5]}						)



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		str												=				"y = x"
						__lowercase				:		str												=				{"x": 3}
						__lowercase				:		int												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result == 3
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 3}						)



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		str												=				"test_list = [x, add_two(x)]\ntest_list[1]"
						__lowercase				:		Optional[Any]												=				{"x": 3}
						__lowercase				:		Any												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_list": [3, 5]}						)

						__lowercase				:		Tuple												=				"test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"
						__lowercase				:		Union[str, Any]												=				{"x": 3}
						__lowercase				:		Tuple												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_dict": {"x": 3, "y": 5}}						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Dict												=				"x = 0\nfor i in range(3):\n    x = i"
						__lowercase				:		str												=				{}
						__lowercase				:		int												=				evaluate(lowercase__    ,   {"range": range}    ,   state=lowercase__						)
						assert result == 2
						self.assertDictEqual(lowercase__    ,   {"x": 2, "i": 2}						)
