





import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
_UpperCAmelCase             =							logging.get_logger(__name__)

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






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:List[str]   ,   SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:Any   ,   SCREAMING_SNAKE_CASE		:Dict   ,   SCREAMING_SNAKE_CASE		:Optional[int]   )      ->   Dict:
 for attribute in key.split("""."""   ):
  __lowerCAmelCase					:	Dict      =       getattr(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )

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

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






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:str   )      ->   Optional[int]:
 __lowerCAmelCase					:	str      =       []
 __lowerCAmelCase					:	Optional[int]      =       fairseq_model.state_dict()

 __lowerCAmelCase					:	Dict      =       hf_model.feature_extractor
 __lowerCAmelCase					:	int      =       hf_model.adapter

 for name, value in fairseq_dict.items():
  __lowerCAmelCase					:	Tuple      =       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					:	List[str]      =       True
  elif any(x in name for x in ["""adaptor""", """w2v_encoder.proj.""", """w2v_proj_ln."""]   ):
   load_adapter(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )
   __lowerCAmelCase					:	Optional[int]      =       True
  else:
   for key, mapped_key in MAPPING.items():
    if key in name or key.split("""w2v_model."""   )[-1] == name.split("""."""   )[0]:
     __lowerCAmelCase					:	List[str]      =       True
     if "*" in mapped_key:
      __lowerCAmelCase					:	int      =       name.split(SCREAMING_SNAKE_CASE   )[0].split("""."""   )[-2]
      __lowerCAmelCase					:	Any      =       mapped_key.replace("""*"""   ,   SCREAMING_SNAKE_CASE   )
     if "weight_g" in name:
      __lowerCAmelCase					:	List[Any]      =       """weight_g"""
     elif "weight_v" in name:
      __lowerCAmelCase					:	int      =       """weight_v"""
     elif "bias" in name:
      __lowerCAmelCase					:	Any      =       """bias"""
     elif "weight" in name:
      __lowerCAmelCase					:	List[Any]      =       """weight"""
     else:
      __lowerCAmelCase					:	Union[str, Any]      =       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}'''   )






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Tuple   ,   SCREAMING_SNAKE_CASE		:Optional[Any]   ,   SCREAMING_SNAKE_CASE		:Union[str, Any]   ,   SCREAMING_SNAKE_CASE		:List[Any]   ,   SCREAMING_SNAKE_CASE		:List[Any]   )      ->   Any:
 __lowerCAmelCase					:	Any      =       full_name.split("""conv_layers."""   )[-1]
 __lowerCAmelCase					:	Optional[Any]      =       name.split("""."""   )
 __lowerCAmelCase					:	List[str]      =       int(items[0]   )
 __lowerCAmelCase					:	Optional[Any]      =       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					:	List[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.'''
   )
   __lowerCAmelCase					:	Optional[int]      =       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					:	Union[str, Any]      =       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					:	Union[str, 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 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:List[str]   ,   SCREAMING_SNAKE_CASE		:List[Any]   ,   SCREAMING_SNAKE_CASE		:int   ,   SCREAMING_SNAKE_CASE		:str   )      ->   Any:
 __lowerCAmelCase					:	Optional[Any]      =       full_name.split("""adaptor."""   )[-1]
 __lowerCAmelCase					:	Optional[Any]      =       name.split("""."""   )

 if items[1].isdigit():
  __lowerCAmelCase					:	Optional[int]      =       int(items[1]   )
 else:
  __lowerCAmelCase					:	Tuple      =       None

 if "adaptor" not in full_name:
  if "proj_ln" in full_name:
   # has to be layer norm
   if "bias" in name:
    assert (
        value.shape == adapter.proj_layer_norm.bias.data.shape
    ), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
    __lowerCAmelCase					:	Union[str, Any]      =       value
    logger.info(F'''Adapter proj layer norm bias was initialized from {full_name}.'''   )
   if "weight" in name:
    assert (
        value.shape == adapter.proj_layer_norm.weight.data.shape
    ), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
    __lowerCAmelCase					:	Optional[int]      =       value
  else:
   # has to be projection layer
   if "bias" in name:
    assert (
        value.shape == adapter.proj.bias.data.shape
    ), F'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
    __lowerCAmelCase					:	Optional[int]      =       value
    logger.info(F'''Adapter proj layer bias was initialized from {full_name}.'''   )
   if "weight" in name:
    assert (
        value.shape == adapter.proj.weight.data.shape
    ), F'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
    __lowerCAmelCase					:	Optional[Any]      =       value
    logger.info(F'''Adapter proj layer weight was initialized from {full_name}.'''   )
 elif isinstance(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ):
  if "bias" in name:
   assert (
       value.shape == adapter.layers[layer_id].conv.bias.data.shape
   ), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
   __lowerCAmelCase					:	List[Any]      =       value
   logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''   )
  elif "weight" in name:
   assert (
       value.shape == adapter.layers[layer_id].conv.weight.data.shape
   ), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
   __lowerCAmelCase					:	int      =       value
   logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''   )
 else:
  unused_weights.append(SCREAMING_SNAKE_CASE   )






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Optional[int]   )      ->   int:
 __lowerCAmelCase ,					__lowerCAmelCase					:	str      =       emb.weight.shape
 __lowerCAmelCase					:	List[str]      =       nn.Linear(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   bias=SCREAMING_SNAKE_CASE   )
 __lowerCAmelCase					:	int      =       emb.weight.data
 return lin_layer






