



'''simple docstring'''





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()
A__    :   Optional[int]           =					logging.get_logger(__name__)

A__    :   List[str]           =					{
    '''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__    :   List[str]           =					[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]





def   a_   (		_UpperCAmelCase    :      Tuple ,_UpperCAmelCase    :      List[str] ,_UpperCAmelCase    :      Optional[Any] ,_UpperCAmelCase    :      List[Any] ,_UpperCAmelCase    :      int				)    ->  str:
      for attribute in key.split('.'				):
            __snake_case :   Optional[int]													=    getattr(UpperCamelCase_ ,UpperCamelCase_				)

      if weight_type is not None:
            __snake_case :   Any													=    getattr(UpperCamelCase_ ,UpperCamelCase_				).shape
      else:
            __snake_case :   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":
            __snake_case :   int													=    value
      elif weight_type == "weight_g":
            __snake_case :   Optional[int]													=    value
      elif weight_type == "weight_v":
            __snake_case :   Any													=    value
      elif weight_type == "bias":
            __snake_case :   Optional[int]													=    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   a_   (		_UpperCAmelCase    :      Union[str, Any] ,_UpperCAmelCase    :      str				)    ->  Any:
      __snake_case :   Optional[Any]													=    []
      __snake_case :   Any													=    fairseq_model.state_dict()

      __snake_case :   Dict													=    hf_model.feature_extractor
      __snake_case :   Dict													=    hf_model.adapter

      for name, value in fairseq_dict.items():
            __snake_case :   Union[str, Any]													=    False
            if "conv_layers" in name:
                  load_conv_layer(
                      UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_ ,hf_model.config.feat_extract_norm == 'group' ,)
                  __snake_case :   Any													=    True
            elif any(x in name for x in ['adaptor', 'w2v_encoder.proj.', 'w2v_proj_ln.']				):
                  load_adapter(UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_				)
                  __snake_case :   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 :   str													=    name.split(UpperCamelCase_				)[0].split('.'				)[-2]
                                    __snake_case :   Optional[int]													=    mapped_key.replace('*' ,UpperCamelCase_				)
                              if "weight_g" in name:
                                    __snake_case :   Dict													=    'weight_g'
                              elif "weight_v" in name:
                                    __snake_case :   str													=    'weight_v'
                              elif "bias" in name:
                                    __snake_case :   Optional[int]													=    'bias'
                              elif "weight" in name:
                                    __snake_case :   List[str]													=    'weight'
                              else:
                                    __snake_case :   Union[str, Any]													=    None
                              set_recursively(UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_				)
                        continue
            if not is_used:
                  unused_weights.append(UpperCamelCase_				)

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





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





def   a_   (		_UpperCAmelCase    :      str ,_UpperCAmelCase    :      Tuple ,_UpperCAmelCase    :      int ,_UpperCAmelCase    :      Optional[Any]				)    ->  List[str]:
      __snake_case :   Any													=    full_name.split('adaptor.'				)[-1]
      __snake_case :   str													=    name.split('.'				)

      if items[1].isdigit():
            __snake_case :   List[str]													=    int(items[1]				)
      else:
            __snake_case :   Dict													=    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.'''
                        __snake_case :   Optional[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.'''
                        __snake_case :   Any													=    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.'''
                        __snake_case :   str													=    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.'''
                        __snake_case :   List[Any]													=    value
                        logger.info(f'''Adapter proj layer weight was initialized from {full_name}.'''				)
      elif isinstance(UpperCamelCase_ ,UpperCamelCase_				):
            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.'''
                  __snake_case :   str													=    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.'''
                  __snake_case :   Dict													=    value
                  logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.'''				)
      else:
            unused_weights.append(UpperCamelCase_				)





def   a_   (		_UpperCAmelCase    :      Any				)    ->  Optional[Any]:
      __snake_case							,     __snake_case :   Tuple													=    emb.weight.shape
      __snake_case :   Dict													=    nn.Linear(UpperCamelCase_ ,UpperCamelCase_ ,bias=UpperCamelCase_				)
      __snake_case :   Dict													=    emb.weight.data
      return lin_layer





@torch.no_grad()
def   a_   (		_UpperCAmelCase    :      List[str] ,_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      Tuple ,_UpperCAmelCase    :      Optional[int] ,_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      Union[str, Any] ,_UpperCAmelCase    :      Dict ,_UpperCAmelCase    :      str ,_UpperCAmelCase    :      str ,_UpperCAmelCase    :      List[Any] ,)    ->  Union[str, Any]:
      __snake_case :   int													=    WavaVecaConfig.from_pretrained(
          UpperCamelCase_ ,add_adapter=UpperCamelCase_ ,adapter_stride=UpperCamelCase_ ,adapter_kernel_size=UpperCamelCase_ ,use_auth_token=UpperCamelCase_ ,output_hidden_size=UpperCamelCase_ ,)
      __snake_case :   Tuple													=    MBartConfig.from_pretrained(UpperCamelCase_				)

      # load model
      __snake_case							,     __snake_case							,     __snake_case :   Optional[Any]													=    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,
          } ,)
      __snake_case :   int													=    model[0].eval()

      # load feature extractor
      __snake_case :   int													=    WavaVecaFeatureExtractor.from_pretrained(UpperCamelCase_ ,use_auth_token=UpperCamelCase_				)

      # set weights for wav2vec2 encoder
      __snake_case :   str													=    WavaVecaModel(UpperCamelCase_				)

      recursively_load_weights_wavaveca(model.encoder ,UpperCamelCase_				)

      # load decoder weights
      __snake_case :   Dict													=    MBartForCausalLM(UpperCamelCase_				)
      __snake_case							,     __snake_case :   Union[str, Any]													=    hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() ,strict=UpperCamelCase_				)
      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 :   Optional[Any]													=    SpeechEncoderDecoderModel(encoder=UpperCamelCase_ ,decoder=UpperCamelCase_				)
      __snake_case :   Union[str, Any]													=    False

      __snake_case :   str													=    MBartaaTokenizer(UpperCamelCase_				)
      tokenizer.save_pretrained(UpperCamelCase_				)

      __snake_case :   Dict													=    hf_wavavec.config.to_dict()
      __snake_case :   List[str]													=    tokenizer.pad_token_id
      __snake_case :   Dict													=    tokenizer.bos_token_id
      __snake_case :   List[Any]													=    tokenizer.eos_token_id
      __snake_case :   List[str]													=    'mbart50'
      __snake_case :   Dict													=    'wav2vec2'

      __snake_case :   List[str]													=    tokenizer.eos_token_id
      __snake_case :   Dict													=    25_00_04
      __snake_case :   str													=    tokenizer.eos_token_id

      __snake_case :   List[Any]													=    SpeechEncoderDecoderConfig.from_dict(UpperCamelCase_				)

      hf_wavavec.save_pretrained(UpperCamelCase_				)
      feature_extractor.save_pretrained(UpperCamelCase_				)


