





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

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

def 		SCREAMING_SNAKE_CASE   (       lowerCAmelCase__  :						List[Any]							,	lowerCAmelCase__  :						Union[str, Any]							,	lowerCAmelCase__  :						Optional[int]							,	lowerCAmelCase__  :						Dict							,	lowerCAmelCase__  :						Optional[int])      ->    Any:

     '''simple docstring'''




     for attribute in key.split('.'):
          _lowercase				:   Union[str, Any]	    =				getattr(_UpperCamelCase							,	_UpperCamelCase)

     if weight_type is not None:
          _lowercase				:   Optional[int]	    =				getattr(_UpperCamelCase							,	_UpperCamelCase).shape
     else:
          _lowercase				:   Any	    =				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				:   List[Any]	    =				value
     elif weight_type == "weight_g":
          _lowercase				:   Tuple	    =				value
     elif weight_type == "weight_v":
          _lowercase				:   Union[str, Any]	    =				value
     elif weight_type == "bias":
          _lowercase				:   List[str]	    =				value
     elif weight_type == "running_mean":
          _lowercase				:   Optional[Any]	    =				value
     elif weight_type == "running_var":
          _lowercase				:   Optional[int]	    =				value
     elif weight_type == "num_batches_tracked":
          _lowercase				:   Optional[Any]	    =				value
     elif weight_type == "inv_freq":
          _lowercase				:   str	    =				value
     else:
          _lowercase				:   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   (       lowerCAmelCase__  :						Any							,	lowerCAmelCase__  :						Any							,	lowerCAmelCase__  :						Optional[int])      ->    List[str]:

     '''simple docstring'''




     _lowercase				:   Optional[int]	    =				[]
     _lowercase				:   Any	    =				fairseq_model.state_dict()

     _lowercase				:   Dict	    =				hf_model.wavaveca_conformer.feature_extractor

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

     '''simple docstring'''




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

@torch.no_grad()
def 		SCREAMING_SNAKE_CASE   (       lowerCAmelCase__  :						List[Any]							,	lowerCAmelCase__  :						Optional[int]							,	lowerCAmelCase__  :						int=None							,	lowerCAmelCase__  :						List[str]=None							,	lowerCAmelCase__  :						int=True)      ->    Dict:

     '''simple docstring'''




     if config_path is not None:
          _lowercase				:   Tuple	    =				WavaVecaConformerConfig.from_pretrained(_UpperCamelCase							,	hidden_act='swish')
     else:
          _lowercase				:   Union[str, Any]	    =				WavaVecaConformerConfig()

     if "rope" in checkpoint_path:
          _lowercase				:   List[str]	    =				'rotary'

     if is_finetuned:
          if dict_path:
               _lowercase				:   Optional[int]	    =				Dictionary.load(_UpperCamelCase)

               # important change bos & pad token id since CTC symbol is <pad> and
               # not <s> as in fairseq
               _lowercase				:   List[Any]	    =				target_dict.pad_index
               _lowercase				:   Union[str, Any]	    =				target_dict.bos_index
               _lowercase				:   int	    =				target_dict.eos_index
               _lowercase				:   Any	    =				len(target_dict.symbols)
               _lowercase				:   Union[str, Any]	    =				os.path.join(_UpperCamelCase							,	'vocab.json')
               if not os.path.isdir(_UpperCamelCase):
                    logger.error('--pytorch_dump_folder_path ({}) should be a directory'.format(_UpperCamelCase))
                    return
               os.makedirs(_UpperCamelCase							,	exist_ok=_UpperCamelCase)
               _lowercase				:   Optional[Any]	    =				target_dict.indices

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

          _lowercase				:   Optional[Any]	    =				WavaVecaConformerForCTC(_UpperCamelCase)
     else:
          _lowercase				:   Optional[Any]	    =				WavaVecaConformerForPreTraining(_UpperCamelCase)

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

          _lowercase    ,			_lowercase    ,			_lowercase				:   Optional[int]	    =				fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]							,	task=_UpperCamelCase)

     _lowercase				:   List[Any]	    =				model[0].eval()

     recursively_load_weights(_UpperCamelCase							,	_UpperCamelCase							,	not is_finetuned)

     hf_wavavec.save_pretrained(_UpperCamelCase)


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


