



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

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



def 		lowerCAmelCase__    (					UpperCamelCase_				:       Optional[Any] ,       UpperCamelCase_				:       int ,       UpperCamelCase_				:       Union[str, Any] ,       UpperCamelCase_				:       int ,       UpperCamelCase_				:       Tuple	)->  int:
    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 		lowerCAmelCase__    (					UpperCamelCase_				:       Union[str, Any] ,       UpperCamelCase_				:       List[str]	)->  Optional[Any]:
    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 		lowerCAmelCase__    (					UpperCamelCase_				:       Optional[int] ,       UpperCamelCase_				:       int ,       UpperCamelCase_				:       str ,       UpperCamelCase_				:       Optional[int] ,       UpperCamelCase_				:       str	)->  List[str]:
    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 		lowerCAmelCase__    (					UpperCamelCase_				:       int	)->  Dict:
    A__ ,	A__		 =					emb.weight.shape
    A__		 =					nn.Linear(UpperCamelCase_ ,       UpperCamelCase_ ,       bias=UpperCamelCase_	)
    A__		 =					emb.weight.data
    return lin_layer



def 		lowerCAmelCase__    (					UpperCamelCase_				:       Dict	)->  List[Any]:
    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 		lowerCAmelCase__    (					UpperCamelCase_				:       List[str] ,       UpperCamelCase_				:       Optional[Any] ,       UpperCamelCase_				:       Tuple ,       UpperCamelCase_				:       Any ,       UpperCamelCase_				:       Union[str, Any] ,       UpperCamelCase_				:       int ,       UpperCamelCase_				:       List[Any] ,       )->  str:
    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=1_6_0_0_0 ,       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__":
   _lowercase						  =		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")

   _lowercase						  =		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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_lowercase						  =		logging.get_logger(__name__)

_lowercase						  =		{
    "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/config.json",
    "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/config.json",
    "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/config.json",
    "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/config.json",
    "bert-base-multilingual-uncased": "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/config.json",
    "bert-base-multilingual-cased": "https://huggingface.co/bert-base-multilingual-cased/resolve/main/config.json",
    "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/config.json",
    "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/config.json",
    "bert-large-uncased-whole-word-masking": (
        "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/config.json"
    ),
    "bert-large-cased-whole-word-masking": (
        "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/config.json"
    ),
    "bert-large-uncased-whole-word-masking-finetuned-squad": (
        "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/config.json"
    ),
    "bert-large-cased-whole-word-masking-finetuned-squad": (
        "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/config.json"
    ),
    "bert-base-cased-finetuned-mrpc": "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/config.json",
    "bert-base-german-dbmdz-cased": "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/config.json",
    "bert-base-german-dbmdz-uncased": "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese": "https://huggingface.co/cl-tohoku/bert-base-japanese/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese-whole-word-masking": (
        "https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking/resolve/main/config.json"
    ),
    "cl-tohoku/bert-base-japanese-char": (
        "https://huggingface.co/cl-tohoku/bert-base-japanese-char/resolve/main/config.json"
    ),
    "cl-tohoku/bert-base-japanese-char-whole-word-masking": (
        "https://huggingface.co/cl-tohoku/bert-base-japanese-char-whole-word-masking/resolve/main/config.json"
    ),
    "TurkuNLP/bert-base-finnish-cased-v1": (
        "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/config.json"
    ),
    "TurkuNLP/bert-base-finnish-uncased-v1": (
        "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/config.json"
    ),
    "wietsedv/bert-base-dutch-cased": "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/config.json",
    # See all BERT models at https://huggingface.co/models?filter=bert
}



class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''bert'''




     def __init__(				self ,						a__=3_0_5_2_2 ,						a__=7_6_8 ,						a__=1_2 ,						a__=1_2 ,						a__=3_0_7_2 ,						a__="gelu" ,						a__=0.1 ,						a__=0.1 ,						a__=5_1_2 ,						a__=2 ,						a__=0.0_2 ,						a__=1e-12 ,						a__=0 ,						a__="absolute" ,						a__=True ,						a__=None ,						**a__ ,						):
         super().__init__(pad_token_id=a__ ,						**a__)

         A__		 =					vocab_size
         A__		 =					hidden_size
         A__		 =					num_hidden_layers
         A__		 =					num_attention_heads
         A__		 =					hidden_act
         A__		 =					intermediate_size
         A__		 =					hidden_dropout_prob
         A__		 =					attention_probs_dropout_prob
         A__		 =					max_position_embeddings
         A__		 =					type_vocab_size
         A__		 =					initializer_range
         A__		 =					layer_norm_eps
         A__		 =					position_embedding_type
         A__		 =					use_cache
         A__		 =					classifier_dropout







class     _UpperCAmelCase			(       A__     ):




     @property
     def 		snake_case_		(				self):
         if self.task == "multiple-choice":
             A__		 =					{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
         else:
             A__		 =					{0: '''batch''', 1: '''sequence'''}
         return OrderedDict(
             [
                 ('''input_ids''', dynamic_axis),
                 ('''attention_mask''', dynamic_axis),
                 ('''token_type_ids''', dynamic_axis),
             ])
