


'''simple docstring'''

import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
a_						:       List[str]      			=				logging.get_logger(__name__)

a_						:       Union[str, Any]      			=				{
    """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""": """ctc_proj""",
    """mask_emb""": """masked_spec_embed""",
}
a_						:       Union[str, Any]      			=				[
    """ctc_proj""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]





def __snake_case      (   UpperCAmelCase_		:		List[Any]      ,       UpperCAmelCase_		:		Tuple      ,       UpperCAmelCase_		:		Dict      ,       UpperCAmelCase_		:		Tuple      ,       UpperCAmelCase_		:		List[str]      ,       UpperCAmelCase_		:		Tuple	):
 for attribute in key.split("."	):
  if is_finetuned:
   if attribute in ["quantizer", "project_q", "project_hid"]:
    # those layers are only relevant for pretraining and should be dropped
    return

   if attribute == "ctc_proj":
    # we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
    lowerCamelCase_  				=						"lm_head"

  lowerCamelCase_  				=						getattr(UpperCAmelCase_      ,       UpperCAmelCase_	)

 if weight_type is not None:
  lowerCamelCase_  				=						getattr(UpperCAmelCase_      ,       UpperCAmelCase_	).shape
 else:
  lowerCamelCase_  				=						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":
  lowerCamelCase_  				=						value
 elif weight_type == "weight_g":
  lowerCamelCase_  				=						value
 elif weight_type == "weight_v":
  lowerCamelCase_  				=						value
 elif weight_type == "bias":
  lowerCamelCase_  				=						value
 else:
  lowerCamelCase_  				=						value

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





def __snake_case      (   UpperCAmelCase_		:		str      ,       UpperCAmelCase_		:		Union[str, Any]      ,       UpperCAmelCase_		:		Optional[Any]	):
 lowerCamelCase_  				=						[]
 lowerCamelCase_  				=						fairseq_model.state_dict()

 lowerCamelCase_  				=						hf_model.unispeech.feature_extractor

 for name, value in fairseq_dict.items():
  lowerCamelCase_  				=						False
  if "conv_layers" in name:
   load_conv_layer(
       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       hf_model.config.feat_extract_norm == "group"      ,       )
   lowerCamelCase_  				=						True
  else:
   for key, mapped_key in MAPPING.items():
    lowerCamelCase_  				=						"unispeech." + 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(UpperCAmelCase_	)[0].split("."	)[-2]
      lowerCamelCase_  				=						mapped_key.replace("*"      ,       UpperCAmelCase_	)
     if "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"
     else:
      lowerCamelCase_  				=						None
     set_recursively(UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_	)
    continue
  if not is_used:
   unused_weights.append(UpperCAmelCase_	)

 logger.warning(F'''Unused weights: {unused_weights}'''	)





def __snake_case      (   UpperCAmelCase_		:		str      ,       UpperCAmelCase_		:		Union[str, Any]      ,       UpperCAmelCase_		:		str      ,       UpperCAmelCase_		:		Optional[int]      ,       UpperCAmelCase_		:		List[str]	):
 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:
   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.'''
   )
   lowerCamelCase_  				=						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.'''
   )
   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:
   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."
   )
   lowerCamelCase_  				=						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.'''
   )
   lowerCamelCase_  				=						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 __snake_case      (   UpperCAmelCase_		:		int      ,       UpperCAmelCase_		:		List[Any]      ,       UpperCAmelCase_		:		Optional[int]=None      ,       UpperCAmelCase_		:		int=None      ,       UpperCAmelCase_		:		Union[str, Any]=True	):
 if config_path is not None:
  lowerCamelCase_  				=						UniSpeechConfig.from_pretrained(UpperCAmelCase_	)
 else:
  lowerCamelCase_  				=						UniSpeechConfig()

 if is_finetuned:
  if dict_path:
   lowerCamelCase_  				=						Dictionary.load_from_json(UpperCAmelCase_	)

   # 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(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_	)
   lowerCamelCase_  				=						target_dict.indices

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

  lowerCamelCase_  				=						UniSpeechForCTC(UpperCAmelCase_	)
 else:
  lowerCamelCase_  				=						UniSpeechForPreTraining(UpperCAmelCase_	)

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

 lowerCamelCase_  				=						model[0].eval()

 recursively_load_weights(UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_	)

 hf_unispeech.save_pretrained(UpperCAmelCase_	)


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_						:       int      			=				parser.parse_args()
     convert_unispeech_checkpoint(
         args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
     )







'''simple docstring'''

import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
     import torch

     from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin





class   snake_case					:





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase    ,						UpperCamelCase=None    ,						UpperCamelCase=None    ,						UpperCamelCase=None    ,						UpperCamelCase="resnet50"    ,						UpperCamelCase=3    ,						UpperCamelCase=32    ,						UpperCamelCase=3    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						):







   """simple docstring"""
   lowerCamelCase_  				=						parent
   lowerCamelCase_  				=						out_indices if out_indices is not None else [4]
   lowerCamelCase_  				=						stage_names
   lowerCamelCase_  				=						out_features
   lowerCamelCase_  				=						backbone
   lowerCamelCase_  				=						batch_size
   lowerCamelCase_  				=						image_size
   lowerCamelCase_  				=						num_channels
   lowerCamelCase_  				=						use_pretrained_backbone
   lowerCamelCase_  				=						is_training






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )
   lowerCamelCase_  				=						self.get_config()

   return config, pixel_values






  def 						snake_case	(					self  ):







   """simple docstring"""
   return TimmBackboneConfig(
       image_size=self.image_size    ,						num_channels=self.num_channels    ,						out_features=self.out_features    ,						out_indices=self.out_indices    ,						stage_names=self.stage_names    ,						use_pretrained_backbone=self.use_pretrained_backbone    ,						backbone=self.backbone    ,						)






  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						TimmBackbone(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   with torch.no_grad():
    lowerCamelCase_  				=						model(UpperCamelCase  )
   self.parent.assertEqual(
       result.feature_map[-1].shape    ,						(self.batch_size, model.channels[-1], 14, 14)    ,						)







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.prepare_config_and_inputs()
   lowerCamelCase_				,lowerCamelCase_  				=						config_and_inputs
   lowerCamelCase_  				=						{"pixel_values": pixel_values}
   return config, inputs_dict





@require_torch
@require_timm
class   snake_case					(					lowercase    ,		lowercase    ,		lowercase    ,		unittest.TestCase    ):





  """simple docstring"""




  _lowerCamelCase						=						(TimmBackbone,) if is_torch_available() else ()
  _lowerCamelCase						=						{"feature-extraction": TimmBackbone} if is_torch_available() else {}
  _lowerCamelCase						=						False
  _lowerCamelCase						=						False
  _lowerCamelCase						=						False
  _lowerCamelCase						=						False






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						TimmBackboneModelTester(self  )
   lowerCamelCase_  				=						ConfigTester(self    ,						config_class=UpperCamelCase    ,						has_text_modality=UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						"resnet18"
   lowerCamelCase_  				=						"microsoft/resnet-18"

   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase    ,						use_timm_backbone=UpperCamelCase  )
   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase  )

   self.assertEqual(len(timm_model.out_features  )    ,						len(transformers_model.out_features  )  )
   self.assertEqual(len(timm_model.stage_names  )    ,						len(transformers_model.stage_names  )  )
   self.assertEqual(timm_model.channels    ,						transformers_model.channels  )
   # Out indices are set to the last layer by default. For timm models, we don't know
   # the number of layers in advance, so we set it to (-1,), whereas for transformers
   # models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
   self.assertEqual(timm_model.out_indices    ,						(-1,)  )
   self.assertEqual(transformers_model.out_indices    ,						[len(timm_model.stage_names  ) - 1]  )

   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase    ,						use_timm_backbone=UpperCamelCase    ,						out_indices=[1, 2, 3]  )
   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase    ,						out_indices=[1, 2, 3]  )

   self.assertEqual(timm_model.out_indices    ,						transformers_model.out_indices  )
   self.assertEqual(len(timm_model.out_features  )    ,						len(transformers_model.out_features  )  )
   self.assertEqual(timm_model.channels    ,						transformers_model.channels  )






  @unittest.skip("TimmBackbone doesn't support feed forward chunking"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone doesn't have num_hidden_layers attribute"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone initialization is managed on the timm side"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone models doesn't have inputs_embeds"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone models doesn't have inputs_embeds"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone model cannot be created without specifying a backbone checkpoint"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Only checkpoints on timm can be loaded into TimmBackbone"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("model weights aren't tied in TimmBackbone."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("model weights aren't tied in TimmBackbone."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Only checkpoints on timm can be loaded into TimmBackbone"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Only checkpoints on timm can be loaded into TimmBackbone"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone doesn't have hidden size info in its configuration."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone doesn't support output_attentions."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Safetensors is not supported by timm."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    lowerCamelCase_  				=						inspect.signature(model.forward  )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    lowerCamelCase_  				=						[*signature.parameters.keys()]

    lowerCamelCase_  				=						["pixel_values"]
    self.assertListEqual(arg_names[:1]    ,						UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()
   lowerCamelCase_  				=						True
   lowerCamelCase_  				=						self.has_attentions

   # no need to test all models as different heads yield the same functionality
   lowerCamelCase_  				=						self.all_model_classes[0]
   lowerCamelCase_  				=						model_class(UpperCamelCase  )
   model.to(UpperCamelCase  )

   lowerCamelCase_  				=						self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase  )
   lowerCamelCase_  				=						model(**UpperCamelCase  )
   lowerCamelCase_  				=						outputs[0][-1]

   # Encoder-/Decoder-only models
   lowerCamelCase_  				=						outputs.hidden_states[0]
   hidden_states.retain_grad()

   if self.has_attentions:
    lowerCamelCase_  				=						outputs.attentions[0]
    attentions.retain_grad()

   output.flatten()[0].backward(retain_graph=UpperCamelCase  )

   self.assertIsNotNone(hidden_states.grad  )

   if self.has_attentions:
    self.assertIsNotNone(attentions.grad  )







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    model.to(UpperCamelCase  )
    model.eval()
    lowerCamelCase_  				=						model(**UpperCamelCase  )

    self.assertEqual(len(result.feature_maps  )    ,						len(config.out_indices  )  )
    self.assertEqual(len(model.channels  )    ,						len(config.out_indices  )  )

    # Check output of last stage is taken if out_features=None, out_indices=None
    lowerCamelCase_  				=						copy.deepcopy(UpperCamelCase  )
    lowerCamelCase_  				=						None
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    model.to(UpperCamelCase  )
    model.eval()
    lowerCamelCase_  				=						model(**UpperCamelCase  )

    self.assertEqual(len(result.feature_maps  )    ,						1  )
    self.assertEqual(len(model.channels  )    ,						1  )

    # Check backbone can be initialized with fresh weights
    lowerCamelCase_  				=						copy.deepcopy(UpperCamelCase  )
    lowerCamelCase_  				=						False
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    model.to(UpperCamelCase  )
    model.eval()
    lowerCamelCase_  				=						model(**UpperCamelCase  )



