"""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()
__lowerCamelCase							 =     logging.get_logger(__name__)

__lowerCamelCase							 =     {
    "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",
}
__lowerCamelCase							 =     [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							for attribute in key.split('.' ):
														A__		   =							getattr(UpperCamelCase__				,			UpperCamelCase__ )

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

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






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							A__		   =							[]
							A__		   =							fairseq_model.state_dict()

							A__		   =							hf_model.feature_extractor

							# if encoder has different dim to decoder -> use proj_weight
							A__		   =							None

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

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

							return proj_weight






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





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






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							A__   ,     A__		   =							emb.weight.shape
							A__		   =							nn.Linear(UpperCamelCase__				,			UpperCamelCase__				,			bias=UpperCamelCase__ )
							A__		   =							emb.weight.data
							return lin_layer






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							with open(UpperCamelCase__				,			'r'				,			encoding='utf-8' ) as f:
														A__		   =							f.readlines()
														A__		   =							[line.split(' ' )[0] for line in lines]

							A__		   =							len(UpperCamelCase__ )

							A__		   =							{
							    '<s>': 0,
							    '<pad>': 1,
							    '</s>': 2,
							    '<unk>': 3,
							}

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






@torch.no_grad()
def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			UpperCamelCase__				,			):






							"""simple docstring"""





							A__		   =							WavaVecaConfig.from_pretrained(UpperCamelCase__ )
							A__		   =							SpeechaTextaConfig.from_pretrained(
							    UpperCamelCase__				,			vocab_size=UpperCamelCase__				,			decoder_layers=UpperCamelCase__				,			do_stable_layer_norm=UpperCamelCase__ )

							A__		   =							WavaVecaFeatureExtractor(
							    feature_size=1				,			sampling_rate=16_000				,			padding_value=0				,			do_normalize=UpperCamelCase__				,			return_attention_mask=UpperCamelCase__				,			)

							A__   ,     A__   ,     A__		   =							fairseq.checkpoint_utils.load_model_ensemble_and_task(
							    [checkpoint_path]				,			arg_overrides={'data': '/'.join(dict_path.split('/' )[:-1] )} )
							A__		   =							model[0].eval()

							# set weights for wav2vec2 encoder
							A__		   =							WavaVecaModel(UpperCamelCase__ )
							A__		   =							recursively_load_weights_wavaveca(model.encoder				,			UpperCamelCase__ )

							A__		   =							SpeechaTextaForCausalLM(UpperCamelCase__ )
							A__   ,     A__		   =							hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()				,			strict=UpperCamelCase__ )

							# set output linear layer
							unexpected_keys.remove('embed_out' )
							A__		   =							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}''' )

							A__		   =							SpeechEncoderDecoderModel(encoder=UpperCamelCase__				,			decoder=UpperCamelCase__ )
							A__		   =							False

							# add projection layer
							A__		   =							nn.Parameter(projection_layer.weight )
							A__		   =							nn.Parameter(projection_layer.bias )

							A__		   =							create_vocab_dict(UpperCamelCase__ )

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

							A__		   =							SpeechaTextaTokenizer(os.path.join(UpperCamelCase__				,			'vocab.json' ) )
							tokenizer.save_pretrained(UpperCamelCase__ )

							A__		   =							hf_wavavec.config.to_dict()
							A__		   =							tokenizer.pad_token_id
							A__		   =							tokenizer.bos_token_id
							A__		   =							tokenizer.eos_token_id
							A__		   =							'speech_to_text_2'
							A__		   =							'wav2vec2'

							A__		   =							SpeechEncoderDecoderConfig.from_dict(UpperCamelCase__ )

							hf_wavavec.save_pretrained(UpperCamelCase__ )
							feature_extractor.save_pretrained(UpperCamelCase__ )


if __name__ == "__main__":
				__lowerCamelCase							 =     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=1_02_24, type=int, help="Vocab size of decoder")
				parser.add_argument("--num_decoder_layers", default=7, type=int, help="Number of decoder layers")

				__lowerCamelCase							 =     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"""




