



"""simple docstring"""

import argparse
import json
import os

import fairseq
import torch
from torch import nn

from transformers import (
    SpeechaTextaConfig,
    SpeechaTextaForCausalLM,
    SpeechaTextaTokenizer,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
a_			   =					logging.get_logger(__name__)

a_			   =					{
    '''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''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}
a_			   =					[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  Any				,    SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  List[str]				,    SCREAMING_SNAKE_CASE__ :  str					):



					"""simple docstring"""





					for attribute in key.split("""."""					):
										snake_case_					:		Optional[int] 					=					getattr(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__					)

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

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







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  int					):



					"""simple docstring"""





					snake_case_					:		int 					=					[]
					snake_case_					:		List[Any] 					=					fairseq_model.state_dict()

					snake_case_					:		Union[str, Any] 					=					hf_model.feature_extractor

					# if encoder has different dim to decoder -> use proj_weight
					snake_case_					:		Any 					=					None

					for name, value in fairseq_dict.items():
										snake_case_					:		List[str] 					=					False
										if "conv_layers" in name:
															load_conv_layer(
															    SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    hf_model.config.feat_extract_norm == """group"""				,    )
															snake_case_					:		Dict 					=					True
										elif name.split("""."""					)[0] == "proj":
															snake_case_					:		Any 					=					fairseq_model.proj
															snake_case_					:		str 					=					True
										else:
															for key, mapped_key in MAPPING.items():
																				if key in name or key.split("""w2v_model."""					)[-1] == name.split("""."""					)[0]:
																									snake_case_					:		Tuple 					=					True
																									if "*" in mapped_key:
																														snake_case_					:		List[Any] 					=					name.split(SCREAMING_SNAKE_CASE__					)[0].split("""."""					)[-2]
																														snake_case_					:		List[Any] 					=					mapped_key.replace("""*"""				,    SCREAMING_SNAKE_CASE__					)
																									if "weight_g" in name:
																														snake_case_					:		Optional[Any] 					=					"""weight_g"""
																									elif "weight_v" in name:
																														snake_case_					:		Any 					=					"""weight_v"""
																									elif "bias" in name:
																														snake_case_					:		Tuple 					=					"""bias"""
																									elif "weight" in name:
																														snake_case_					:		int 					=					"""weight"""
																									else:
																														snake_case_					:		Optional[Any] 					=					None
																									set_recursively(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__					)
																				continue
										if not is_used:
															unused_weights.append(SCREAMING_SNAKE_CASE__					)

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

					return proj_weight







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[int]				,    SCREAMING_SNAKE_CASE__ :  Dict				,    SCREAMING_SNAKE_CASE__ :  Tuple				,    SCREAMING_SNAKE_CASE__ :  Union[str, Any]				,    SCREAMING_SNAKE_CASE__ :  int					):



					"""simple docstring"""





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







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  str					):



					"""simple docstring"""





					snake_case_					:		List[str] 					=					emb.weight.shape
					snake_case_					:		Optional[int] 					=					nn.Linear(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__				,    bias=SCREAMING_SNAKE_CASE__					)
					snake_case_					:		Optional[int] 					=					emb.weight.data
					return lin_layer







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Dict					):



					"""simple docstring"""





					with open(SCREAMING_SNAKE_CASE__				,    """r"""				,    encoding="""utf-8"""					) as f:
										snake_case_					:		List[Any] 					=					f.readlines()
										snake_case_					:		List[str] 					=					[line.split(""" """					)[0] for line in lines]

					snake_case_					:		Optional[Any] 					=					len(SCREAMING_SNAKE_CASE__					)

					snake_case_					:		str 					=					{
					    """<s>""": 0,
					    """<pad>""": 1,
					    """</s>""": 2,
					    """<unk>""": 3,
					}

					vocab_dict.update(dict(zip(SCREAMING_SNAKE_CASE__				,    range(4				,    num_words + 4					)					)					)					)
					return vocab_dict







@torch.no_grad()
def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  int				,    SCREAMING_SNAKE_CASE__ :  Any				,    SCREAMING_SNAKE_CASE__ :  List[Any]				,    SCREAMING_SNAKE_CASE__ :  Union[str, Any]				,    SCREAMING_SNAKE_CASE__ :  Any				,    SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    SCREAMING_SNAKE_CASE__ :  Optional[Any]				,    ):



					"""simple docstring"""





					snake_case_					:		List[Any] 					=					WavaVecaConfig.from_pretrained(SCREAMING_SNAKE_CASE__					)
					snake_case_					:		Dict 					=					SpeechaTextaConfig.from_pretrained(
					    SCREAMING_SNAKE_CASE__				,    vocab_size=SCREAMING_SNAKE_CASE__				,    decoder_layers=SCREAMING_SNAKE_CASE__				,    do_stable_layer_norm=SCREAMING_SNAKE_CASE__					)

					snake_case_					:		Optional[int] 					=					WavaVecaFeatureExtractor(
					    feature_size=1				,    sampling_rate=1_6_0_0_0				,    padding_value=0				,    do_normalize=SCREAMING_SNAKE_CASE__				,    return_attention_mask=SCREAMING_SNAKE_CASE__				,    )

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

					# set weights for wav2vec2 encoder
					snake_case_					:		Any 					=					WavaVecaModel(SCREAMING_SNAKE_CASE__					)
					snake_case_					:		Optional[Any] 					=					recursively_load_weights_wavaveca(model.encoder				,    SCREAMING_SNAKE_CASE__					)

					snake_case_					:		Tuple 					=					SpeechaTextaForCausalLM(SCREAMING_SNAKE_CASE__					)
					snake_case_					:		List[str] 					=					hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()				,    strict=SCREAMING_SNAKE_CASE__					)

					# set output linear layer
					unexpected_keys.remove("""embed_out"""					)
					snake_case_					:		Any 					=					nn.Parameter(model.decoder.embed_out.detach()					)

					# layer norm is init to identity matrix so leaving it is fine
					logger.warning(f'The following keys are missing when loading the decoder weights: {missing_keys}'					)
					logger.warning(f'The following keys are unexpected when loading the decoder weights: {unexpected_keys}'					)

					snake_case_					:		Dict 					=					SpeechEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE__				,    decoder=SCREAMING_SNAKE_CASE__					)
					snake_case_					:		List[Any] 					=					False

					# add projection layer
					snake_case_					:		Dict 					=					nn.Parameter(projection_layer.weight					)
					snake_case_					:		List[Any] 					=					nn.Parameter(projection_layer.bias					)

					snake_case_					:		List[Any] 					=					create_vocab_dict(SCREAMING_SNAKE_CASE__					)

					with open(os.path.join(SCREAMING_SNAKE_CASE__				,    """vocab.json"""					)				,    """w"""					) as fp:
										json.dump(SCREAMING_SNAKE_CASE__				,    SCREAMING_SNAKE_CASE__					)

					snake_case_					:		Optional[int] 					=					SpeechaTextaTokenizer(os.path.join(SCREAMING_SNAKE_CASE__				,    """vocab.json"""					)					)
					tokenizer.save_pretrained(SCREAMING_SNAKE_CASE__					)

					snake_case_					:		int 					=					hf_wavavec.config.to_dict()
					snake_case_					:		Optional[Any] 					=					tokenizer.pad_token_id
					snake_case_					:		Optional[Any] 					=					tokenizer.bos_token_id
					snake_case_					:		Union[str, Any] 					=					tokenizer.eos_token_id
					snake_case_					:		List[str] 					=					"""speech_to_text_2"""
					snake_case_					:		str 					=					"""wav2vec2"""

					snake_case_					:		Optional[Any] 					=					SpeechEncoderDecoderConfig.from_dict(SCREAMING_SNAKE_CASE__					)

					hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE__					)
					feature_extractor.save_pretrained(SCREAMING_SNAKE_CASE__					)


if __name__ == "__main__":
			a_			   =					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(
			    '''--encoder_config_path''',
			    default='''facebook/wav2vec2-large-lv60''',
			    type=str,
			    help='''Path to hf encoder wav2vec2 checkpoint config''',
			)
			parser.add_argument(
			    '''--decoder_config_path''',
			    default='''facebook/s2t-small-mustc-en-fr-st''',
			    type=str,
			    help='''Path to hf decoder s2t checkpoint config''',
			)
			parser.add_argument('''--vocab_size''', default=10224, type=int, help='''Vocab size of decoder''')
			parser.add_argument('''--num_decoder_layers''', default=7, type=int, help='''Number of decoder layers''')

			a_			   =					parser.parse_args()
			convert_wavaveca_checkpoint(
			    args.checkpoint_path,
			    args.pytorch_dump_folder_path,
			    args.dict_path,
			    encoder_config_path=args.encoder_config_path,
			    decoder_config_path=args.decoder_config_path,
			    vocab_size=args.vocab_size,
			    num_decoder_layers=args.num_decoder_layers,
			)









"""simple docstring"""

from manim import *







class    __lowercase     (					_UpperCAmelCase):







						"""simple docstring"""







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					Rectangle(height=0.5   ,		width=0.5       )
											snake_case_					:		str 					=					Rectangle(height=0.46   ,		width=0.46       ).set_stroke(width=0       )

											snake_case_					:		Optional[Any] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		str 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		str 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		Any 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[str] 					=					VGroup(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[Any] 					=					Text("""CPU"""   ,		font_size=24       )
											snake_case_					:		Tuple 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											cpu.move_to([-2.5, -0.5, 0]       )
											self.add(lowercase__       )

											snake_case_					:		List[Any] 					=					[mem.copy() for i in range(4       )]
											snake_case_					:		Tuple 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[str] 					=					Text("""GPU"""   ,		font_size=24       )
											snake_case_					:		Any 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											gpu.move_to([-1, -1, 0]       )
											self.add(lowercase__       )

											snake_case_					:		Optional[Any] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		List[Any] 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )

											snake_case_					:		Dict 					=					Text("""Model"""   ,		font_size=24       )
											snake_case_					:		int 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											model.move_to([3, -1.0, 0]       )
											self.add(lowercase__       )

											snake_case_					:		Dict 					=					[]
											for i, rect in enumerate(lowercase__       ):
																rect.set_stroke(lowercase__       )
																# target = fill.copy().set_fill(YELLOW, opacity=0.7)
																# target.move_to(rect)
																# self.add(target)

																snake_case_					:		List[str] 					=					Rectangle(height=0.46 / 4   ,		width=0.46 / 3       ).set_stroke(width=0.0       ).set_fill(lowercase__   ,		opacity=0.7       )

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

											snake_case_					:		List[str] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		List[str] 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )

											snake_case_					:		str 					=					Text("""Loaded Checkpoint"""   ,		font_size=24       )
											snake_case_					:		Any 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		aligned_edge=lowercase__   ,		buff=0.4       )
											checkpoint.move_to([3, 0.5, 0]       )

											snake_case_					:		Optional[Any] 					=					Square(side_length=2.2       )
											key.move_to([-5, 2, 0]       )

											snake_case_					:		Union[str, Any] 					=					MarkupText(
											    f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'   ,		font_size=18   ,		)

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

											self.add(lowercase__   ,		lowercase__       )

											snake_case_					:		List[Any] 					=					MarkupText(
											    f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'   ,		font_size=18   ,		)

											blue_text.next_to(lowercase__   ,		DOWN * 2.4   ,		aligned_edge=key_text.get_left()       )

											snake_case_					:		List[Any] 					=					MarkupText(
											    f'Next, a <i><span fgcolor="{BLUE}">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor="{BLUE}">single shard</span>.'   ,		font_size=24   ,		)
											step_a.move_to([2, 2, 0]       )
											self.play(Write(lowercase__       )   ,		Write(lowercase__       )       )

											self.play(Write(lowercase__   ,		run_time=1       )   ,		Create(lowercase__   ,		run_time=1       )       )

											snake_case_					:		Optional[int] 					=					[]
											snake_case_					:		List[str] 					=					[]
											for i, rect in enumerate(lowercase__       ):
																snake_case_					:		Optional[Any] 					=					fill.copy().set_fill(lowercase__   ,		opacity=0.7       )
																target.move_to(lowercase__       )
																first_animations.append(GrowFromCenter(lowercase__   ,		run_time=1       )       )

																snake_case_					:		List[Any] 					=					target.copy()
																cpu_target.generate_target()
																if i < 5:
																					cpu_target.target.move_to(cpu_left_col_base[i + 1]       )
																else:
																					cpu_target.target.move_to(cpu_right_col_base[i - 5]       )
																second_animations.append(MoveToTarget(lowercase__   ,		run_time=1.5       )       )

											self.play(*lowercase__       )
											self.play(*lowercase__       )
											self.wait()




