





import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    HubertConfig,
    HubertForCTC,
    HubertModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
UpperCamelCase__  			=logging.get_logger(__name__)

UpperCamelCase__  			={
    '''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''',
    '''w2v_encoder.proj''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}






def      lowerCamelCase__      (__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase						):





     for attribute in key.split("."						):
          _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				getattr(UpperCAmelCase_,							UpperCAmelCase_						)

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

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






def      lowerCamelCase__      (__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase						):
     _SCREAMING_SNAKE_CASE							:						Any         =				[]
     _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				fairseq_model.state_dict()

     _SCREAMING_SNAKE_CASE							:						List[Any]         =				hf_model.hubert.feature_extractor if is_finetuned else hf_model.feature_extractor

     for name, value in fairseq_dict.items():
          _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				False
          if "conv_layers" in name:
               load_conv_layer(
                   UpperCAmelCase_,							UpperCAmelCase_,							UpperCAmelCase_,							UpperCAmelCase_,							hf_model.config.feat_extract_norm == "group",							)
               _SCREAMING_SNAKE_CASE							:						Any         =				True
          else:
               for key, mapped_key in MAPPING.items():
                    _SCREAMING_SNAKE_CASE							:						Optional[int]         =				'''hubert.''' + mapped_key if (is_finetuned and mapped_key != '''lm_head''') else mapped_key

                    if key in name or (key.split("w2v_model."						)[-1] == name.split("."						)[0] and not is_finetuned):
                         _SCREAMING_SNAKE_CASE							:						int         =				True
                         if "*" in mapped_key:
                              _SCREAMING_SNAKE_CASE							:						int         =				name.split(UpperCAmelCase_						)[0].split("."						)[-2]
                              _SCREAMING_SNAKE_CASE							:						Optional[int]         =				mapped_key.replace("*",							UpperCAmelCase_						)
                         if "weight_g" in name:
                              _SCREAMING_SNAKE_CASE							:						Union[str, Any]         =				'''weight_g'''
                         elif "weight_v" in name:
                              _SCREAMING_SNAKE_CASE							:						Union[str, Any]         =				'''weight_v'''
                         elif "weight" in name:
                              _SCREAMING_SNAKE_CASE							:						List[str]         =				'''weight'''
                         elif "bias" in name:
                              _SCREAMING_SNAKE_CASE							:						Union[str, Any]         =				'''bias'''
                         else:
                              _SCREAMING_SNAKE_CASE							:						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      lowerCamelCase__      (__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase						):
     _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				full_name.split("conv_layers."						)[-1]
     _SCREAMING_SNAKE_CASE							:						Optional[int]         =				name.split("."						)
     _SCREAMING_SNAKE_CASE							:						List[Any]         =				int(items[0]						)
     _SCREAMING_SNAKE_CASE							:						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."""
               )
               _SCREAMING_SNAKE_CASE							:						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."""
               )
               _SCREAMING_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."
               )
               _SCREAMING_SNAKE_CASE							:						Dict         =				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."""
               )
               _SCREAMING_SNAKE_CASE							:						Tuple         =				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      lowerCamelCase__      (__lowerCamelCase,							__lowerCamelCase,							__lowerCamelCase=None,							__lowerCamelCase=None,							__lowerCamelCase=True						):





     if config_path is not None:
          _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				HubertConfig.from_pretrained(UpperCAmelCase_						)
     else:
          _SCREAMING_SNAKE_CASE							:						Dict         =				HubertConfig()

     if is_finetuned:
          if dict_path:
               _SCREAMING_SNAKE_CASE							:						Union[str, Any]         =				Dictionary.load(UpperCAmelCase_						)

               # important change bos & pad token id since CTC symbol is <pad> and
               # not <s> as in fairseq
               _SCREAMING_SNAKE_CASE							:						Tuple         =				target_dict.pad_index
               _SCREAMING_SNAKE_CASE							:						int         =				target_dict.bos_index
               _SCREAMING_SNAKE_CASE							:						Dict         =				target_dict.eos_index
               _SCREAMING_SNAKE_CASE							:						Dict         =				len(target_dict.symbols						)
               _SCREAMING_SNAKE_CASE							:						Dict         =				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_						)
               with open(UpperCAmelCase_,							"w",							encoding="utf-8"						) as vocab_handle:
                    json.dump(target_dict.indices,							UpperCAmelCase_						)
               _SCREAMING_SNAKE_CASE							:						Union[str, 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_,							)
               _SCREAMING_SNAKE_CASE							:						List[str]         =				True if config.feat_extract_norm == '''layer''' else False
               _SCREAMING_SNAKE_CASE							:						List[Any]         =				WavaVecaFeatureExtractor(
                   feature_size=1,							sampling_rate=16000,							padding_value=0,							do_normalize=UpperCAmelCase_,							return_attention_mask=UpperCAmelCase_,							)
               _SCREAMING_SNAKE_CASE							:						List[str]         =				WavaVecaProcessor(feature_extractor=UpperCAmelCase_,							tokenizer=UpperCAmelCase_						)
               processor.save_pretrained(UpperCAmelCase_						)

          _SCREAMING_SNAKE_CASE							:						List[Any]         =				HubertForCTC(UpperCAmelCase_						)
     else:
          _SCREAMING_SNAKE_CASE							:						int         =				HubertModel(UpperCAmelCase_						)

     if is_finetuned:
          _SCREAMING_SNAKE_CASE							:						int         =				fairseq.checkpoint_utils.load_model_ensemble_and_task(
              [checkpoint_path],							arg_overrides={"data": "/".join(dict_path.split("/"						)[:-1]						)}						)
     else:
          _SCREAMING_SNAKE_CASE							:						Any         =				fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]						)

     _SCREAMING_SNAKE_CASE							:						Any         =				model[0].eval()

     recursively_load_weights(UpperCAmelCase_,							UpperCAmelCase_,							UpperCAmelCase_						)

     hf_wavavec.save_pretrained(UpperCAmelCase_						)


if __name__ == "__main__":
       UpperCamelCase__  			=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'
       )
       UpperCamelCase__  			=parser.parse_args()
       convert_hubert_checkpoint(
           args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
       )



import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.distilbert.modeling_flax_distilbert import (
         FlaxDistilBertForMaskedLM,
         FlaxDistilBertForMultipleChoice,
         FlaxDistilBertForQuestionAnswering,
         FlaxDistilBertForSequenceClassification,
         FlaxDistilBertForTokenClassification,
         FlaxDistilBertModel,
     )







class 						_snake_case    ( unittest.TestCase				):




 def __init__(						self   ,    _lowerCamelCase   ,    _lowerCamelCase=13   ,    _lowerCamelCase=7   ,    _lowerCamelCase=True   ,    _lowerCamelCase=True   ,    _lowerCamelCase=True   ,    _lowerCamelCase=True   ,    _lowerCamelCase=99   ,    _lowerCamelCase=32   ,    _lowerCamelCase=5   ,    _lowerCamelCase=4   ,    _lowerCamelCase=37   ,    _lowerCamelCase="gelu"   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=512   ,    _lowerCamelCase=16   ,    _lowerCamelCase=2   ,    _lowerCamelCase=0.02   ,    _lowerCamelCase=4   ,    ):
       a		:Optional[Any]									=			parent
       a		:str									=			batch_size
       a		:Tuple									=			seq_length
       a		:List[Any]									=			is_training
       a		:Optional[int]									=			use_attention_mask
       a		:List[str]									=			use_token_type_ids
       a		:str									=			use_labels
       a		:Optional[Any]									=			vocab_size
       a		:Optional[int]									=			hidden_size
       a		:Tuple									=			num_hidden_layers
       a		:Union[str, Any]									=			num_attention_heads
       a		:int									=			intermediate_size
       a		:int									=			hidden_act
       a		:int									=			hidden_dropout_prob
       a		:Union[str, Any]									=			attention_probs_dropout_prob
       a		:str									=			max_position_embeddings
       a		:Dict									=			type_vocab_size
       a		:str									=			type_sequence_label_size
       a		:List[str]									=			initializer_range
       a		:Optional[Any]									=			num_choices




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Tuple									=			ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size       )

       a		:Any									=			None
       if self.use_attention_mask:
             a		:Any									=			random_attention_mask([self.batch_size, self.seq_length]       )

       a		:Any									=			DistilBertConfig(
           vocab_size=self.vocab_size   ,    dim=self.hidden_size   ,    n_layers=self.num_hidden_layers   ,    n_heads=self.num_attention_heads   ,    hidden_dim=self.intermediate_size   ,    hidden_act=self.hidden_act   ,    dropout=self.hidden_dropout_prob   ,    attention_dropout=self.attention_probs_dropout_prob   ,    max_position_embeddings=self.max_position_embeddings   ,    initializer_range=self.initializer_range   ,    tie_weights_=_lowerCamelCase   ,    )

       return config, input_ids, attention_mask




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:List[str]									=			self.prepare_config_and_inputs()
       a				,							a				,							a		:str									=			config_and_inputs
       a		:List[Any]									=			{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict







@require_flax
class 						_snake_case    ( _snake_case		,       unittest.TestCase				):
 SCREAMING_SNAKE_CASE__ =			(
     (
         FlaxDistilBertModel,
         FlaxDistilBertForMaskedLM,
         FlaxDistilBertForMultipleChoice,
         FlaxDistilBertForQuestionAnswering,
         FlaxDistilBertForSequenceClassification,
         FlaxDistilBertForTokenClassification,
         FlaxDistilBertForQuestionAnswering,
     )
     if is_flax_available()
     else ()
 )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:List[Any]									=			FlaxDistilBertModelTester(self       )




 @slow
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       for model_class_name in self.all_model_classes:
             a		:int									=			model_class_name.from_pretrained('''distilbert-base-uncased'''       )
             a		:List[str]									=			model(np.ones((1, 1)       )       )
             self.assertIsNotNone(_lowerCamelCase       )







@require_flax
class 						_snake_case    ( unittest.TestCase				):




 @slow
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:List[str]									=			FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased'''       )
       a		:Optional[Any]									=			np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]       )
       a		:List[Any]									=			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
       a		:List[str]									=			model(_lowerCamelCase   ,    attention_mask=_lowerCamelCase       )[0]
       a		:Union[str, Any]									=			(1, 11, 768)
       self.assertEqual(output.shape   ,    _lowerCamelCase       )
       a		:int									=			np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]       )

       self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]   ,    _lowerCamelCase   ,    atol=1e-4       )       )





