



'''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__    :   List[str]           =					logging.get_logger(__name__)

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





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

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

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





def   a_   (		_UpperCAmelCase    :      int ,_UpperCAmelCase    :      int				)    ->  Any:
      __snake_case :   List[Any]													=    []
      __snake_case :   Union[str, Any]													=    fairseq_model.state_dict()

      __snake_case :   Tuple													=    hf_model.feature_extractor

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

      for name, value in fairseq_dict.items():
            __snake_case :   Tuple													=    False
            if "conv_layers" in name:
                  load_conv_layer(
                      _UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,hf_model.config.feat_extract_norm == 'group' ,)
                  __snake_case :   Optional[int]													=    True
            elif name.split('.'				)[0] == "proj":
                  __snake_case :   Optional[Any]													=    fairseq_model.proj
                  __snake_case :   str													=    True
            else:
                  for key, mapped_key in MAPPING.items():
                        if key in name or key.split('w2v_model.'				)[-1] == name.split('.'				)[0]:
                              __snake_case :   str													=    True
                              if "*" in mapped_key:
                                    __snake_case :   List[Any]													=    name.split(_UpperCAmelCase				)[0].split('.'				)[-2]
                                    __snake_case :   Optional[int]													=    mapped_key.replace('*' ,_UpperCAmelCase				)
                              if "weight_g" in name:
                                    __snake_case :   int													=    'weight_g'
                              elif "weight_v" in name:
                                    __snake_case :   Optional[Any]													=    'weight_v'
                              elif "bias" in name:
                                    __snake_case :   str													=    '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}'''				)

      return proj_weight





def   a_   (		_UpperCAmelCase    :      Tuple ,_UpperCAmelCase    :      Optional[int] ,_UpperCAmelCase    :      Optional[Any] ,_UpperCAmelCase    :      Dict ,_UpperCAmelCase    :      int				)    ->  Union[str, Any]:
      __snake_case :   List[str]													=    full_name.split('conv_layers.'				)[-1]
      __snake_case :   List[str]													=    name.split('.'				)
      __snake_case :   Optional[int]													=    int(items[0]				)
      __snake_case :   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.'''
                  )
                  __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 :   Tuple													=    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 :   Optional[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.'''
                  )
                  __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(_UpperCAmelCase				)





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





def   a_   (		_UpperCAmelCase    :      Optional[int]				)    ->  str:
      with open(_UpperCAmelCase ,'r' ,encoding='utf-8'				) as f:
            __snake_case :   str													=    f.readlines()
            __snake_case :   Tuple													=    [line.split(' '				)[0] for line in lines]

      __snake_case :   List[Any]													=    len(_UpperCAmelCase				)

      __snake_case :   List[str]													=    {
          '<s>': 0,
          '<pad>': 1,
          '</s>': 2,
          '<unk>': 3,
      }

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





@torch.no_grad()
def   a_   (		_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      List[Any] ,_UpperCAmelCase    :      int ,_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      List[str] ,_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      Optional[Any] ,)    ->  Tuple:
      __snake_case :   List[Any]													=    WavaVecaConfig.from_pretrained(_UpperCAmelCase				)
      __snake_case :   Optional[int]													=    SpeechaTextaConfig.from_pretrained(
          _UpperCAmelCase ,vocab_size=_UpperCAmelCase ,decoder_layers=_UpperCAmelCase ,do_stable_layer_norm=_UpperCAmelCase				)

      __snake_case :   List[str]													=    WavaVecaFeatureExtractor(
          feature_size=1 ,sampling_rate=1_60_00 ,padding_value=0 ,do_normalize=_UpperCAmelCase ,return_attention_mask=_UpperCAmelCase ,)

      __snake_case							,     __snake_case							,     __snake_case :   Union[str, Any]													=    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(_UpperCAmelCase				)
      __snake_case :   Any													=    recursively_load_weights_wavaveca(model.encoder ,_UpperCAmelCase				)

      __snake_case :   str													=    SpeechaTextaForCausalLM(_UpperCAmelCase				)
      __snake_case							,     __snake_case :   Optional[int]													=    hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() ,strict=_UpperCAmelCase				)

      # 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=_UpperCAmelCase ,decoder=_UpperCAmelCase				)
      __snake_case :   Union[str, Any]													=    False

      # add projection layer
      __snake_case :   Tuple													=    nn.Parameter(projection_layer.weight				)
      __snake_case :   Tuple													=    nn.Parameter(projection_layer.bias				)

      __snake_case :   List[str]													=    create_vocab_dict(_UpperCAmelCase				)

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

      __snake_case :   Dict													=    SpeechaTextaTokenizer(os.path.join(_UpperCAmelCase ,'vocab.json'				)				)
      tokenizer.save_pretrained(_UpperCAmelCase				)

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

      __snake_case :   Union[str, Any]													=    SpeechEncoderDecoderConfig.from_dict(_UpperCAmelCase				)

      hf_wavavec.save_pretrained(_UpperCAmelCase				)
      feature_extractor.save_pretrained(_UpperCAmelCase				)


if __name__ == "__main__":
    A__    :   List[str]           =					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_0_2_2_4, type=int, help='''Vocab size of decoder''')
    parser.add_argument('''--num_decoder_layers''', default=7, type=int, help='''Number of decoder layers''')

    A__    :   Any           =					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 copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


A__    :   Union[str, Any]           =					logging.get_logger(__name__)



class  snake_case__						(    SCREAMING_SNAKE_CASE_				):
 A__				= '''upernet'''
 def __init__(   self						:							Optional[int]      ,     __a						:							Dict=None      ,     __a						:							List[str]=512      ,     __a						:							int=0.0_2      ,     __a						:							Optional[Any]=[1, 2, 3, 6]      ,     __a						:							Union[str, Any]=True      ,     __a						:							Any=0.4      ,     __a						:							int=384      ,     __a						:							int=256      ,     __a						:							str=1      ,     __a						:							List[Any]=False      ,     __a						:							Dict=255      ,     **__a						:							Dict      ,     )      ->   Optional[int]:







       '''simple docstring'''
       super().__init__(**__a  )

       if backbone_config is None:
             logger.info('`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'  )
             __snake_case :   Optional[int]													=    CONFIG_MAPPING['resnet'](out_features=['stage1', 'stage2', 'stage3', 'stage4']  )
       elif isinstance(__a      ,     __a  ):
             __snake_case :   List[Any]													=    backbone_config.get('model_type'  )
             __snake_case :   List[str]													=    CONFIG_MAPPING[backbone_model_type]
             __snake_case :   int													=    config_class.from_dict(__a  )

       __snake_case :   Any													=    backbone_config
       __snake_case :   Tuple													=    hidden_size
       __snake_case :   str													=    initializer_range
       __snake_case :   Dict													=    pool_scales
       __snake_case :   Optional[Any]													=    use_auxiliary_head
       __snake_case :   Optional[Any]													=    auxiliary_loss_weight
       __snake_case :   Optional[int]													=    auxiliary_in_channels
       __snake_case :   Dict													=    auxiliary_channels
       __snake_case :   Any													=    auxiliary_num_convs
       __snake_case :   Optional[Any]													=    auxiliary_concat_input
       __snake_case :   List[str]													=    loss_ignore_index





 def A_   (   self						:							List[Any]  )      ->   Union[str, Any]:







       '''simple docstring'''
       __snake_case :   Union[str, Any]													=    copy.deepcopy(self.__dict__  )
       __snake_case :   str													=    self.backbone_config.to_dict()
       __snake_case :   Optional[int]													=    self.__class__.model_type
       return output



