



"""simple docstring"""






import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
__snake_case				:  List[Any] =       logging.get_logger(__name__)

__snake_case				:  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',
}
__snake_case				:  Tuple =       [
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'label_embeddings_concat',
    'speaker_proj',
    'layer_norm_for_extract',
]


def      _lowercase							(    __snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case      )		->  Optional[int]:
			for attribute in key.split("."      ):
						__lowerCAmelCase :    List[Any]				=			getattr(__snake_case   ,__snake_case      )

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

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

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

			__lowerCAmelCase :    List[Any]				=			""

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

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

			recursively_load_weights(__snake_case   ,__snake_case      )

			hf_wavavec.save_pretrained(__snake_case      )


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




"""simple docstring"""






def      _lowercase							(    __snake_case   ,__snake_case      )		->  int:
			return int((input_a, input_a).count(1      ) != 0      )


def      _lowercase							(    )		->  None:
			assert or_gate(0   ,0      ) == 0
			assert or_gate(0   ,1      ) == 1
			assert or_gate(1   ,0      ) == 1
			assert or_gate(1   ,1      ) == 1


if __name__ == "__main__":
	print(or_gate(0, 1))
	print(or_gate(1, 0))
	print(or_gate(0, 0))
	print(or_gate(1, 1))