






"""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()
snake_case : Union[str, Any]     		=		logging.get_logger(__name__)

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


def 							A						(				__snake_case: Dict ,     __snake_case: str ,     __snake_case: Any ,     __snake_case: List[str] ,     __snake_case: Union[str, Any] ,     __snake_case: str  )    ->		Optional[Any]:

       """simple docstring"""





       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
                            __magic_name__      				=    'lm_head'

              __magic_name__      				=    getattr(__snake_case ,     __snake_case  )

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

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


def 							A						(				__snake_case: Optional[int] ,     __snake_case: Tuple ,     __snake_case: List[str]  )    ->		Optional[int]:

       """simple docstring"""





       __magic_name__      				=    []
       __magic_name__      				=    fairseq_model.state_dict()

       __magic_name__      				=    hf_model.unispeech.feature_extractor

       for name, value in fairseq_dict.items():
              __magic_name__      				=    False
              if "conv_layers" in name:
                     load_conv_layer(
                         __snake_case ,     __snake_case ,     __snake_case ,     __snake_case ,     hf_model.config.feat_extract_norm == 'group' ,     )
                     __magic_name__      				=    True
              else:
                     for key, mapped_key in MAPPING.items():
                            __magic_name__      				=    '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]:
                                   __magic_name__      				=    True
                                   if "*" in mapped_key:
                                          __magic_name__      				=    name.split(__snake_case  )[0].split('.'  )[-2]
                                          __magic_name__      				=    mapped_key.replace('*' ,     __snake_case  )
                                   if "weight_g" in name:
                                          __magic_name__      				=    'weight_g'
                                   elif "weight_v" in name:
                                          __magic_name__      				=    'weight_v'
                                   elif "bias" in name:
                                          __magic_name__      				=    'bias'
                                   elif "weight" in name:
                                          # TODO: don't match quantizer.weight_proj
                                          __magic_name__      				=    'weight'
                                   else:
                                          __magic_name__      				=    None
                                   set_recursively(__snake_case ,     __snake_case ,     __snake_case ,     __snake_case ,     __snake_case ,     __snake_case  )
                            continue
              if not is_used:
                     unused_weights.append(__snake_case  )

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


def 							A						(				__snake_case: Tuple ,     __snake_case: Dict ,     __snake_case: List[str] ,     __snake_case: List[str] ,     __snake_case: Optional[Any]  )    ->		Dict:

       """simple docstring"""





       __magic_name__      				=    full_name.split('conv_layers.'  )[-1]
       __magic_name__      				=    name.split('.'  )
       __magic_name__      				=    int(items[0]  )
       __magic_name__      				=    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."""
                     )
                     __magic_name__      				=    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."""
                     )
                     __magic_name__      				=    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."
                     )
                     __magic_name__      				=    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."""
                     )
                     __magic_name__      				=    value
                     logger.info(F"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""  )
       else:
              unused_weights.append(__snake_case  )


@torch.no_grad()
def 							A						(				__snake_case: List[str] ,     __snake_case: List[Any] ,     __snake_case: int=None ,     __snake_case: Any=None ,     __snake_case: Tuple=True  )    ->		Dict:

       """simple docstring"""





       if config_path is not None:
              __magic_name__      				=    UniSpeechConfig.from_pretrained(__snake_case  )
       else:
              __magic_name__      				=    UniSpeechConfig()

       if is_finetuned:
              if dict_path:
                     __magic_name__      				=    Dictionary.load_from_json(__snake_case  )

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

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

              __magic_name__      				=    UniSpeechForCTC(__snake_case  )
       else:
              __magic_name__      				=    UniSpeechForPreTraining(__snake_case  )

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

       __magic_name__      				=    model[0].eval()

       recursively_load_weights(__snake_case ,     __snake_case ,     __snake_case  )

       hf_unispeech.save_pretrained(__snake_case  )


if __name__ == "__main__":
 snake_case : 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"""
 )
 snake_case : Optional[Any]     		=		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"""
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case : int     		=		{}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 snake_case : str     		=		["""NllbTokenizer"""]

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 snake_case : Dict     		=		["""NllbTokenizerFast"""]


if TYPE_CHECKING:
 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_nllb import NllbTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_nllb_fast import NllbTokenizerFast

else:
 import sys

 snake_case : Dict     		=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)