






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

__a      :       Tuple				  =   {
    "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      :       Union[str, Any]				  =   [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str			,				__lowercase   :			List[str]			,				__lowercase   :			Union[str, Any]			,				__lowercase   :			List[str]			,				__lowercase   :			Optional[int]    )						->      int:





    """simple docstring"""

    for attribute in key.split("""."""    ):
        __A         =	getattr(__lowercase			,				__lowercase    )

    if weight_type is not None:
        __A         =	getattr(__lowercase			,				__lowercase    ).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 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Any			,				__lowercase   :			Union[str, Any]    )						->      Any:





    """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(
                __lowercase			,				__lowercase			,				__lowercase			,				__lowercase			,				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(__lowercase    )[0].split("""."""    )[-2]
                        __A         =	mapped_key.replace("""*"""			,				__lowercase    )
                    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(__lowercase			,				__lowercase			,				__lowercase			,				__lowercase			,				__lowercase    )
                continue
        if not is_used:
            unused_weights.append(__lowercase    )

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

    return proj_weight



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[int]			,				__lowercase   :			Optional[Any]			,				__lowercase   :			Optional[Any]			,				__lowercase   :			Dict			,				__lowercase   :			Tuple    )						->      Any:





    """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(__lowercase    )



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[Any]    )						->      str:





    """simple docstring"""

    __A    ,	__A         =	emb.weight.shape
    __A         =	nn.Linear(__lowercase			,				__lowercase			,				bias=__lowercase    )
    __A         =	emb.weight.data
    return lin_layer



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Union[str, Any]    )						->      List[Any]:





    """simple docstring"""

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

    __A         =	len(__lowercase    )

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

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



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





    """simple docstring"""

    __A         =	WavaVecaConfig.from_pretrained(__lowercase    )
    __A         =	SpeechaTextaConfig.from_pretrained(
        __lowercase			,				vocab_size=__lowercase			,				decoder_layers=__lowercase			,				do_stable_layer_norm=__lowercase    )

    __A         =	WavaVecaFeatureExtractor(
        feature_size=1			,				sampling_rate=1_6_0_0_0			,				padding_value=0			,				do_normalize=__lowercase			,				return_attention_mask=__lowercase			,				)

    __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(__lowercase    )
    __A         =	recursively_load_weights_wavaveca(model.encoder			,				__lowercase    )

    __A         =	SpeechaTextaForCausalLM(__lowercase    )
    __A    ,	__A         =	hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()			,				strict=__lowercase    )

    # 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=__lowercase			,				decoder=__lowercase    )
    __A         =	False

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

    __A         =	create_vocab_dict(__lowercase    )

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

    __A         =	SpeechaTextaTokenizer(os.path.join(__lowercase			,				"""vocab.json"""    )    )
    tokenizer.save_pretrained(__lowercase    )

    __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(__lowercase    )

    hf_wavavec.save_pretrained(__lowercase    )
    feature_extractor.save_pretrained(__lowercase    )


if __name__ == "__main__":
  __a      :       Tuple				  =   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      :       Dict				  =   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 Any, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


__a      :       List[Any]				  =   logging.get_logger(__name__)

__a      :       List[str]				  =   {
    "EleutherAI/gpt-neo-1.3B": "https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json",
    # See all GPTNeo models at https://huggingface.co/models?filter=gpt_neo
}


