






import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


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

__UpperCAmelCase							:   Tuple							=  {
    "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",
}
__UpperCAmelCase							:   List[str]							=  [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
    "label_embeddings_concat",
    "speaker_proj",
    "layer_norm_for_extract",
]





def 					lowercase_ (    __snake_case     :							int	,	__snake_case     :							Any	,	__snake_case     :							int	,	__snake_case     :							List[Any]	,	__snake_case     :							Union[str, Any]      )     -> Optional[int]:





				'''simple docstring'''


				for attribute in key.split("."      ):
								snake_case__      :Any         =   getattr(__snake_case	,	__snake_case      )

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

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





def 					lowercase_ (    __snake_case     :							str	,	__snake_case     :							Optional[int]      )     -> Optional[Any]:





				'''simple docstring'''


				snake_case__      :Dict         =   []
				snake_case__      :List[str]         =   fairseq_model.state_dict()

				snake_case__      :List[str]         =   hf_model.unispeech_sat.feature_extractor

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

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





def 					lowercase_ (    __snake_case     :							Dict	,	__snake_case     :							List[str]	,	__snake_case     :							Union[str, Any]	,	__snake_case     :							List[str]	,	__snake_case     :							Dict      )     -> List[str]:





				'''simple docstring'''


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





@torch.no_grad()
def 					lowercase_ (    __snake_case     :							Optional[int]	,	__snake_case     :							Union[str, Any]	,	__snake_case     :							Union[str, Any]=None	,	__snake_case     :							List[str]=None	,	__snake_case     :							Optional[int]=True      )     -> Optional[Any]:





				'''simple docstring'''


				if config_path is not None:
								snake_case__      :int         =   UniSpeechSatConfig.from_pretrained(__snake_case      )
				else:
								snake_case__      :Tuple         =   UniSpeechSatConfig()

				snake_case__      :Optional[Any]         =   ""

				if is_finetuned:
								snake_case__      :Optional[Any]         =   UniSpeechSatForCTC(__snake_case      )
				else:
								snake_case__      :Optional[Any]         =   UniSpeechSatForPreTraining(__snake_case      )

				snake_case__    ,						snake_case__    ,						snake_case__      :Union[str, Any]         =   fairseq.checkpoint_utils.load_model_ensemble_and_task(
				    [checkpoint_path]	,	arg_overrides={"data": "/".join(dict_path.split("/"      )[:-1]      )}      )
				snake_case__      :Tuple         =   model[0].eval()

				recursively_load_weights(__snake_case	,	__snake_case      )

				hf_wavavec.save_pretrained(__snake_case      )


if __name__ == "__main__":
	__UpperCAmelCase							:   Optional[Any]							=  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							:   Optional[Any]							=  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
	)







from typing import Any





def 					lowercase_ (    __snake_case     :							list	,	__snake_case     :							list	,	__snake_case     :							dict	,	__snake_case     :							dict	,	__snake_case     :							dict	,	)     -> list:





				'''simple docstring'''


				_validation(
				    __snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	)
				# Creates data structures and fill initial step
				snake_case__      :dict         =   {}
				snake_case__      :dict         =   {}
				for state in states_space:
								snake_case__      :List[Any]         =   observations_space[0]
								snake_case__      :str         =   (
								    initial_probabilities[state] * emission_probabilities[state][observation]
								)
								snake_case__      :str         =   None

				# Fills the data structure with the probabilities of
				# different transitions and pointers to previous states
				for o in range(1	,	len(__snake_case      )      ):
								snake_case__      :Any         =   observations_space[o]
								snake_case__      :Tuple         =   observations_space[o - 1]
								for state in states_space:
												# Calculates the argmax for probability function
												snake_case__      :Tuple         =   ""
												snake_case__      :Union[str, Any]         =   -1
												for k_state in states_space:
																snake_case__      :int         =   (
																    probabilities[(k_state, prior_observation)]
																    * transition_probabilities[k_state][state]
																    * emission_probabilities[state][observation]
																)
																if probability > max_probability:
																				snake_case__      :str         =   probability
																				snake_case__      :Tuple         =   k_state

            # Update probabilities and pointers dicts
												snake_case__      :List[str]         =   (
												    probabilities[(arg_max, prior_observation)]
												    * transition_probabilities[arg_max][state]
												    * emission_probabilities[state][observation]
												)

												snake_case__      :List[str]         =   arg_max

    # The final observation
				snake_case__      :str         =   observations_space[len(__snake_case      ) - 1]

				# argmax for given final observation
				snake_case__      :Optional[int]         =   ""
				snake_case__      :List[str]         =   -1
				for k_state in states_space:
								snake_case__      :List[str]         =   probabilities[(k_state, final_observation)]
								if probability > max_probability:
												snake_case__      :List[str]         =   probability
												snake_case__      :int         =   k_state
				snake_case__      :Any         =   arg_max

				# Process pointers backwards
				snake_case__      :int         =   last_state
				snake_case__      :List[str]         =   []
				for o in range(len(__snake_case      ) - 1	,	-1	,	-1      ):
								result.append(__snake_case      )
								snake_case__      :List[str]         =   pointers[previous, observations_space[o]]
				result.reverse()

				return result





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	)     -> None:





				'''simple docstring'''


				_validate_not_empty(
				    __snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	)
				_validate_lists(__snake_case	,	__snake_case      )
				_validate_dicts(
				    __snake_case	,	__snake_case	,	__snake_case      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	)     -> None:





				'''simple docstring'''


				if not all(
				    [
				        observations_space,
				        states_space,
				        initial_probabilities,
				        transition_probabilities,
				        emission_probabilities,
				    ]      ):
								raise ValueError("There's an empty parameter"      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any      )     -> None:





				'''simple docstring'''


				_validate_list(__snake_case	,	"observations_space"      )
				_validate_list(__snake_case	,	"states_space"      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							str      )     -> None:





				'''simple docstring'''


				if not isinstance(_object	,	__snake_case      ):
								snake_case__      :Optional[int]         =   F'{var_name} must be a list'
								raise ValueError(__snake_case      )
				else:
								for x in _object:
												if not isinstance(__snake_case	,	__snake_case      ):
																snake_case__      :Any         =   F'{var_name} must be a list of strings'
																raise ValueError(__snake_case      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	)     -> None:





				'''simple docstring'''


				_validate_dict(__snake_case	,	"initial_probabilities"	,	__snake_case      )
				_validate_nested_dict(__snake_case	,	"transition_probabilities"      )
				_validate_nested_dict(__snake_case	,	"emission_probabilities"      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							str      )     -> None:





				'''simple docstring'''


				_validate_dict(_object	,	__snake_case	,	__snake_case      )
				for x in _object.values():
								_validate_dict(__snake_case	,	__snake_case	,	__snake_case	,	__snake_case      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							str	,	__snake_case     :							type	,	__snake_case     :							bool = False      )     -> None:





				'''simple docstring'''


				if not isinstance(_object	,	__snake_case      ):
								snake_case__      :str         =   F'{var_name} must be a dict'
								raise ValueError(__snake_case      )
				if not all(isinstance(__snake_case	,	__snake_case      ) for x in _object      ):
								snake_case__      :List[Any]         =   F'{var_name} all keys must be strings'
								raise ValueError(__snake_case      )
				if not all(isinstance(__snake_case	,	__snake_case      ) for x in _object.values()      ):
								snake_case__      :Optional[int]         =   "nested dictionary " if nested else ""
								snake_case__      :int         =   F'{var_name} {nested_text}all values must be {value_type.__name__}'
								raise ValueError(__snake_case      )


if __name__ == "__main__":
	from doctest import testmod

	testmod()