






'''simple docstring'''







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()
UpperCamelCase    :      List[Any]  			=		logging.get_logger(__name__)

UpperCamelCase    :      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",
    "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",
}
UpperCamelCase    :      Optional[int]  			=		[
    "ctc_proj",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]
def 						A__			(				__lowerCAmelCase		:			int   ,    __lowerCAmelCase		:			Tuple   ,    __lowerCAmelCase		:			str   ,    __lowerCAmelCase		:			str   ,    __lowerCAmelCase		:			str   ,    __lowerCAmelCase		:			List[Any]						):


						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
																								lowerCamelCase__         =  "lm_head"

												lowerCamelCase__         =  getattr(lowerCamelCase__   ,    lowerCamelCase__						)

						if weight_type is not None:
												lowerCamelCase__         =  getattr(lowerCamelCase__   ,    lowerCamelCase__						).shape
						else:
												lowerCamelCase__         =  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":
												lowerCamelCase__         =  value
						elif weight_type == "weight_g":
												lowerCamelCase__         =  value
						elif weight_type == "weight_v":
												lowerCamelCase__         =  value
						elif weight_type == "bias":
												lowerCamelCase__         =  value
						else:
												lowerCamelCase__         =  value

						logger.info(F'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''						)
def 						A__			(				__lowerCAmelCase		:			Dict   ,    __lowerCAmelCase		:			Optional[Any]   ,    __lowerCAmelCase		:			Union[str, Any]						):
						lowerCamelCase__         =  []
						lowerCamelCase__         =  fairseq_model.state_dict()

						lowerCamelCase__         =  hf_model.unispeech.feature_extractor

						for name, value in fairseq_dict.items():
												lowerCamelCase__         =  False
												if "conv_layers" in name:
																		load_conv_layer(
																		    lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__   ,    hf_model.config.feat_extract_norm == """group"""   ,    )
																		lowerCamelCase__         =  True
												else:
																		for key, mapped_key in MAPPING.items():
																								lowerCamelCase__         =  "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]:
																														lowerCamelCase__         =  True
																														if "*" in mapped_key:
																																				lowerCamelCase__         =  name.split(lowerCamelCase__						)[0].split("""."""						)[-2]
																																				lowerCamelCase__         =  mapped_key.replace("""*"""   ,    lowerCamelCase__						)
																														if "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"
																														else:
																																				lowerCamelCase__         =  None
																														set_recursively(lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__						)
																								continue
												if not is_used:
																		unused_weights.append(lowerCamelCase__						)

						logger.warning(F'''Unused weights: {unused_weights}'''						)
def 						A__			(				__lowerCAmelCase		:			List[Any]   ,    __lowerCAmelCase		:			Optional[Any]   ,    __lowerCAmelCase		:			Tuple   ,    __lowerCAmelCase		:			Dict   ,    __lowerCAmelCase		:			Dict						):
						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:
																		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.'''
																		)
																		lowerCamelCase__         =  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.'''
																		)
																		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:
																		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."
																		)
																		lowerCamelCase__         =  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.'''
																		)
																		lowerCamelCase__         =  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 						A__			(				__lowerCAmelCase		:			int   ,    __lowerCAmelCase		:			Dict   ,    __lowerCAmelCase		:			str=None   ,    __lowerCAmelCase		:			Optional[int]=None   ,    __lowerCAmelCase		:			str=True						):


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

						if is_finetuned:
												if dict_path:
																		lowerCamelCase__         =  Dictionary.load_from_json(lowerCamelCase__						)

																		# 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(lowerCamelCase__   ,    """vocab.json"""						)
																		if not os.path.isdir(lowerCamelCase__						):
																								logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(lowerCamelCase__						)						)
																								return
																		os.makedirs(lowerCamelCase__   ,    exist_ok=lowerCamelCase__						)
																		lowerCamelCase__         =  target_dict.indices

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

												lowerCamelCase__         =  UniSpeechForCTC(lowerCamelCase__						)
						else:
												lowerCamelCase__         =  UniSpeechForPreTraining(lowerCamelCase__						)

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

						lowerCamelCase__         =  model[0].eval()

						recursively_load_weights(lowerCamelCase__   ,    lowerCamelCase__   ,    lowerCamelCase__						)

						hf_unispeech.save_pretrained(lowerCamelCase__						)


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









'''simple docstring'''







from manim import *




class 		UpperCamelCase__    (a			):




		'''simple docstring'''





		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  Rectangle(height=0.5   ,width=0.5       )
								lowerCamelCase__         =  Rectangle(height=0.46   ,width=0.46       ).set_stroke(width=0       )

								lowerCamelCase__         =  [mem.copy() for i in range(6       )]
								lowerCamelCase__         =  [mem.copy() for i in range(6       )]
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  VGroup(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  Text("""CPU"""   ,font_size=24       )
								lowerCamelCase__         =  Group(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0.5   ,aligned_edge=_lowerCAmelCase       )
								cpu.move_to([-2.5, -0.5, 0]       )
								self.add(_lowerCAmelCase       )

								lowerCamelCase__         =  [mem.copy() for i in range(1       )]
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )
								lowerCamelCase__         =  Text("""GPU"""   ,font_size=24       )
								lowerCamelCase__         =  Group(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0.5   ,aligned_edge=_lowerCAmelCase       )
								gpu.align_to(_lowerCAmelCase   ,_lowerCAmelCase       )
								gpu.set_x(gpu.get_x() - 1       )

								self.add(_lowerCAmelCase       )

								lowerCamelCase__         =  [mem.copy() for i in range(6       )]
								lowerCamelCase__         =  VGroup(*_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0       )

								lowerCamelCase__         =  Text("""Model"""   ,font_size=24       )
								lowerCamelCase__         =  Group(_lowerCAmelCase   ,_lowerCAmelCase       ).arrange(_lowerCAmelCase   ,buff=0.5   ,aligned_edge=_lowerCAmelCase       )
								model.move_to([3, -1.0, 0]       )

								self.play(
								    Create(_lowerCAmelCase   ,run_time=1       )   ,Create(_lowerCAmelCase   ,run_time=1       )   ,Create(_lowerCAmelCase   ,run_time=1       )   ,)

								lowerCamelCase__         =  MarkupText(
								    F'''First, an empty model skeleton is loaded\ninto <span fgcolor=\'{YELLOW}\'>memory</span> without using much RAM.'''   ,font_size=24   ,)

								lowerCamelCase__         =  Square(side_length=2.2       )
								key.move_to([-5, 2, 0]       )

								lowerCamelCase__         =  MarkupText(
								    F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''   ,font_size=18   ,)

								key_text.move_to([-5, 2.4, 0]       )

								step_a.move_to([2, 2, 0]       )
								self.play(Write(_lowerCAmelCase   ,run_time=2.5       )   ,Write(_lowerCAmelCase       )   ,Write(_lowerCAmelCase       )       )

								self.add(_lowerCAmelCase       )

								lowerCamelCase__         =  []
								lowerCamelCase__         =  []
								lowerCamelCase__         =  []
								for i, rect in enumerate(_lowerCAmelCase       ):
														lowerCamelCase__         =  Rectangle(height=0.46   ,width=0.46       ).set_stroke(width=0.0       ).set_fill(_lowerCAmelCase   ,opacity=0.7       )
														cpu_target.move_to(_lowerCAmelCase       )
														cpu_target.generate_target()
														lowerCamelCase__         =  0.46 / 4
														lowerCamelCase__         =  0.46 / 3

														if i == 0:
																				cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT       )   ,buff=0.02   ,direction=_lowerCAmelCase       )
																				cpu_target.target.set_x(cpu_target.target.get_x() + 0.1       )
														elif i == 3:
																				cpu_target.target.next_to(cpu_targs[0].target   ,direction=_lowerCAmelCase   ,buff=0.0       )
														else:
																				cpu_target.target.next_to(cpu_targs[i - 1].target   ,direction=_lowerCAmelCase   ,buff=0.0       )
														cpu_targs.append(_lowerCAmelCase       )

														first_animations.append(rect.animate(run_time=0.5       ).set_stroke(_lowerCAmelCase       )       )
														second_animations.append(MoveToTarget(_lowerCAmelCase   ,run_time=1.5       )       )

								self.play(*_lowerCAmelCase       )
								self.play(*_lowerCAmelCase       )

								self.wait()

