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

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


def 			UpperCamelCase					(    __lowerCamelCase     : str				,					__lowerCamelCase     : Dict				,					__lowerCamelCase     : Any				,					__lowerCamelCase     : str				,					__lowerCamelCase     : str    ):

       for attribute in key.split("."    ):
              snake_case       :   Optional[Any]							=      getattr(snake_case__				,					snake_case__    )

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

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


def 			UpperCamelCase					(    __lowerCamelCase     : str				,					__lowerCamelCase     : List[str]    ):
       snake_case       :   List[Any]							=      []
       snake_case       :   Optional[Any]							=      fairseq_model.state_dict()

       snake_case       :   str							=      hf_model.feature_extractor

       # if encoder has different dim to decoder -> use proj_weight
       snake_case       :   Optional[Any]							=      None

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


def 			UpperCamelCase					(    __lowerCamelCase     : List[Any]    ):

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

       snake_case       :   str							=      len(snake_case__    )

       snake_case       :   Tuple							=      {
           """<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 			UpperCamelCase					(    __lowerCamelCase     : int				,					__lowerCamelCase     : List[str]				,					__lowerCamelCase     : int				,					__lowerCamelCase     : Dict				,					__lowerCamelCase     : List[Any]				,					__lowerCamelCase     : str				,					__lowerCamelCase     : Union[str, Any]				,					):
       snake_case       :   Optional[int]							=      WavaVecaConfig.from_pretrained(snake_case__    )
       snake_case       :   List[str]							=      SpeechaTextaConfig.from_pretrained(
           snake_case__				,					vocab_size=snake_case__				,					decoder_layers=snake_case__				,					do_stable_layer_norm=snake_case__    )

       snake_case       :   Dict							=      WavaVecaFeatureExtractor(
           feature_size=1				,					sampling_rate=16000				,					padding_value=0				,					do_normalize=snake_case__				,					return_attention_mask=snake_case__				,					)

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

       # set weights for wav2vec2 encoder
       snake_case       :   Any							=      WavaVecaModel(snake_case__    )
       snake_case       :   Optional[Any]							=      recursively_load_weights_wavaveca(model.encoder				,					snake_case__    )

       snake_case       :   Optional[Any]							=      SpeechaTextaForCausalLM(snake_case__    )
       snake_case       :   List[str]							=      hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()				,					strict=snake_case__    )

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

       snake_case       :   Any							=      SpeechEncoderDecoderModel(encoder=snake_case__				,					decoder=snake_case__    )
       snake_case       :   Any							=      False

       # add projection layer
       snake_case       :   int							=      nn.Parameter(projection_layer.weight    )
       snake_case       :   Any							=      nn.Parameter(projection_layer.bias    )

       snake_case       :   Any							=      create_vocab_dict(snake_case__    )

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

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

       snake_case       :   str							=      hf_wavavec.config.to_dict()
       snake_case       :   List[str]							=      tokenizer.pad_token_id
       snake_case       :   Union[str, Any]							=      tokenizer.bos_token_id
       snake_case       :   Union[str, Any]							=      tokenizer.eos_token_id
       snake_case       :   Optional[Any]							=      """speech_to_text_2"""
       snake_case       :   Optional[int]							=      """wav2vec2"""

       snake_case       :   Tuple							=      SpeechEncoderDecoderConfig.from_dict(snake_case__    )

       hf_wavavec.save_pretrained(snake_case__    )
       feature_extractor.save_pretrained(snake_case__    )


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

   __lowerCamelCase							=       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"""


import re
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin






class 					lowercase(     __a ):







 '''simple docstring'''





 lowercase__ 					=					["image_processor", "tokenizer"]
 lowercase__ 					=					"AutoImageProcessor"
 lowercase__ 					=					"AutoTokenizer"


 def __init__(				self:		List[str],      a_:		List[str]=None,      a_:		Tuple=None,      **a_:		Tuple		):




   '''simple docstring'''



   _snake_case							:		str								=  None
   if "feature_extractor" in kwargs:
     warnings.warn(
         """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
         """ instead.""",      a_,      )
     _snake_case							:		str								=  kwargs.pop("""feature_extractor"""		)

   _snake_case							:		Union[str, Any]								=  image_processor if image_processor is not None else feature_extractor
   if image_processor is None:
     raise ValueError("""You need to specify an `image_processor`."""		)
   if tokenizer is None:
     raise ValueError("""You need to specify a `tokenizer`."""		)

   super().__init__(a_,      a_		)
   _snake_case							:		Dict								=  self.image_processor
   _snake_case							:		Any								=  False


 def __call__(				self:		Any,      *a_:		Any,      **a_:		Tuple		):




   '''simple docstring'''



   if self._in_target_context_manager:
     return self.current_processor(*a_,      **a_		)

   _snake_case							:		Dict								=  kwargs.pop("""images""",      a_		)
   _snake_case							:		Optional[Any]								=  kwargs.pop("""text""",      a_		)
   if len(a_		) > 0:
     _snake_case							:		Optional[int]								=  args[0]
     _snake_case							:		Tuple								=  args[1:]

   if images is None and text is None:
     raise ValueError("""You need to specify either an `images` or `text` input to process."""		)

   if images is not None:
     _snake_case							:		Tuple								=  self.image_processor(a_,      *a_,      **a_		)
   if text is not None:
     _snake_case							:		Tuple								=  self.tokenizer(a_,      **a_		)

   if text is None:
     return inputs
   elif images is None:
     return encodings
   else:
     _snake_case							:		List[str]								=  encodings["""input_ids"""]
     return inputs


 def  UpperCamelCase_   (				self:		Optional[int],      *a_:		Tuple,      **a_:		List[str]		):




   '''simple docstring'''



   return self.tokenizer.batch_decode(*a_,      **a_		)


 def  UpperCamelCase_   (				self:		int,      *a_:		List[str],      **a_:		int		):




   '''simple docstring'''



   return self.tokenizer.decode(*a_,      **a_		)


 @contextmanager
 def  UpperCamelCase_   (				self:		Dict		):




   '''simple docstring'''



   warnings.warn(
       """`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your """
       """labels by using the argument `text` of the regular `__call__` method (either in the same call as """
       """your images inputs, or in a separate call."""		)
   _snake_case							:		Any								=  True
   _snake_case							:		Optional[int]								=  self.tokenizer
   yield
   _snake_case							:		int								=  self.image_processor
   _snake_case							:		Optional[int]								=  False


 def  UpperCamelCase_   (				self:		Dict,      a_:		Optional[Any],      a_:		str=False,      a_:		Optional[Any]=None		):




   '''simple docstring'''



   if added_vocab is None:
     _snake_case							:		Dict								=  self.tokenizer.get_added_vocab()

   _snake_case							:		str								=  {}

   while tokens:
     _snake_case							:		Union[str, Any]								=  re.search(r"""<s_(.*?)>""",      a_,      re.IGNORECASE		)
     if start_token is None:
       break
     _snake_case							:		List[Any]								=  start_token.group(1		)
     _snake_case							:		str								=  re.search(rf"</s_{key}>",      a_,      re.IGNORECASE		)
     _snake_case							:		Dict								=  start_token.group()
     if end_token is None:
       _snake_case							:		List[Any]								=  tokens.replace(a_,      """"""		)
     else:
       _snake_case							:		List[str]								=  end_token.group()
       _snake_case							:		str								=  re.escape(a_		)
       _snake_case							:		str								=  re.escape(a_		)
       _snake_case							:		Union[str, Any]								=  re.search(f"{start_token_escaped}(.*?){end_token_escaped}",      a_,      re.IGNORECASE		)
       if content is not None:
         _snake_case							:		int								=  content.group(1		).strip()
         if r"<s_" in content and r"</s_" in content:  # non-leaf node
           _snake_case							:		List[Any]								=  self.tokenajson(a_,      is_inner_value=a_,      added_vocab=a_		)
           if value:
             if len(a_		) == 1:
               _snake_case							:		List[str]								=  value[0]
             _snake_case							:		List[str]								=  value
         else:  # leaf nodes
           _snake_case							:		Tuple								=  []
           for leaf in content.split(r"""<sep/>"""		):
             _snake_case							:		Tuple								=  leaf.strip()
             if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
               _snake_case							:		int								=  leaf[1:-2]  # for categorical special tokens
             output[key].append(a_		)
           if len(output[key]		) == 1:
             _snake_case							:		int								=  output[key][0]

       _snake_case							:		Any								=  tokens[tokens.find(a_		) + len(a_		) :].strip()
       if tokens[:6] == r"<sep/>":  # non-leaf nodes
         return [output] + self.tokenajson(tokens[6:],      is_inner_value=a_,      added_vocab=a_		)

   if len(a_		):
     return [output] if is_inner_value else output
   else:
     return [] if is_inner_value else {"text_sequence": tokens}


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




   '''simple docstring'''



   warnings.warn(
       """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.""",      a_,      )
   return self.image_processor_class


 @property
 def  UpperCamelCase_   (				self:		Tuple		):




   '''simple docstring'''



   warnings.warn(
       """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.""",      a_,      )
   return self.image_processor