if __name__ == "__main__":
    A__    :   Optional[int]           =					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=1_0_2_4, type=int, help='''encoder output dim''')
    parser.add_argument('''--start_token_id''', default=2_5_0_0_0_4, type=int, help='''`decoder_start_token_id` of model config''')

    A__    :   Tuple           =					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,
    )







'''simple docstring'''
__SCREAMING_SNAKE_CASE					:				Optional[int]									=		{
    '''a''': '''AAAAA''',
    '''b''': '''AAAAB''',
    '''c''': '''AAABA''',
    '''d''': '''AAABB''',
    '''e''': '''AABAA''',
    '''f''': '''AABAB''',
    '''g''': '''AABBA''',
    '''h''': '''AABBB''',
    '''i''': '''ABAAA''',
    '''j''': '''BBBAA''',
    '''k''': '''ABAAB''',
    '''l''': '''ABABA''',
    '''m''': '''ABABB''',
    '''n''': '''ABBAA''',
    '''o''': '''ABBAB''',
    '''p''': '''ABBBA''',
    '''q''': '''ABBBB''',
    '''r''': '''BAAAA''',
    '''s''': '''BAAAB''',
    '''t''': '''BAABA''',
    '''u''': '''BAABB''',
    '''v''': '''BBBAB''',
    '''w''': '''BABAA''',
    '''x''': '''BABAB''',
    '''y''': '''BABBA''',
    '''z''': '''BABBB''',
    ''' ''': ''' ''',
}


__SCREAMING_SNAKE_CASE					:				Tuple									=		{value: key for key, value in encode_dict.items()}







def       a_				(	UpperCamelCase_ ):
							A_			    =		""
							for letter in word.lower():
														if letter.isalpha() or letter == " ":
																					encoded += encode_dict[letter]
														else:
																					raise Exception("encode() accepts only letters of the alphabet and spaces" )
							return encoded







def       a_				(	UpperCamelCase_ ):
							if set(UpperCamelCase_ ) - {"A", "B", " "} != set():
														raise Exception("decode() accepts only 'A', 'B' and spaces" )
							A_			    =		""
							for word in coded.split():
														while len(UpperCamelCase_ ) != 0:
																					decoded += decode_dict[word[:5]]
																					A_			    =		word[5:]
														decoded += " "
							return decoded.strip()


if __name__ == "__main__":
							from doctest import testmod

							testmod()





