





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

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

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :Any			,							snake_case__  :List[Any]			,							snake_case__  :Optional[Any]			,							snake_case__  :Dict			,							snake_case__  :List[Any]     )      ->						str:
  for attribute in key.split('.'     ):
    _lowercase							= getattr(lowerCAmelCase__			,							lowerCAmelCase__     )

  if weight_type is not None:
    _lowercase							= getattr(lowerCAmelCase__			,							lowerCAmelCase__     ).shape
  else:
    _lowercase							= 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":
    _lowercase							= value
  elif weight_type == "weight_g":
    _lowercase							= value
  elif weight_type == "weight_v":
    _lowercase							= value
  elif weight_type == "bias":
    _lowercase							= value
  elif weight_type == "running_mean":
    _lowercase							= value
  elif weight_type == "running_var":
    _lowercase							= value
  elif weight_type == "num_batches_tracked":
    _lowercase							= value
  elif weight_type == "inv_freq":
    _lowercase							= value
  else:
    _lowercase							= value

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

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :List[Any]			,							snake_case__  :Union[str, Any]			,							snake_case__  :Optional[int]     )      ->						Optional[Any]:
  _lowercase							= []
  _lowercase							= fairseq_model.state_dict()

  _lowercase							= hf_model.wavaveca_conformer.feature_extractor

  for name, value in fairseq_dict.items():
    _lowercase							= False
    if "conv_layers" in name:
      load_conv_layer(
          lowerCAmelCase__			,							lowerCAmelCase__			,							lowerCAmelCase__			,							lowerCAmelCase__			,							hf_model.config.feat_extract_norm == 'group'			,							)
      _lowercase							= True
    else:
      for key, mapped_key in MAPPING.items():
        _lowercase							= '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]:
          _lowercase							= True
          if "*" in mapped_key:
            _lowercase							= name.split(lowerCAmelCase__     )[0].split('.'     )[-2]
            _lowercase							= mapped_key.replace('*'			,							lowerCAmelCase__     )
          if "pos_bias_u" in name:
            _lowercase							= None
          elif "pos_bias_v" in name:
            _lowercase							= None
          elif "weight_g" in name:
            _lowercase							= 'weight_g'
          elif "weight_v" in name:
            _lowercase							= 'weight_v'
          elif "bias" in name:
            _lowercase							= 'bias'
          elif "weight" in name:
            # TODO: don't match quantizer.weight_proj
            _lowercase							= 'weight'
          elif "running_mean" in name:
            _lowercase							= 'running_mean'
          elif "inv_freq" in name:
            _lowercase							= 'inv_freq'
          elif "running_var" in name:
            _lowercase							= 'running_var'
          elif "num_batches_tracked" in name:
            _lowercase							= 'num_batches_tracked'
          else:
            _lowercase							= 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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :List[Any]			,							snake_case__  :Any			,							snake_case__  :List[Any]			,							snake_case__  :Any			,							snake_case__  :List[str]     )      ->						List[Any]:
  _lowercase							= full_name.split('conv_layers.'     )[-1]
  _lowercase							= name.split('.'     )
  _lowercase							= int(items[0]     )
  _lowercase							= 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."""     )
      _lowercase							= 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."""     )
      _lowercase							= 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."""     )
      _lowercase							= 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."""     )
      _lowercase							= 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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :List[str]			,							snake_case__  :Optional[Any]			,							snake_case__  :List[str]=None			,							snake_case__  :List[str]=None			,							snake_case__  :List[str]=True     )      ->						Tuple:
  if config_path is not None:
    _lowercase							= WavaVecaConformerConfig.from_pretrained(lowerCAmelCase__			,							hidden_act='swish'     )
  else:
    _lowercase							= WavaVecaConformerConfig()

  if "rope" in checkpoint_path:
    _lowercase							= 'rotary'

  if is_finetuned:
    if dict_path:
      _lowercase							= Dictionary.load(lowerCAmelCase__     )

      # important change bos & pad token id since CTC symbol is <pad> and
      # not <s> as in fairseq
      _lowercase							= target_dict.pad_index
      _lowercase							= target_dict.bos_index
      _lowercase							= target_dict.eos_index
      _lowercase							= len(target_dict.symbols     )
      _lowercase							= 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__     )
      _lowercase							= target_dict.indices

      # fairseq has the <pad> and <s> switched
      _lowercase							= 0
      _lowercase							= 1
      with open(lowerCAmelCase__			,							'w'			,							encoding='utf-8'     ) as vocab_handle:
        json.dump(lowerCAmelCase__			,							lowerCAmelCase__     )
      _lowercase							= 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__			,							)
      _lowercase							= True if config.feat_extract_norm == 'layer' else False
      _lowercase							= WavaVecaFeatureExtractor(
          feature_size=1			,							sampling_rate=1_6000			,							padding_value=0			,							do_normalize=lowerCAmelCase__			,							return_attention_mask=lowerCAmelCase__			,							)
      _lowercase							= WavaVecaProcessor(feature_extractor=lowerCAmelCase__			,							tokenizer=lowerCAmelCase__     )
      processor.save_pretrained(lowerCAmelCase__     )

    _lowercase							= WavaVecaConformerForCTC(lowerCAmelCase__     )
  else:
    _lowercase							= WavaVecaConformerForPreTraining(lowerCAmelCase__     )

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

    _lowercase						,		_lowercase						,		_lowercase							= fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]			,							task=lowerCAmelCase__     )

  _lowercase							= model[0].eval()

  recursively_load_weights(lowerCAmelCase__			,							lowerCAmelCase__			,							not is_finetuned     )

  hf_wavavec.save_pretrained(lowerCAmelCase__     )


