





import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


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

SCREAMING_SNAKE_CASE_						   =       {
    """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""",
    """encoder.layer_norm_for_extract""": """layer_norm_for_extract""",
    """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""",
    """label_embs_concat""": """label_embeddings_concat""",
    """mask_emb""": """masked_spec_embed""",
    """spk_proj""": """speaker_proj""",
}
SCREAMING_SNAKE_CASE_						   =       [
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
    """label_embeddings_concat""",
    """speaker_proj""",
    """layer_norm_for_extract""",
]



def    __lowercase   (      _SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	) ->   Any:


      '''simple docstring'''


      for attribute in key.split("""."""	):
            SCREAMING_SNAKE_CASE		=						getattr(_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	)

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

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



def    __lowercase   (      _SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	) ->   Any:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						[]
      SCREAMING_SNAKE_CASE		=						fairseq_model.state_dict()

      SCREAMING_SNAKE_CASE		=						hf_model.unispeech_sat.feature_extractor

      for name, value in fairseq_dict.items():
            SCREAMING_SNAKE_CASE		=						False
            if "conv_layers" in name:
                  load_conv_layer(
                      _SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						hf_model.config.feat_extract_norm == """group""" ,						)
                  SCREAMING_SNAKE_CASE		=						True
            else:
                  for key, mapped_key in MAPPING.items():
                        SCREAMING_SNAKE_CASE		=						"""unispeech_sat.""" + 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]:
                              if "layer_norm_for_extract" in name and (".".join(name.split("""."""	)[:-1]	) != key):
                                    # special case since naming is very similar
                                    continue
                              SCREAMING_SNAKE_CASE		=						True
                              if "*" in mapped_key:
                                    SCREAMING_SNAKE_CASE		=						name.split(_SCREAMING_SNAKE_CASE	)[0].split("""."""	)[-2]
                                    SCREAMING_SNAKE_CASE		=						mapped_key.replace("""*""" ,						_SCREAMING_SNAKE_CASE	)
                              if "weight_g" in name:
                                    SCREAMING_SNAKE_CASE		=						"""weight_g"""
                              elif "weight_v" in name:
                                    SCREAMING_SNAKE_CASE		=						"""weight_v"""
                              elif "bias" in name:
                                    SCREAMING_SNAKE_CASE		=						"""bias"""
                              elif "weight" in name:
                                    # TODO: don't match quantizer.weight_proj
                                    SCREAMING_SNAKE_CASE		=						"""weight"""
                              else:
                                    SCREAMING_SNAKE_CASE		=						None
                              set_recursively(_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	)
                        continue
            if not is_used:
                  unused_weights.append(_SCREAMING_SNAKE_CASE	)

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



def    __lowercase   (      _SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	) ->   Optional[int]:


      '''simple docstring'''


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



@torch.no_grad()
def    __lowercase   (      _SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE=None ,						_SCREAMING_SNAKE_CASE=None ,						_SCREAMING_SNAKE_CASE=True	) ->   List[Any]:


      '''simple docstring'''


      if config_path is not None:
            SCREAMING_SNAKE_CASE		=						UniSpeechSatConfig.from_pretrained(_SCREAMING_SNAKE_CASE	)
      else:
            SCREAMING_SNAKE_CASE		=						UniSpeechSatConfig()

      SCREAMING_SNAKE_CASE		=						""""""

      if is_finetuned:
            SCREAMING_SNAKE_CASE		=						UniSpeechSatForCTC(_SCREAMING_SNAKE_CASE	)
      else:
            SCREAMING_SNAKE_CASE		=						UniSpeechSatForPreTraining(_SCREAMING_SNAKE_CASE	)

      SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE		=						fairseq.checkpoint_utils.load_model_ensemble_and_task(
          [checkpoint_path] ,						arg_overrides={"""data""": """/""".join(dict_path.split("""/"""	)[:-1]	)}	)
      SCREAMING_SNAKE_CASE		=						model[0].eval()

      recursively_load_weights(_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	)

      hf_wavavec.save_pretrained(_SCREAMING_SNAKE_CASE	)


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









import importlib
import shutil
import threading
import warnings
from typing import List

import fsspec
import fsspec.asyn

from . import compression
from .hffilesystem import HfFileSystem


SCREAMING_SNAKE_CASE_						   =       importlib.util.find_spec("""s3fs""") is not None

if _has_safs:
      from .safilesystem import SaFileSystem  # noqa: F401

SCREAMING_SNAKE_CASE_						   =       [
    compression.BzaFileSystem,
    compression.GzipFileSystem,
    compression.LzaFileSystem,
    compression.XzFileSystem,
    compression.ZstdFileSystem,
]

# Register custom filesystems
for fs_class in COMPRESSION_FILESYSTEMS + [HfFileSystem]:
      if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
            warnings.warn(F'''A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.''')
      fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)



def    __lowercase   (      _SCREAMING_SNAKE_CASE	) ->   str:


      '''simple docstring'''


      if "://" in dataset_path:
            SCREAMING_SNAKE_CASE		=						dataset_path.split("""://"""	)[1]
      return dataset_path



def    __lowercase   (      _SCREAMING_SNAKE_CASE	) ->   bool:


      '''simple docstring'''


      if fs is not None and fs.protocol != "file":
            return True
      else:
            return False



def    __lowercase   (      _SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE	) ->   Dict:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						not is_remote_filesystem(_SCREAMING_SNAKE_CASE	)
      if is_local:
            # LocalFileSystem.mv does copy + rm, it is more efficient to simply move a local directory
            shutil.move(fs._strip_protocol(_SCREAMING_SNAKE_CASE	) ,						fs._strip_protocol(_SCREAMING_SNAKE_CASE	)	)
      else:
            fs.mv(_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE ,						recursive=_SCREAMING_SNAKE_CASE	)



def    __lowercase   (      ) ->   None:


      '''simple docstring'''


      if hasattr(fsspec.asyn ,						"""reset_lock"""	):
            # for future fsspec>2022.05.0
            fsspec.asyn.reset_lock()
      else:
            SCREAMING_SNAKE_CASE		=						None
            SCREAMING_SNAKE_CASE		=						None
            SCREAMING_SNAKE_CASE		=						threading.Lock()


