





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

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





def      _UpperCamelCase  (a__ :Dict ,      a__ :Union[str, Any] ,      a__ :List[str] ,      a__ :Union[str, Any] ,      a__ :Dict							):






							"""simple docstring"""





							for attribute in key.split("""."""							):
														UpperCamelCase__					  =      getattr(a__ ,      a__							)

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

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





def      _UpperCamelCase  (a__ :Tuple ,      a__ :Optional[int]							):






							"""simple docstring"""





							UpperCamelCase__					  =      []
							UpperCamelCase__					  =      fairseq_model.state_dict()

							UpperCamelCase__					  =      hf_model.feature_extractor

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

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

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

							return proj_weight





def      _UpperCamelCase  (a__ :Union[str, Any] ,      a__ :Union[str, Any] ,      a__ :Optional[Any] ,      a__ :int ,      a__ :Optional[int]							):






							"""simple docstring"""





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





def      _UpperCamelCase  (a__ :Optional[Any]							):






							"""simple docstring"""





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





def      _UpperCamelCase  (a__ :Union[str, Any]							):






							"""simple docstring"""





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

							UpperCamelCase__					  =      len(a__							)

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

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





@torch.no_grad()
def      _UpperCamelCase  (a__ :Dict ,      a__ :Optional[int] ,      a__ :List[str] ,      a__ :Any ,      a__ :Optional[Any] ,      a__ :Dict ,      a__ :Optional[int] ,      ):






							"""simple docstring"""





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

							UpperCamelCase__					  =      WavaVecaFeatureExtractor(
							    feature_size=1 ,      sampling_rate=1_6000 ,      padding_value=0 ,      do_normalize=a__ ,      return_attention_mask=a__ ,      )

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

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

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

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

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

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

							UpperCamelCase__					  =      create_vocab_dict(a__							)

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

							UpperCamelCase__					  =      SpeechaTextaTokenizer(os.path.join(a__ ,      """vocab.json"""							)							)
							tokenizer.save_pretrained(a__							)

							UpperCamelCase__					  =      hf_wavavec.config.to_dict()
							UpperCamelCase__					  =      tokenizer.pad_token_id
							UpperCamelCase__					  =      tokenizer.bos_token_id
							UpperCamelCase__					  =      tokenizer.eos_token_id
							UpperCamelCase__					  =      """speech_to_text_2"""
							UpperCamelCase__					  =      """wav2vec2"""

							UpperCamelCase__					  =      SpeechEncoderDecoderConfig.from_dict(a__							)

							hf_wavavec.save_pretrained(a__							)
							feature_extractor.save_pretrained(a__							)


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

					UpperCamelCase__           =			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,
					)











from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput

@dataclass
class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Optional[torch.FloatTensor]      =							None
							snake_case  :   torch.FloatTensor      =							None
							snake_case  :   Optional[Tuple[torch.FloatTensor]]      =							None
							snake_case  :   Optional[Tuple[torch.FloatTensor]]      =							None






class 							__SCREAMING_SNAKE_CASE       (  _a							):






							def __init__(    self	,			__lowerCAmelCase=1	,			__lowerCAmelCase=0	,			__lowerCAmelCase=2	,			__lowerCAmelCase=512	,			__lowerCAmelCase="cls"	,			__lowerCAmelCase=False	,			__lowerCAmelCase=True	,			**__lowerCAmelCase	,			):
														super().__init__(pad_token_id=__lowerCAmelCase	,			bos_token_id=__lowerCAmelCase	,			eos_token_id=__lowerCAmelCase	,			**__lowerCAmelCase	)
														UpperCamelCase__					  =      project_dim
														UpperCamelCase__					  =      pooler_fn
														UpperCamelCase__					  =      learn_encoder
														UpperCamelCase__					  =      use_attention_mask






class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   int      =							[r"""pooler""", r"""logit_scale"""]
							snake_case  :   Tuple      =							[r"""position_ids""", r"""predictions.decoder.bias"""]
							snake_case  :   str      =							"""roberta"""
							snake_case  :   Dict      =							RobertaSeriesConfig






							def __init__(    self	,			__lowerCAmelCase	):
														super().__init__(__lowerCAmelCase	)
														UpperCamelCase__					  =      XLMRobertaModel(__lowerCAmelCase	)
														UpperCamelCase__					  =      nn.Linear(config.hidden_size	,			config.project_dim	)
														UpperCamelCase__					  =      getattr(__lowerCAmelCase	,			"""has_pre_transformation"""	,			__lowerCAmelCase	)
														if self.has_pre_transformation:
																					UpperCamelCase__					  =      nn.Linear(config.hidden_size	,			config.project_dim	)
																					UpperCamelCase__					  =      nn.LayerNorm(config.hidden_size	,			eps=config.layer_norm_eps	)
														self.post_init()







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			__lowerCAmelCase = None	,			):
														UpperCamelCase__					  =      return_dict if return_dict is not None else self.config.use_return_dict

														UpperCamelCase__					  =      self.base_model(
														    input_ids=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	,			position_ids=__lowerCAmelCase	,			head_mask=__lowerCAmelCase	,			inputs_embeds=__lowerCAmelCase	,			encoder_hidden_states=__lowerCAmelCase	,			encoder_attention_mask=__lowerCAmelCase	,			output_attentions=__lowerCAmelCase	,			output_hidden_states=True if self.has_pre_transformation else output_hidden_states	,			return_dict=__lowerCAmelCase	,			)

														if self.has_pre_transformation:
																					UpperCamelCase__					  =      outputs["""hidden_states"""][-2]
																					UpperCamelCase__					  =      self.pre_LN(__lowerCAmelCase	)
																					UpperCamelCase__					  =      self.transformation_pre(__lowerCAmelCase	)

																					return TransformationModelOutput(
																					    projection_state=__lowerCAmelCase	,			last_hidden_state=outputs.last_hidden_state	,			hidden_states=outputs.hidden_states	,			attentions=outputs.attentions	,			)
														else:
																					UpperCamelCase__					  =      self.transformation(outputs.last_hidden_state	)
																					return TransformationModelOutput(
																					    projection_state=__lowerCAmelCase	,			last_hidden_state=outputs.last_hidden_state	,			hidden_states=outputs.hidden_states	,			attentions=outputs.attentions	,			)




