


'''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 functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


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

a_						:       Dict      			=				{
    """microsoft/unispeech-large-1500h-cv""": (
        """https://huggingface.co/microsoft/unispeech-large-1500h-cv/resolve/main/config.json"""
    ),
    # See all UniSpeech models at https://huggingface.co/models?filter=unispeech
}





class   snake_case					(					lowercase    ):





  """simple docstring"""




  _lowerCamelCase						=						"unispeech"






  def __init__(					self    ,						UpperCamelCase=32    ,						UpperCamelCase=768    ,						UpperCamelCase=12    ,						UpperCamelCase=12    ,						UpperCamelCase=3072    ,						UpperCamelCase="gelu"    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.0    ,						UpperCamelCase=0.0    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.02    ,						UpperCamelCase=1e-5    ,						UpperCamelCase="group"    ,						UpperCamelCase="gelu"    ,						UpperCamelCase=(512, 512, 512, 512, 512, 512, 512)    ,						UpperCamelCase=(5, 2, 2, 2, 2, 2, 2)    ,						UpperCamelCase=(10, 3, 3, 3, 3, 2, 2)    ,						UpperCamelCase=False    ,						UpperCamelCase=128    ,						UpperCamelCase=16    ,						UpperCamelCase=False    ,						UpperCamelCase=True    ,						UpperCamelCase=0.05    ,						UpperCamelCase=10    ,						UpperCamelCase=2    ,						UpperCamelCase=0.0    ,						UpperCamelCase=10    ,						UpperCamelCase=0    ,						UpperCamelCase=320    ,						UpperCamelCase=2    ,						UpperCamelCase=0.1    ,						UpperCamelCase=100    ,						UpperCamelCase=256    ,						UpperCamelCase=256    ,						UpperCamelCase=0.1    ,						UpperCamelCase="mean"    ,						UpperCamelCase=False    ,						UpperCamelCase=False    ,						UpperCamelCase=256    ,						UpperCamelCase=80    ,						UpperCamelCase=0    ,						UpperCamelCase=1    ,						UpperCamelCase=2    ,						UpperCamelCase=0.5    ,						**UpperCamelCase    ,						):







   """simple docstring"""
   super().__init__(**UpperCamelCase    ,						pad_token_id=UpperCamelCase    ,						bos_token_id=UpperCamelCase    ,						eos_token_id=UpperCamelCase  )
   lowerCamelCase_  				=						hidden_size
   lowerCamelCase_  				=						feat_extract_norm
   lowerCamelCase_  				=						feat_extract_activation
   lowerCamelCase_  				=						list(UpperCamelCase  )
   lowerCamelCase_  				=						list(UpperCamelCase  )
   lowerCamelCase_  				=						list(UpperCamelCase  )
   lowerCamelCase_  				=						conv_bias
   lowerCamelCase_  				=						num_conv_pos_embeddings
   lowerCamelCase_  				=						num_conv_pos_embedding_groups
   lowerCamelCase_  				=						len(self.conv_dim  )
   lowerCamelCase_  				=						num_hidden_layers
   lowerCamelCase_  				=						intermediate_size
   lowerCamelCase_  				=						hidden_act
   lowerCamelCase_  				=						num_attention_heads
   lowerCamelCase_  				=						hidden_dropout
   lowerCamelCase_  				=						attention_dropout
   lowerCamelCase_  				=						activation_dropout
   lowerCamelCase_  				=						feat_proj_dropout
   lowerCamelCase_  				=						final_dropout
   lowerCamelCase_  				=						layerdrop
   lowerCamelCase_  				=						layer_norm_eps
   lowerCamelCase_  				=						initializer_range
   lowerCamelCase_  				=						num_ctc_classes
   lowerCamelCase_  				=						vocab_size
   lowerCamelCase_  				=						do_stable_layer_norm
   lowerCamelCase_  				=						use_weighted_layer_sum
   lowerCamelCase_  				=						classifier_proj_size

   if (
       (len(self.conv_stride  ) != self.num_feat_extract_layers)
       or (len(self.conv_kernel  ) != self.num_feat_extract_layers)
       or (len(self.conv_dim  ) != self.num_feat_extract_layers)
   ):
    raise ValueError(
        "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
        " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
        f''' {len(self.conv_dim  )}`, `len(config.conv_stride) = {len(self.conv_stride  )}`,'''
        f''' `len(config.conv_kernel) = {len(self.conv_kernel  )}`.'''  )

   # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
   lowerCamelCase_  				=						apply_spec_augment
   lowerCamelCase_  				=						mask_time_prob
   lowerCamelCase_  				=						mask_time_length
   lowerCamelCase_  				=						mask_time_min_masks
   lowerCamelCase_  				=						mask_feature_prob
   lowerCamelCase_  				=						mask_feature_length
   lowerCamelCase_  				=						mask_feature_min_masks

   # parameters for pretraining with codevector quantized representations
   lowerCamelCase_  				=						num_codevectors_per_group
   lowerCamelCase_  				=						num_codevector_groups
   lowerCamelCase_  				=						contrastive_logits_temperature
   lowerCamelCase_  				=						feat_quantizer_dropout
   lowerCamelCase_  				=						num_negatives
   lowerCamelCase_  				=						codevector_dim
   lowerCamelCase_  				=						proj_codevector_dim
   lowerCamelCase_  				=						diversity_loss_weight

   # ctc loss
   lowerCamelCase_  				=						ctc_loss_reduction
   lowerCamelCase_  				=						ctc_zero_infinity

   # pretraining loss
   lowerCamelCase_  				=						replace_prob







  @property
  def 						snake_case	(					self  ):







   """simple docstring"""
   return functools.reduce(operator.mul    ,						self.conv_stride    ,						1  )



