





'''simple docstring'''






import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConformerConfig,
    WavaVecaConformerForCTC,
    WavaVecaConformerForPreTraining,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
A__							:  Tuple    =logging.get_logger(__name__)

A__							:  List[str]    ={
    '''post_extract_proj''': '''feature_projection.projection''',
    '''encoder.pos_conv.0''': '''encoder.pos_conv_embed.conv''',
    '''self_attn.linear_k''': '''encoder.layers.*.self_attn.linear_k''',
    '''self_attn.linear_v''': '''encoder.layers.*.self_attn.linear_v''',
    '''self_attn.linear_q''': '''encoder.layers.*.self_attn.linear_q''',
    '''self_attn.pos_bias_u''': '''encoder.layers.*.self_attn.pos_bias_u''',
    '''self_attn.pos_bias_v''': '''encoder.layers.*.self_attn.pos_bias_v''',
    '''self_attn.linear_out''': '''encoder.layers.*.self_attn.linear_out''',
    '''self_attn.linear_pos''': '''encoder.layers.*.self_attn.linear_pos''',
    '''self_attn.rotary_emb''': '''encoder.embed_positions''',
    '''self_attn_layer_norm''': '''encoder.layers.*.self_attn_layer_norm''',
    '''conv_module.pointwise_conv1''': '''encoder.layers.*.conv_module.pointwise_conv1''',
    '''conv_module.pointwise_conv2''': '''encoder.layers.*.conv_module.pointwise_conv2''',
    '''conv_module.depthwise_conv''': '''encoder.layers.*.conv_module.depthwise_conv''',
    '''conv_module.batch_norm''': '''encoder.layers.*.conv_module.batch_norm''',
    '''conv_module.layer_norm''': '''encoder.layers.*.conv_module.layer_norm''',
    '''ffn1.w_1''': '''encoder.layers.*.ffn1.intermediate_dense''',
    '''ffn1.w_2''': '''encoder.layers.*.ffn1.output_dense''',
    '''ffn1.layer_norm''': '''encoder.layers.*.ffn1_layer_norm''',
    '''ffn2.w_1''': '''encoder.layers.*.ffn2.intermediate_dense''',
    '''ffn2.w_2''': '''encoder.layers.*.ffn2.output_dense''',
    '''ffn2.layer_norm''': '''encoder.layers.*.ffn2_layer_norm''',
    '''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__							:  int    =[
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase    ):





      """simple docstring"""


      for attribute in key.split("""."""    ):
            _lowerCAmelCase					=				getattr(lowerCAmelCase				,      lowerCAmelCase    )

      if weight_type is not None:
            _lowerCAmelCase					=				getattr(lowerCAmelCase				,      lowerCAmelCase    ).shape
      else:
            _lowerCAmelCase					=				hf_pointer.shape

      if hf_shape != value.shape:
            raise ValueError(
                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					=				value
      elif weight_type == "weight_g":
            _lowerCAmelCase					=				value
      elif weight_type == "weight_v":
            _lowerCAmelCase					=				value
      elif weight_type == "bias":
            _lowerCAmelCase					=				value
      elif weight_type == "running_mean":
            _lowerCAmelCase					=				value
      elif weight_type == "running_var":
            _lowerCAmelCase					=				value
      elif weight_type == "num_batches_tracked":
            _lowerCAmelCase					=				value
      elif weight_type == "inv_freq":
            _lowerCAmelCase					=				value
      else:
            _lowerCAmelCase					=				value

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





      """simple docstring"""


      _lowerCAmelCase					=				[]
      _lowerCAmelCase					=				fairseq_model.state_dict()

      _lowerCAmelCase					=				hf_model.wavaveca_conformer.feature_extractor

      for name, value in fairseq_dict.items():
            _lowerCAmelCase					=				False
            if "conv_layers" in name:
                  load_conv_layer(
                      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      hf_model.config.feat_extract_norm == """group"""				,      )
                  _lowerCAmelCase					=				True
            else:
                  for key, mapped_key in MAPPING.items():
                        _lowerCAmelCase					=				"""wav2vec2_conformer.""" + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                        if key in name or key.split("""w2v_model."""    )[-1] == name.split("""."""    )[0]:
                              _lowerCAmelCase					=				True
                              if "*" in mapped_key:
                                    _lowerCAmelCase					=				name.split(lowerCAmelCase    )[0].split("""."""    )[-2]
                                    _lowerCAmelCase					=				mapped_key.replace("""*"""				,      lowerCAmelCase    )
                              if "pos_bias_u" in name:
                                    _lowerCAmelCase					=				None
                              elif "pos_bias_v" in name:
                                    _lowerCAmelCase					=				None
                              elif "weight_g" in name:
                                    _lowerCAmelCase					=				"""weight_g"""
                              elif "weight_v" in name:
                                    _lowerCAmelCase					=				"""weight_v"""
                              elif "bias" in name:
                                    _lowerCAmelCase					=				"""bias"""
                              elif "weight" in name:
                                    # TODO: don't match quantizer.weight_proj
                                    _lowerCAmelCase					=				"""weight"""
                              elif "running_mean" in name:
                                    _lowerCAmelCase					=				"""running_mean"""
                              elif "inv_freq" in name:
                                    _lowerCAmelCase					=				"""inv_freq"""
                              elif "running_var" in name:
                                    _lowerCAmelCase					=				"""running_var"""
                              elif "num_batches_tracked" in name:
                                    _lowerCAmelCase					=				"""num_batches_tracked"""
                              else:
                                    _lowerCAmelCase					=				None
                              set_recursively(lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase    )
                        continue
            if not is_used:
                  unused_weights.append(lowerCAmelCase    )

      logger.warning(f"Unused weights: {unused_weights}"    )
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase    ):





      """simple docstring"""


      _lowerCAmelCase					=				full_name.split("""conv_layers."""    )[-1]
      _lowerCAmelCase					=				name.split("""."""    )
      _lowerCAmelCase					=				int(items[0]    )
      _lowerCAmelCase					=				int(items[1]    )

      if type_id == 0:
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."    )
                  _lowerCAmelCase					=				value
                  logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}."    )
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."    )
                  _lowerCAmelCase					=				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:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found."    )
                  _lowerCAmelCase					=				value
                  logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."    )
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                        raise ValueError(
                            f"{full_name} has size {value.shape}, but"
                            f" {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found."    )
                  _lowerCAmelCase					=				value
                  logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."    )
      else:
            unused_weights.append(lowerCAmelCase    )
@torch.no_grad()
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase=None				,      lowerCAmelCase=None				,      lowerCAmelCase=True    ):





      """simple docstring"""


      if config_path is not None:
            _lowerCAmelCase					=				WavaVecaConformerConfig.from_pretrained(lowerCAmelCase				,      hidden_act="""swish"""    )
      else:
            _lowerCAmelCase					=				WavaVecaConformerConfig()

      if "rope" in checkpoint_path:
            _lowerCAmelCase					=				"""rotary"""

      if is_finetuned:
            if dict_path:
                  _lowerCAmelCase					=				Dictionary.load(lowerCAmelCase    )

                  # important change bos & pad token id since CTC symbol is <pad> and
                  # not <s> as in fairseq
                  _lowerCAmelCase					=				target_dict.pad_index
                  _lowerCAmelCase					=				target_dict.bos_index
                  _lowerCAmelCase					=				target_dict.eos_index
                  _lowerCAmelCase					=				len(target_dict.symbols    )
                  _lowerCAmelCase					=				os.path.join(lowerCAmelCase				,      """vocab.json"""    )
                  if not os.path.isdir(lowerCAmelCase    ):
                        logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(lowerCAmelCase    )    )
                        return
                  os.makedirs(lowerCAmelCase				,      exist_ok=lowerCAmelCase    )
                  _lowerCAmelCase					=				target_dict.indices

                  # fairseq has the <pad> and <s> switched
                  _lowerCAmelCase					=				0
                  _lowerCAmelCase					=				1
                  with open(lowerCAmelCase				,      """w"""				,      encoding="""utf-8"""    ) as vocab_handle:
                        json.dump(lowerCAmelCase				,      lowerCAmelCase    )
                  _lowerCAmelCase					=				WavaVecaCTCTokenizer(
                      lowerCAmelCase				,      unk_token=target_dict.unk_word				,      pad_token=target_dict.pad_word				,      bos_token=target_dict.bos_word				,      eos_token=target_dict.eos_word				,      word_delimiter_token="""|"""				,      do_lower_case=lowerCAmelCase				,      )
                  _lowerCAmelCase					=				True if config.feat_extract_norm == """layer""" else False
                  _lowerCAmelCase					=				WavaVecaFeatureExtractor(
                      feature_size=1				,      sampling_rate=1_60_00				,      padding_value=0				,      do_normalize=lowerCAmelCase				,      return_attention_mask=lowerCAmelCase				,      )
                  _lowerCAmelCase					=				WavaVecaProcessor(feature_extractor=lowerCAmelCase				,      tokenizer=lowerCAmelCase    )
                  processor.save_pretrained(lowerCAmelCase    )

            _lowerCAmelCase					=				WavaVecaConformerForCTC(lowerCAmelCase    )
      else:
            _lowerCAmelCase					=				WavaVecaConformerForPreTraining(lowerCAmelCase    )

      if is_finetuned:
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				fairseq.checkpoint_utils.load_model_ensemble_and_task(
                [checkpoint_path]				,      arg_overrides={"""data""": """/""".join(dict_path.split("""/"""    )[:-1]    )}    )
      else:
            _lowerCAmelCase					=				argparse.Namespace(task="""audio_pretraining"""    )
            _lowerCAmelCase					=				fairseq.tasks.setup_task(lowerCAmelCase    )

            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]				,      task=lowerCAmelCase    )

      _lowerCAmelCase					=				model[0].eval()

      recursively_load_weights(lowerCAmelCase				,      lowerCAmelCase				,      not is_finetuned    )

      hf_wavavec.save_pretrained(lowerCAmelCase    )


if __name__ == "__main__":
      A__							:  List[Any]    =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_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
      parser.add_argument(
          '''--not_finetuned''', action='''store_true''', help='''Whether the model to convert is a fine-tuned model or not'''
      )
      A__							:  Optional[int]    =parser.parse_args()
      convert_wavaveca_conformer_checkpoint(
          args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
      )













'''simple docstring'''






from ..utils import DummyObject, requires_backends




class 					UpperCAmelCase	(     metaclass=snake_case_  ):
      _lowercase:     List[Any]			=  ['''torch''', '''scipy''']




      def __init__(     self     :					Tuple     ,     *__snake_case     :					Dict     ,     **__snake_case     :					List[str]    )							-> str:
            requires_backends(self     ,     ["""torch""", """scipy"""]    )




      @classmethod
      def 							lowercase__       (     cls     :					str     ,     *__snake_case     :					List[str]     ,     **__snake_case     :					str    )							-> Tuple:
            requires_backends(cls     ,     ["""torch""", """scipy"""]    )




      @classmethod
      def 							lowercase__       (     cls     :					Optional[int]     ,     *__snake_case     :					int     ,     **__snake_case     :					Union[str, Any]    )							-> Optional[int]:
            requires_backends(cls     ,     ["""torch""", """scipy"""]    )






