





"""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 		__lowercase					(       snake_case_  :  List[str]    ,snake_case_  :  Dict    ,snake_case_  :  Optional[int]    ,snake_case_  :  Dict    ,snake_case_  :  str						)       ->Optional[int]:



    '''simple docstring'''







    for attribute in key.split('''.'''						):
        __A   :						List[str]									=					getattr(snake_case_    ,snake_case_						)

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

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







def 		__lowercase					(       snake_case_  :  List[Any]    ,snake_case_  :  int						)       ->List[str]:



    '''simple docstring'''







    __A   :						Any									=					[]
    __A   :						Optional[int]									=					fairseq_model.state_dict()

    __A   :						Optional[int]									=					hf_model.feature_extractor

    # if encoder has different dim to decoder -> use proj_weight
    __A   :						Union[str, Any]									=					None

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

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

    return proj_weight







def 		__lowercase					(       snake_case_  :  Optional[Any]    ,snake_case_  :  List[Any]    ,snake_case_  :  str    ,snake_case_  :  Optional[Any]    ,snake_case_  :  List[str]						)       ->Dict:



    '''simple docstring'''







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







def 		__lowercase					(       snake_case_  :  List[str]						)       ->int:



    '''simple docstring'''







    __A			,	__A   :						Union[str, Any]									=					emb.weight.shape
    __A   :						Union[str, Any]									=					nn.Linear(snake_case_    ,snake_case_    ,bias=snake_case_						)
    __A   :						Optional[Any]									=					emb.weight.data
    return lin_layer







def 		__lowercase					(       snake_case_  :  Optional[Any]						)       ->List[Any]:



    '''simple docstring'''







    with open(snake_case_    ,'''r'''    ,encoding='''utf-8'''						) as f:
        __A   :						Tuple									=					f.readlines()
        __A   :						Optional[int]									=					[line.split(''' '''						)[0] for line in lines]

    __A   :						Optional[Any]									=					len(snake_case_						)

    __A   :						Optional[Any]									=					{
        '''<s>''': 0,
        '''<pad>''': 1,
        '''</s>''': 2,
        '''<unk>''': 3,
    }

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







@torch.no_grad()
def 		__lowercase					(       snake_case_  :  List[Any]    ,snake_case_  :  int    ,snake_case_  :  str    ,snake_case_  :  Tuple    ,snake_case_  :  Tuple    ,snake_case_  :  Dict    ,snake_case_  :  Any    ,)       ->List[str]:



    '''simple docstring'''







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

    __A   :						Tuple									=					WavaVecaFeatureExtractor(
        feature_size=1    ,sampling_rate=16000    ,padding_value=0    ,do_normalize=snake_case_    ,return_attention_mask=snake_case_    ,)

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

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

    __A   :						Optional[int]									=					SpeechaTextaForCausalLM(snake_case_						)
    __A			,	__A   :						Dict									=					hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()    ,strict=snake_case_						)

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

    __A   :						Tuple									=					SpeechEncoderDecoderModel(encoder=snake_case_    ,decoder=snake_case_						)
    __A   :						Optional[Any]									=					False

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

    __A   :						Tuple									=					create_vocab_dict(snake_case_						)

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

    __A   :						List[Any]									=					SpeechaTextaTokenizer(os.path.join(snake_case_    ,'''vocab.json'''						)						)
    tokenizer.save_pretrained(snake_case_						)

    __A   :						Dict									=					hf_wavavec.config.to_dict()
    __A   :						List[str]									=					tokenizer.pad_token_id
    __A   :						Any									=					tokenizer.bos_token_id
    __A   :						int									=					tokenizer.eos_token_id
    __A   :						Tuple									=					'''speech_to_text_2'''
    __A   :						Tuple									=					'''wav2vec2'''

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

    hf_wavavec.save_pretrained(snake_case_						)
    feature_extractor.save_pretrained(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 typing import TYPE_CHECKING

from ..utils import _LazyModule


a_			=      {
    """config""": [
        """EXTERNAL_DATA_FORMAT_SIZE_LIMIT""",
        """OnnxConfig""",
        """OnnxConfigWithPast""",
        """OnnxSeq2SeqConfigWithPast""",
        """PatchingSpec""",
    ],
    """convert""": ["""export""", """validate_model_outputs"""],
    """features""": ["""FeaturesManager"""],
    """utils""": ["""ParameterFormat""", """compute_serialized_parameters_size"""],
}


if TYPE_CHECKING:
       from .config import (
           EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
           OnnxConfig,
           OnnxConfigWithPast,
           OnnxSeqaSeqConfigWithPast,
           PatchingSpec,
       )
       from .convert import export, validate_model_outputs
       from .features import FeaturesManager
       from .utils import ParameterFormat, compute_serialized_parameters_size

else:
       import sys

       a_			=      _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




