




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()
SCREAMING_SNAKE_CASE__       :      Union[str, Any]  =   logging.get_logger(__name__)

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

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> Optional[Any]:
       for attribute in key.split("."    ):
              lowerCamelCase   :      Dict			 =					getattr(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )

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

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

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> Dict:
       lowerCamelCase   :      Optional[int]			 =					[]
       lowerCamelCase   :      Optional[Any]			 =					fairseq_model.state_dict()

       lowerCamelCase   :      Dict			 =					hf_model.feature_extractor

       # if encoder has different dim to decoder -> use proj_weight
       lowerCamelCase   :      str			 =					None

       for name, value in fairseq_dict.items():
              lowerCamelCase   :      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" ,)
                     lowerCamelCase   :      Union[str, Any]			 =					True
              elif name.split("."    )[0] == "proj":
                     lowerCamelCase   :      List[str]			 =					fairseq_model.proj
                     lowerCamelCase   :      Optional[Any]			 =					True
              else:
                     for key, mapped_key in MAPPING.items():
                            if key in name or key.split("w2v_model."    )[-1] == name.split("."    )[0]:
                                   lowerCamelCase   :      int			 =					True
                                   if "*" in mapped_key:
                                          lowerCamelCase   :      Union[str, Any]			 =					name.split(_SCREAMING_SNAKE_CASE    )[0].split("."    )[-2]
                                          lowerCamelCase   :      List[str]			 =					mapped_key.replace("*" ,_SCREAMING_SNAKE_CASE    )
                                   if "weight_g" in name:
                                          lowerCamelCase   :      Optional[int]			 =					"weight_g"
                                   elif "weight_v" in name:
                                          lowerCamelCase   :      int			 =					"weight_v"
                                   elif "bias" in name:
                                          lowerCamelCase   :      int			 =					"bias"
                                   elif "weight" in name:
                                          lowerCamelCase   :      List[str]			 =					"weight"
                                   else:
                                          lowerCamelCase   :      Dict			 =					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      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> str:
       lowerCamelCase   :      List[Any]			 =					full_name.split("conv_layers."    )[-1]
       lowerCamelCase   :      Any			 =					name.split("."    )
       lowerCamelCase   :      int			 =					int(items[0]    )
       lowerCamelCase   :      Tuple			 =					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   :      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.'''
                     )
                     lowerCamelCase   :      Dict			 =					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   :      Dict			 =					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   :      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      A      (      _SCREAMING_SNAKE_CASE    )	-> str:
       lowerCamelCase ,      lowerCamelCase   :      List[str]			 =					emb.weight.shape
       lowerCamelCase   :      int			 =					nn.Linear(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,bias=_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Optional[int]			 =					emb.weight.data
       return lin_layer

def      A      (      _SCREAMING_SNAKE_CASE    )	-> Dict:
       with open(_SCREAMING_SNAKE_CASE ,"r" ,encoding="utf-8"    ) as f:
              lowerCamelCase   :      Union[str, Any]			 =					f.readlines()
              lowerCamelCase   :      Optional[int]			 =					[line.split(" "    )[0] for line in lines]

       lowerCamelCase   :      Dict			 =					len(_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      List[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      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,)	-> Optional[int]:
       lowerCamelCase   :      Dict			 =					WavaVecaConfig.from_pretrained(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Tuple			 =					SpeechaTextaConfig.from_pretrained(
           _SCREAMING_SNAKE_CASE ,vocab_size=_SCREAMING_SNAKE_CASE ,decoder_layers=_SCREAMING_SNAKE_CASE ,do_stable_layer_norm=_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      str			 =					WavaVecaFeatureExtractor(
           feature_size=1 ,sampling_rate=1_6000 ,padding_value=0 ,do_normalize=_SCREAMING_SNAKE_CASE ,return_attention_mask=_SCREAMING_SNAKE_CASE ,)

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

       # set weights for wav2vec2 encoder
       lowerCamelCase   :      str			 =					WavaVecaModel(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Dict			 =					recursively_load_weights_wavaveca(model.encoder ,_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Any			 =					SpeechaTextaForCausalLM(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase ,      lowerCamelCase   :      List[Any]			 =					hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() ,strict=_SCREAMING_SNAKE_CASE    )

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

       lowerCamelCase   :      str			 =					SpeechEncoderDecoderModel(encoder=_SCREAMING_SNAKE_CASE ,decoder=_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Optional[int]			 =					False

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

       lowerCamelCase   :      Dict			 =					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    )

       lowerCamelCase   :      Dict			 =					SpeechaTextaTokenizer(os.path.join(_SCREAMING_SNAKE_CASE ,"vocab.json"    )    )
       tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Optional[Any]			 =					hf_wavavec.config.to_dict()
       lowerCamelCase   :      str			 =					tokenizer.pad_token_id
       lowerCamelCase   :      List[Any]			 =					tokenizer.bos_token_id
       lowerCamelCase   :      List[Any]			 =					tokenizer.eos_token_id
       lowerCamelCase   :      str			 =					"speech_to_text_2"
       lowerCamelCase   :      Dict			 =					"wav2vec2"

       lowerCamelCase   :      Dict			 =					SpeechEncoderDecoderConfig.from_dict(_SCREAMING_SNAKE_CASE    )

       hf_wavavec.save_pretrained(_SCREAMING_SNAKE_CASE    )
       feature_extractor.save_pretrained(_SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
       SCREAMING_SNAKE_CASE__       :      Optional[Any]  =   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')

       SCREAMING_SNAKE_CASE__       :      Tuple  =   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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase           =					{'''configuration_xlnet''': ['''XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''XLNetConfig''']}

try:
       if not is_sentencepiece_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase           =					['''XLNetTokenizer''']

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase           =					['''XLNetTokenizerFast''']

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase           =					[
           '''XLNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''XLNetForMultipleChoice''',
           '''XLNetForQuestionAnswering''',
           '''XLNetForQuestionAnsweringSimple''',
           '''XLNetForSequenceClassification''',
           '''XLNetForTokenClassification''',
           '''XLNetLMHeadModel''',
           '''XLNetModel''',
           '''XLNetPreTrainedModel''',
           '''load_tf_weights_in_xlnet''',
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase           =					[
           '''TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''TFXLNetForMultipleChoice''',
           '''TFXLNetForQuestionAnsweringSimple''',
           '''TFXLNetForSequenceClassification''',
           '''TFXLNetForTokenClassification''',
           '''TFXLNetLMHeadModel''',
           '''TFXLNetMainLayer''',
           '''TFXLNetModel''',
           '''TFXLNetPreTrainedModel''',
       ]


if TYPE_CHECKING:
       from .configuration_xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfig

       try:
              if not is_sentencepiece_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_xlnet import XLNetTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_xlnet_fast import XLNetTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_xlnet import (
                  XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                  XLNetForMultipleChoice,
                  XLNetForQuestionAnswering,
                  XLNetForQuestionAnsweringSimple,
                  XLNetForSequenceClassification,
                  XLNetForTokenClassification,
                  XLNetLMHeadModel,
                  XLNetModel,
                  XLNetPreTrainedModel,
                  load_tf_weights_in_xlnet,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_xlnet import (
                  TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFXLNetForMultipleChoice,
                  TFXLNetForQuestionAnsweringSimple,
                  TFXLNetForSequenceClassification,
                  TFXLNetForTokenClassification,
                  TFXLNetLMHeadModel,
                  TFXLNetMainLayer,
                  TFXLNetModel,
                  TFXLNetPreTrainedModel,
              )

else:
       import sys

       lowerCamelCase           =					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