@torch.no_grad()
def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:int   ,   SCREAMING_SNAKE_CASE		:Optional[int]   ,   SCREAMING_SNAKE_CASE		:Dict   ,   SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:List[Any]   ,   SCREAMING_SNAKE_CASE		:Optional[Any]   ,   SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:int   ,   SCREAMING_SNAKE_CASE		:Optional[Any]   ,   SCREAMING_SNAKE_CASE		:Any   ,   )      ->   str:
 __lowerCAmelCase					:	List[str]      =       WavaVecaConfig.from_pretrained(
     SCREAMING_SNAKE_CASE   ,   add_adapter=SCREAMING_SNAKE_CASE   ,   adapter_stride=SCREAMING_SNAKE_CASE   ,   adapter_kernel_size=SCREAMING_SNAKE_CASE   ,   use_auth_token=SCREAMING_SNAKE_CASE   ,   output_hidden_size=SCREAMING_SNAKE_CASE   ,   )
 __lowerCAmelCase					:	str      =       MBartConfig.from_pretrained(SCREAMING_SNAKE_CASE   )

 # load model
 __lowerCAmelCase ,					__lowerCAmelCase ,					__lowerCAmelCase					:	Dict      =       fairseq.checkpoint_utils.load_model_ensemble_and_task(
     [checkpoint_path]   ,   arg_overrides={
         """config_yaml""": config_yaml_path,
         """data""": """/""".join(dict_path.split("""/"""   )[:-1]   ),
         """w2v_path""": checkpoint_path,
         """load_pretrained_decoder_from""": None,
     }   ,   )
 __lowerCAmelCase					:	Tuple      =       model[0].eval()

 # load feature extractor
 __lowerCAmelCase					:	Tuple      =       WavaVecaFeatureExtractor.from_pretrained(SCREAMING_SNAKE_CASE   ,   use_auth_token=SCREAMING_SNAKE_CASE   )

 # set weights for wav2vec2 encoder
 __lowerCAmelCase					:	List[str]      =       WavaVecaModel(SCREAMING_SNAKE_CASE   )

 recursively_load_weights_wavaveca(model.encoder   ,   SCREAMING_SNAKE_CASE   )

 # load decoder weights
 __lowerCAmelCase					:	Dict      =       MBartForCausalLM(SCREAMING_SNAKE_CASE   )
 __lowerCAmelCase ,					__lowerCAmelCase					:	str      =       hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()   ,   strict=SCREAMING_SNAKE_CASE   )
 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					:	Any      =       SpeechEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE   ,   decoder=SCREAMING_SNAKE_CASE   )
 __lowerCAmelCase					:	int      =       False

 __lowerCAmelCase					:	Union[str, Any]      =       MBartaaTokenizer(SCREAMING_SNAKE_CASE   )
 tokenizer.save_pretrained(SCREAMING_SNAKE_CASE   )

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

 __lowerCAmelCase					:	int      =       tokenizer.eos_token_id
 __lowerCAmelCase					:	Optional[int]      =       250_004
 __lowerCAmelCase					:	Optional[int]      =       tokenizer.eos_token_id

 __lowerCAmelCase					:	int      =       SpeechEncoderDecoderConfig.from_dict(SCREAMING_SNAKE_CASE   )

 hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE   )
 feature_extractor.save_pretrained(SCREAMING_SNAKE_CASE   )