if __name__ == "__main__":
      snake_case        =       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"""
      )
      snake_case        =       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
      )






from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import numpy
      import tensorflow as tf

      from transformers import (
          TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
          BertConfig,
          DPRConfig,
          TFDPRContextEncoder,
          TFDPRQuestionEncoder,
          TFDPRReader,
      )





class 				A_				:







       """simple docstring"""





       def __init__(    self	:	Union[str, Any]						,__A	:	Tuple						,__A	:	Dict=13						,__A	:	Union[str, Any]=7						,__A	:	Optional[Any]=True						,__A	:	Dict=True						,__A	:	List[str]=True						,__A	:	int=True						,__A	:	str=99						,__A	:	Tuple=32						,__A	:	List[Any]=2						,__A	:	Dict=4						,__A	:	List[str]=37						,__A	:	Dict="gelu"						,__A	:	Optional[Any]=0.1						,__A	:	Dict=0.1						,__A	:	Dict=512						,__A	:	Any=16						,__A	:	Tuple=2						,__A	:	List[Any]=0.02						,__A	:	Any=3						,__A	:	Tuple=4						,__A	:	Dict=None						,__A	:	List[Any]=0						,)						->					Optional[Any]:
         _lowercase							= parent
         _lowercase							= batch_size
         _lowercase							= seq_length
         _lowercase							= is_training
         _lowercase							= use_input_mask
         _lowercase							= use_token_type_ids
         _lowercase							= use_labels
         _lowercase							= vocab_size
         _lowercase							= hidden_size
         _lowercase							= num_hidden_layers
         _lowercase							= num_attention_heads
         _lowercase							= intermediate_size
         _lowercase							= hidden_act
         _lowercase							= hidden_dropout_prob
         _lowercase							= attention_probs_dropout_prob
         _lowercase							= max_position_embeddings
         _lowercase							= type_vocab_size
         _lowercase							= type_sequence_label_size
         _lowercase							= initializer_range
         _lowercase							= num_labels
         _lowercase							= num_choices
         _lowercase							= scope
         _lowercase							= projection_dim







       def    __UpperCAmelCase (    self	:	List[Any]  )						->					List[str]:
         _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,self.vocab_size  )

         _lowercase							= None
         if self.use_input_mask:
           # follow test_modeling_tf_ctrl.py
           _lowercase							= random_attention_mask([self.batch_size, self.seq_length]  )

         _lowercase							= None
         if self.use_token_type_ids:
           _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,self.type_vocab_size  )

         _lowercase							= None
         _lowercase							= None
         _lowercase							= None
         if self.use_labels:
           _lowercase							= ids_tensor([self.batch_size]						,self.type_sequence_label_size  )
           _lowercase							= ids_tensor([self.batch_size, self.seq_length]						,self.num_labels  )
           _lowercase							= ids_tensor([self.batch_size]						,self.num_choices  )

         _lowercase							= BertConfig(
             vocab_size=self.vocab_size						,hidden_size=self.hidden_size						,num_hidden_layers=self.num_hidden_layers						,num_attention_heads=self.num_attention_heads						,intermediate_size=self.intermediate_size						,hidden_act=self.hidden_act						,hidden_dropout_prob=self.hidden_dropout_prob						,attention_probs_dropout_prob=self.attention_probs_dropout_prob						,max_position_embeddings=self.max_position_embeddings						,type_vocab_size=self.type_vocab_size						,is_decoder=__A						,initializer_range=self.initializer_range						,)
         _lowercase							= DPRConfig(projection_dim=self.projection_dim						,**config.to_dict()  )

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







       def    __UpperCAmelCase (    self	:	List[Any]						,__A	:	Dict						,__A	:	str						,__A	:	Any						,__A	:	Dict						,__A	:	Tuple						,__A	:	Tuple						,__A	:	int  )						->					List[str]:
         _lowercase							= TFDPRContextEncoder(config=__A  )
         _lowercase							= model(__A						,attention_mask=__A						,token_type_ids=__A  )
         _lowercase							= model(__A						,token_type_ids=__A  )
         _lowercase							= model(__A  )
         self.parent.assertEqual(result.pooler_output.shape						,(self.batch_size, self.projection_dim or self.hidden_size)  )







       def    __UpperCAmelCase (    self	:	List[Any]						,__A	:	Optional[Any]						,__A	:	Tuple						,__A	:	Optional[Any]						,__A	:	Union[str, Any]						,__A	:	Tuple						,__A	:	Tuple						,__A	:	Any  )						->					str:
         _lowercase							= TFDPRQuestionEncoder(config=__A  )
         _lowercase							= model(__A						,attention_mask=__A						,token_type_ids=__A  )
         _lowercase							= model(__A						,token_type_ids=__A  )
         _lowercase							= model(__A  )
         self.parent.assertEqual(result.pooler_output.shape						,(self.batch_size, self.projection_dim or self.hidden_size)  )







       def    __UpperCAmelCase (    self	:	int						,__A	:	List[str]						,__A	:	Optional[int]						,__A	:	List[str]						,__A	:	str						,__A	:	str						,__A	:	int						,__A	:	int  )						->					Optional[Any]:
         _lowercase							= TFDPRReader(config=__A  )
         _lowercase							= model(__A						,attention_mask=__A  )

         self.parent.assertEqual(result.start_logits.shape						,(self.batch_size, self.seq_length)  )
         self.parent.assertEqual(result.end_logits.shape						,(self.batch_size, self.seq_length)  )
         self.parent.assertEqual(result.relevance_logits.shape						,(self.batch_size,)  )







       def    __UpperCAmelCase (    self	:	str  )						->					str:
         _lowercase							= self.prepare_config_and_inputs()
         (
             (
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		(
             _lowercase
         )						,		
         )							= config_and_inputs
         _lowercase							= {'input_ids': input_ids}
         return config, inputs_dict





@require_tf
class 				A_				(      UpperCAmelCase					,	UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Union[str, Any]      =  (
           (
               TFDPRContextEncoder,
               TFDPRQuestionEncoder,
               TFDPRReader,
           )
           if is_tf_available()
           else ()
       )
       SCREAMING_SNAKE_CASE_      : Optional[Any]      =  {'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

       SCREAMING_SNAKE_CASE_      : List[Any]      =  False
       SCREAMING_SNAKE_CASE_      : Dict      =  False
       SCREAMING_SNAKE_CASE_      : int      =  False
       SCREAMING_SNAKE_CASE_      : List[Any]      =  False
       SCREAMING_SNAKE_CASE_      : Optional[int]      =  False







       def    __UpperCAmelCase (    self	:	List[str]  )						->					Optional[int]:
         _lowercase							= TFDPRModelTester(self  )
         _lowercase							= ConfigTester(self						,config_class=__A						,hidden_size=37  )







       def    __UpperCAmelCase (    self	:	Any  )						->					List[str]:
         self.config_tester.run_common_tests()







       def    __UpperCAmelCase (    self	:	Dict  )						->					Optional[int]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_context_encoder(*__A  )







       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					Optional[Any]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_question_encoder(*__A  )







       def    __UpperCAmelCase (    self	:	Dict  )						->					List[Any]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_reader(*__A  )







       @slow
       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					str:
         for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRContextEncoder.from_pretrained(__A  )
           self.assertIsNotNone(__A  )

         for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRContextEncoder.from_pretrained(__A  )
           self.assertIsNotNone(__A  )

         for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRQuestionEncoder.from_pretrained(__A  )
           self.assertIsNotNone(__A  )

         for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= TFDPRReader.from_pretrained(__A  )
           self.assertIsNotNone(__A  )





@require_tf
class 				A_				(      unittest.TestCase	):







       """simple docstring"""





       @slow
       def    __UpperCAmelCase (    self	:	List[Any]  )						->					int:
         _lowercase							= TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'  )

         _lowercase							= tf.constant(
             [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]]  )  # [CLS] hello, is my dog cute? [SEP]
         _lowercase							= model(__A  )[0]  # embedding shape = (1, 768)
         # compare the actual values for a slice.
         _lowercase							= tf.constant(
             [
                 [
                     0.03236253,
                     0.12753335,
                     0.16818509,
                     0.00279786,
                     0.3896933,
                     0.24264945,
                     0.2178971,
                     -0.02335227,
                     -0.08481959,
                     -0.14324117,
                 ]
             ]  )
         self.assertTrue(numpy.allclose(output[:, :10].numpy()						,expected_slice.numpy()						,atol=1e-4  )  )