class       __lowercase      (	lowercase_  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				"gpt_neo"
     SCREAMING_SNAKE_CASE   						=				["past_key_values"]
     SCREAMING_SNAKE_CASE   						=				{"num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}




     def __init__(				self       :			Optional[Any] ,	UpperCamelCase_       :			List[Any]=50_257 ,	UpperCamelCase_       :			Union[str, Any]=2_048 ,	UpperCamelCase_       :			List[str]=2_048 ,	UpperCamelCase_       :			Optional[int]=24 ,	UpperCamelCase_       :			Any=[[["global", "local"], 12]] ,	UpperCamelCase_       :			Tuple=16 ,	UpperCamelCase_       :			List[Any]=None ,	UpperCamelCase_       :			Optional[Any]=256 ,	UpperCamelCase_       :			List[str]="gelu_new" ,	UpperCamelCase_       :			str=0.0 ,	UpperCamelCase_       :			Optional[Any]=0.0 ,	UpperCamelCase_       :			str=0.0 ,	UpperCamelCase_       :			List[Any]=0.1 ,	UpperCamelCase_       :			str=1e-5 ,	UpperCamelCase_       :			Optional[Any]=0.02 ,	UpperCamelCase_       :			Union[str, Any]=True ,	UpperCamelCase_       :			Union[str, Any]=50_256 ,	UpperCamelCase_       :			Dict=50_256 ,	**UpperCamelCase_       :			Union[str, Any] ,	):


         """simple docstring"""



         __A         =	vocab_size
         __A         =	max_position_embeddings
         __A         =	hidden_size
         __A         =	num_layers
         __A         =	num_heads
         __A         =	intermediate_size
         __A         =	window_size
         __A         =	activation_function
         __A         =	resid_dropout
         __A         =	embed_dropout
         __A         =	attention_dropout
         __A         =	classifier_dropout
         __A         =	layer_norm_epsilon
         __A         =	initializer_range
         __A         =	use_cache

         __A         =	bos_token_id
         __A         =	eos_token_id

         __A         =	attention_types
         __A         =	self.expand_attention_types_params(UpperCamelCase_				)

         if len(self.attention_layers				) != self.num_layers:
             raise ValueError(
                 """Configuration for convolutional module is incorrect. """
                 """It is required that `len(config.attention_layers)` == `config.num_layers` """
                 F"but is `len(config.attention_layers) = {len(self.attention_layers				)}`, "
                 F"`config.num_layers = {self.num_layers}`. "
                 """`config.attention_layers` is prepared using `config.attention_types`. """
                 """Please verify the value of `config.attention_types` argument."""				)

         super().__init__(bos_token_id=UpperCamelCase_ ,	eos_token_id=UpperCamelCase_ ,	**UpperCamelCase_				)





     @staticmethod
     def 		lowerCAmelCase_     (				UpperCamelCase_       :			str				):


         """simple docstring"""



         __A         =	[]
         for item in attention_types:
             for _ in range(item[1]				):
                 attentions.extend(item[0]				)
         return attentions



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





    """simple docstring"""

    import torch

    __A         =	input.size()
    __A         =	len(__lowercase    )
    __A         =	shape[dimension]

    __A         =	torch.arange(0			,				__lowercase			,				__lowercase    )
    __A         =	torch.div(sizedim - size			,				__lowercase			,				rounding_mode="""floor"""    ) + 1
    __A         =	torch.arange(__lowercase    ) + low_indices[:min_length][:, None]

    __A         =	[slice(__lowercase    )] * rank
    __A         =	indices
    __A         =	input[s]

    __A         =	list(range(0			,				rank + 1    )    )
    perm.append(perm.pop(dimension + 1    )    )

    return sliced.permute(__lowercase    )



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Dict			,				__lowercase   :			Any    )						->      Union[str, Any]:





    """simple docstring"""

    import torch

    __A         =	torch.arange(1			,				__lowercase    )
    __A         =	torch.remainder(__lowercase			,				__lowercase    )
    __A         =	remainders == 0
    __A         =	candidates[divisor_indices]
    __A         =	torch.max(__lowercase    )
    return largest_divisor, torch.div(__lowercase			,				__lowercase			,				rounding_mode="""floor"""    )






class       __lowercase      (	lowercase_  ):







     '''simple docstring'''




     @property
     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A         =	OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}				)
         if self.use_past:
             self.fill_with_past_key_values_(UpperCamelCase_ ,	direction="""inputs"""				)
             __A         =	{0: """batch""", 1: """past_sequence + sequence"""}
         else:
             __A         =	{0: """batch""", 1: """sequence"""}

         return common_inputs




     @property
     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         return self._config.num_heads




     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			PreTrainedTokenizer ,	UpperCamelCase_       :			int = -1 ,	UpperCamelCase_       :			int = -1 ,	UpperCamelCase_       :			bool = False ,	UpperCamelCase_       :			Optional[TensorType] = None ,	):


         """simple docstring"""



         __A         =	super(UpperCamelCase_ ,	self				).generate_dummy_inputs(
             UpperCamelCase_ ,	batch_size=UpperCamelCase_ ,	seq_length=UpperCamelCase_ ,	is_pair=UpperCamelCase_ ,	framework=UpperCamelCase_				)

         # We need to order the input in the way they appears in the forward()
         __A         =	OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}				)

         # Need to add the past_keys
         if self.use_past:
             if not is_torch_available():
                 raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""				)
             else:
                 import torch

                 __A    ,	__A         =	common_inputs["""input_ids"""].shape
                 # Not using the same length for past_key_values
                 __A         =	seqlen + 2
                 __A         =	(
                     batch,
                     self.num_attention_heads,
                     past_key_values_length,
                     self._config.hidden_size // self.num_attention_heads,
                 )
                 __A         =	[
                     (torch.zeros(UpperCamelCase_				), torch.zeros(UpperCamelCase_				)) for _ in range(self.num_layers				)
                 ]

         __A         =	common_inputs["""attention_mask"""]
         if self.use_past:
             __A         =	ordered_inputs["""attention_mask"""].dtype
             __A         =	torch.cat(
                 [ordered_inputs["""attention_mask"""], torch.ones(UpperCamelCase_ ,	UpperCamelCase_ ,	dtype=UpperCamelCase_				)] ,	dim=1				)

         return ordered_inputs





     @property
     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         return 13