if __name__ == "__main__":
      _UpperCAmelCase             =							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('--config_yaml_path', default=None, type=str, help='Path to yaml file of fine-tuned model')
      parser.add_argument(
          '--encoder_config_path',
          default='facebook/wav2vec2-xls-r-1b',
          type=str,
          help='Path to hf encoder wav2vec2 checkpoint config',
      )
      parser.add_argument(
          '--decoder_config_path',
          default='facebook/mbart-large-50-one-to-many-mmt',
          type=str,
          help='Path to hf decoder checkpoint config',
      )
      parser.add_argument('--add_adapter', default=True, type=bool, help='whethere to add model adapter layers')
      parser.add_argument('--adapter_stride', default=2, type=int, help='stride of adapter layers')
      parser.add_argument('--adapter_kernel_size', default=3, type=int, help='kernel size of adapter layers')
      parser.add_argument('--encoder_output_dim', default=1024, type=int, help='encoder output dim')
      parser.add_argument('--start_token_id', default=25_0004, type=int, help='`decoder_start_token_id` of model config')

      _UpperCAmelCase             =							parser.parse_args()
      convert_wavaveca_checkpoint(
          args.checkpoint_path,
          args.pytorch_dump_folder_path,
          args.dict_path,
          args.config_yaml_path,
          encoder_config_path=args.encoder_config_path,
          decoder_config_path=args.decoder_config_path,
          add_adapter=args.add_adapter,
          adapter_kernel_size=args.adapter_kernel_size,
          adapter_stride=args.adapter_stride,
          decoder_start_token_id=args.start_token_id,
          encoder_output_dim=args.encoder_output_dim,
      )






import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding






class        snake_case_    (		__lowercase						):
  A_      							=   ['image_processor', 'tokenizer']
  A_      							=   'ChineseCLIPImageProcessor'
  A_      							=   ('BertTokenizer', 'BertTokenizerFast')






  def __init__(							self			:					Optional[Any]       ,  _snake_case			:					List[Any]=None       ,  _snake_case			:					str=None       ,  **_snake_case			:					int      )->List[str]:






   '''simple docstring'''

   __lowerCAmelCase					:	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."""       ,  _snake_case       ,  )
    __lowerCAmelCase					:	List[str]      =       kwargs.pop("""feature_extractor"""      )

   __lowerCAmelCase					:	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__(_snake_case       ,  _snake_case      )
   __lowerCAmelCase					:	Any      =       self.image_processor






  def __call__(							self			:					Optional[Any]       ,  _snake_case			:					Tuple=None       ,  _snake_case			:					Tuple=None       ,  _snake_case			:					List[str]=None       ,  **_snake_case			:					Any      )->Union[str, Any]:






   '''simple docstring'''

   if text is None and images is None:
    raise ValueError("""You have to specify either text or images. Both cannot be none."""      )

   if text is not None:
    __lowerCAmelCase					:	List[Any]      =       self.tokenizer(_snake_case       ,  return_tensors=_snake_case       ,  **_snake_case      )

   if images is not None:
    __lowerCAmelCase					:	Tuple      =       self.image_processor(_snake_case       ,  return_tensors=_snake_case       ,  **_snake_case      )

   if text is not None and images is not None:
    __lowerCAmelCase					:	Optional[Any]      =       image_features.pixel_values
    return encoding
   elif text is not None:
    return encoding
   else:
    return BatchEncoding(data=dict(**_snake_case      )       ,  tensor_type=_snake_case      )






  def 				UpperCAmelCase__						(							self			:					Optional[int]       ,  *_snake_case			:					Union[str, Any]       ,  **_snake_case			:					int      )->Optional[Any]:






   '''simple docstring'''

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






  def 				UpperCAmelCase__						(							self			:					Dict       ,  *_snake_case			:					Dict       ,  **_snake_case			:					Any      )->Dict:






   '''simple docstring'''

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






  @property
  def 				UpperCAmelCase__						(							self			:					Optional[int]      )->str:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       self.tokenizer.model_input_names
   __lowerCAmelCase					:	Tuple      =       self.image_processor.model_input_names
   return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )






  @property
  def 				UpperCAmelCase__						(							self			:					Dict      )->Union[str, Any]:






   '''simple docstring'''

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