import unittest

from transformers import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING, is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
	from PIL import Image
else:
	class     _UpperCamelCase    :
								'''simple docstring'''







								@staticmethod
								def      snake_case							(	*__a       ,  **__a      ):
											pass
@is_pipeline_test
@require_torch
@require_vision
class     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							__UpperCAmelCase		:       str							 =MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     pipeline("visual-question-answering"       ,  model="hf-internal-testing/tiny-vilt-random-vqa"      )
										__lowerCAmelCase							       =     [
										    {
										        "image": Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      ),
										        "question": "How many cats are there?",
										    },
										    {
										        "image": "./tests/fixtures/tests_samples/COCO/000000039769.png",
										        "question": "How many cats are there?",
										    },
										]
										return vqa_pipeline, examples

							def      snake_case							(	self       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     vqa_pipeline(__a       ,  top_k=1      )
										self.assertEqual(
										    __a       ,  [
										        [{"score": ANY(__a      ), "answer": ANY(__a      )}],
										        [{"score": ANY(__a      ), "answer": ANY(__a      )}],
										    ]       ,  )

							@require_torch
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     pipeline("visual-question-answering"       ,  model="hf-internal-testing/tiny-vilt-random-vqa"      )
										__lowerCAmelCase							       =     "./tests/fixtures/tests_samples/COCO/000000039769.png"
										__lowerCAmelCase							       =     "How many cats are there?"

										__lowerCAmelCase							       =     vqa_pipeline(image=__a       ,  question="How many cats are there?"       ,  top_k=2      )
										self.assertEqual(
										    __a       ,  [{"score": ANY(__a      ), "answer": ANY(__a      )}, {"score": ANY(__a      ), "answer": ANY(__a      )}]      )

										__lowerCAmelCase							       =     vqa_pipeline({"image": image, "question": question}       ,  top_k=2      )
										self.assertEqual(
										    __a       ,  [{"score": ANY(__a      ), "answer": ANY(__a      )}, {"score": ANY(__a      ), "answer": ANY(__a      )}]      )

							@slow
							@require_torch
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     pipeline("visual-question-answering"       ,  model="dandelin/vilt-b32-finetuned-vqa"      )
										__lowerCAmelCase							       =     "./tests/fixtures/tests_samples/COCO/000000039769.png"
										__lowerCAmelCase							       =     "How many cats are there?"

										__lowerCAmelCase							       =     vqa_pipeline(image=__a       ,  question=__a       ,  top_k=2      )
										self.assertEqual(
										    nested_simplify(__a       ,  decimals=4      )       ,  [{"score": 0.8_7_9_9, "answer": "2"}, {"score": 0.2_9_6, "answer": "1"}]      )

										__lowerCAmelCase							       =     vqa_pipeline({"image": image, "question": question}       ,  top_k=2      )
										self.assertEqual(
										    nested_simplify(__a       ,  decimals=4      )       ,  [{"score": 0.8_7_9_9, "answer": "2"}, {"score": 0.2_9_6, "answer": "1"}]      )

										__lowerCAmelCase							       =     vqa_pipeline(
										    [{"image": image, "question": question}, {"image": image, "question": question}]       ,  top_k=2      )
										self.assertEqual(
										    nested_simplify(__a       ,  decimals=4      )       ,  [[{"score": 0.8_7_9_9, "answer": "2"}, {"score": 0.2_9_6, "answer": "1"}]] * 2       ,  )



							@require_tf
							@unittest.skip("Visual question answering not implemented in TF"      )
							def      snake_case							(	self      ):
										pass