import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
				from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
				from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
				from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


__lowerCamelCase							 =     DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

__lowerCamelCase							 =     "main"
# Default branch name
__lowerCamelCase							 =     "f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"
# One particular commit (not the top of `main`)
__lowerCamelCase							 =     "aaaaaaa"
# This commit does not exist, so we should 404.

__lowerCamelCase							 =     "d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"
# Sha-1 of config.json on the top of `main`, for checking purposes
__lowerCamelCase							 =     "4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"






@contextlib.contextmanager
def      UpperCAmelCase    (		):






							"""simple docstring"""





							print('Welcome!' )
							yield
							print('Bye!' )






@contextlib.contextmanager
def      UpperCAmelCase    (		):






							"""simple docstring"""





							print('Bonjour!' )
							yield
							print('Au revoir!' )
class 			UpperCamelCase__(	unittest.TestCase       ):





						def  snake_case__		(				self   )    -> List[str]:
													# If the spec is missing, importlib would not be able to import the module dynamically.
													assert transformers.__spec__ is not None
													assert importlib.util.find_spec('transformers'   ) is not None






class 			UpperCamelCase__(	unittest.TestCase       ):





						@unittest.mock.patch('sys.stdout'						,new_callable=io.StringIO   )
						def  snake_case__		(				self						,__UpperCAmelCase   )    -> Dict:
													with ContextManagers([]   ):
																				print('Transformers are awesome!'   )
													# The print statement adds a new line at the end of the output
													self.assertEqual(mock_stdout.getvalue()						,'Transformers are awesome!\n'   )





						@unittest.mock.patch('sys.stdout'						,new_callable=io.StringIO   )
						def  snake_case__		(				self						,__UpperCAmelCase   )    -> List[str]:
													with ContextManagers([context_en()]   ):
																				print('Transformers are awesome!'   )
													# The output should be wrapped with an English welcome and goodbye
													self.assertEqual(mock_stdout.getvalue()						,'Welcome!\nTransformers are awesome!\nBye!\n'   )





						@unittest.mock.patch('sys.stdout'						,new_callable=io.StringIO   )
						def  snake_case__		(				self						,__UpperCAmelCase   )    -> Any:
													with ContextManagers([context_fr(), context_en()]   ):
																				print('Transformers are awesome!'   )
													# The output should be wrapped with an English and French welcome and goodbye
													self.assertEqual(mock_stdout.getvalue()						,'Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n'   )





						@require_torch
						def  snake_case__		(				self   )    -> Union[str, Any]:
													self.assertEqual(find_labels(__UpperCAmelCase   )						,['labels']   )
													self.assertEqual(find_labels(__UpperCAmelCase   )						,['labels', 'next_sentence_label']   )
													self.assertEqual(find_labels(__UpperCAmelCase   )						,['start_positions', 'end_positions']   )
													class 			UpperCamelCase__(	__A       ):
																			pass






													self.assertEqual(find_labels(__UpperCAmelCase   )						,['labels']   )





						@require_tf
						def  snake_case__		(				self   )    -> str:
													self.assertEqual(find_labels(__UpperCAmelCase   )						,['labels']   )
													self.assertEqual(find_labels(__UpperCAmelCase   )						,['labels', 'next_sentence_label']   )
													self.assertEqual(find_labels(__UpperCAmelCase   )						,['start_positions', 'end_positions']   )
													class 			UpperCamelCase__(	__A       ):
																			pass






													self.assertEqual(find_labels(__UpperCAmelCase   )						,['labels']   )





						@require_flax
						def  snake_case__		(				self   )    -> List[Any]:
													# Flax models don't have labels
													self.assertEqual(find_labels(__UpperCAmelCase   )						,[]   )
													self.assertEqual(find_labels(__UpperCAmelCase   )						,[]   )
													self.assertEqual(find_labels(__UpperCAmelCase   )						,[]   )
													class 			UpperCamelCase__(	__A       ):
																			pass






													self.assertEqual(find_labels(__UpperCAmelCase   )						,[]   )





