


'''simple docstring'''






import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


logging.set_verbosity_info()
__A   :    List[Any]											= logging.get_logger(__name__)

__A   :    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",
    "adapter_layer": "encoder.layers.*.adapter_layer",
    "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",
    "mask_emb": "masked_spec_embed",
    "pooling_layer.linear": "projector",
    "pooling_layer.projection": "classifier",
}
__A   :    List[str]											= [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
    "projector",
    "classifier",
]







def 							UpperCamelCase_     (       A__   :  str   ):
  '''simple docstring'''





  lowerCAmelCase_     : Dict     =  {}
  with open(A__				,       """r"""   ) as file:
    for line_number, line in enumerate(A__   ):
      lowerCAmelCase_     : Dict     =  line.strip()
      if line:
        lowerCAmelCase_     : Union[str, Any]     =  line.split()
        lowerCAmelCase_     : Dict     =  line_number
        lowerCAmelCase_     : Any     =  words[0]
        lowerCAmelCase_     : int     =  value
  return result







def 							UpperCamelCase_     (       A__   :  Optional[Any]				,       A__   :  List[Any]				,       A__   :  Dict				,       A__   :  Union[str, Any]				,       A__   :  Dict   ):
  '''simple docstring'''





  for attribute in key.split("""."""   ):
    lowerCAmelCase_     : int     =  getattr(A__				,       A__   )

  lowerCAmelCase_     : Any     =  None
  for param_key in PARAM_MAPPING.keys():
    if full_name.endswith(A__   ):
      lowerCAmelCase_     : int     =  PARAM_MAPPING[full_name.split("""."""   )[-1]]
      lowerCAmelCase_     : List[Any]     =  """param"""

  if weight_type is not None and weight_type != "param":
    lowerCAmelCase_     : int     =  getattr(A__				,       A__   ).shape
  elif weight_type is not None and weight_type == "param":
    lowerCAmelCase_     : int     =  hf_pointer
    for attribute in hf_param_name.split("""."""   ):
      lowerCAmelCase_     : Union[str, Any]     =  getattr(A__				,       A__   )
    lowerCAmelCase_     : List[str]     =  shape_pointer.shape

    # let's reduce dimension
    lowerCAmelCase_     : int     =  value[0]
  else:
    lowerCAmelCase_     : Any     =  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":
    lowerCAmelCase_     : Tuple     =  value
  elif weight_type == "weight_g":
    lowerCAmelCase_     : int     =  value
  elif weight_type == "weight_v":
    lowerCAmelCase_     : List[str]     =  value
  elif weight_type == "bias":
    lowerCAmelCase_     : Union[str, Any]     =  value
  elif weight_type == "param":
    for attribute in hf_param_name.split("""."""   ):
      lowerCAmelCase_     : List[Any]     =  getattr(A__				,       A__   )
    lowerCAmelCase_     : List[Any]     =  value
  else:
    lowerCAmelCase_     : Any     =  value

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







def 							UpperCamelCase_     (       A__   :  List[str]				,       A__   :  Optional[Any]				,       A__   :  Optional[int]				,       A__   :  Any				,       A__   :  List[str]   ):
  '''simple docstring'''





  lowerCAmelCase_     : Optional[int]     =  None
  for param_key in PARAM_MAPPING.keys():
    if full_name.endswith(A__   ):
      lowerCAmelCase_     : List[Any]     =  PARAM_MAPPING[full_name.split("""."""   )[-1]]
      lowerCAmelCase_     : Optional[Any]     =  """param"""

  if weight_type is not None and weight_type != "param":
    lowerCAmelCase_     : Any     =  """.""".join([key, weight_type]   )
  elif weight_type is not None and weight_type == "param":
    lowerCAmelCase_     : Union[str, Any]     =  """.""".join([key, hf_param_name]   )
  else:
    lowerCAmelCase_     : Union[str, Any]     =  key

  lowerCAmelCase_     : List[Any]     =  value if """lm_head""" in full_key else value[0]


__A   :    Optional[int]											= {
    "W_a": "linear_1.weight",
    "W_b": "linear_2.weight",
    "b_a": "linear_1.bias",
    "b_b": "linear_2.bias",
    "ln_W": "norm.weight",
    "ln_b": "norm.bias",
}







def 							UpperCamelCase_     (       A__   :  Any				,       A__   :  List[Any]				,       A__   :  Optional[Any]=None				,       A__   :  List[str]=None   ):
  '''simple docstring'''





  lowerCAmelCase_     : str     =  False
  for key, mapped_key in MAPPING.items():
    lowerCAmelCase_     : List[Any]     =  """wav2vec2.""" + 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_     : Optional[int]     =  True
      if "*" in mapped_key:
        lowerCAmelCase_     : Any     =  name.split(A__   )[0].split("""."""   )[-2]
        lowerCAmelCase_     : str     =  mapped_key.replace("""*"""				,       A__   )
      if "weight_g" in name:
        lowerCAmelCase_     : Optional[int]     =  """weight_g"""
      elif "weight_v" in name:
        lowerCAmelCase_     : Optional[Any]     =  """weight_v"""
      elif "bias" in name:
        lowerCAmelCase_     : Dict     =  """bias"""
      elif "weight" in name:
        # TODO: don't match quantizer.weight_proj
        lowerCAmelCase_     : Optional[Any]     =  """weight"""
      else:
        lowerCAmelCase_     : List[Any]     =  None
      if hf_dict is not None:
        rename_dict(A__				,       A__				,       A__				,       A__				,       A__   )
      else:
        set_recursively(A__				,       A__				,       A__				,       A__				,       A__   )
      return is_used
  return is_used







def 							UpperCamelCase_     (       A__   :  List[Any]				,       A__   :  int				,       A__   :  List[str]   ):
  '''simple docstring'''





  lowerCAmelCase_     : Tuple     =  []
  lowerCAmelCase_     : Optional[Any]     =  fairseq_model.state_dict()

  lowerCAmelCase_     : Tuple     =  hf_model.wavaveca.feature_extractor

  for name, value in fairseq_dict.items():
    lowerCAmelCase_     : int     =  False
    if "conv_layers" in name:
      load_conv_layer(
          A__				,       A__				,       A__				,       A__				,       hf_model.config.feat_extract_norm == """group"""				,       )
      lowerCAmelCase_     : List[Any]     =  True
    else:
      lowerCAmelCase_     : int     =  load_wavaveca_layer(A__				,       A__				,       A__   )
    if not is_used:
      unused_weights.append(A__   )

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







def 							UpperCamelCase_     (       A__   :  Optional[Any]				,       A__   :  Tuple				,       A__   :  Optional[int]				,       A__   :  List[str]				,       A__   :  Optional[Any]   ):
  '''simple docstring'''





  lowerCAmelCase_     : Optional[int]     =  full_name.split("""conv_layers."""   )[-1]
  lowerCAmelCase_     : Tuple     =  name.split("""."""   )
  lowerCAmelCase_     : Tuple     =  int(items[0]   )
  lowerCAmelCase_     : Any     =  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.'   )
      lowerCAmelCase_     : Any     =  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.'   )
      lowerCAmelCase_     : List[str]     =  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.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'   )
      lowerCAmelCase_     : Dict     =  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.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'   )
      lowerCAmelCase_     : List[str]     =  value
      logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'   )
  else:
    unused_weights.append(A__   )







@torch.no_grad()
def 							UpperCamelCase_     (       A__   :  Tuple				,       A__   :  Union[str, Any]				,       A__   :  str=None				,       A__   :  List[Any]=None				,       A__   :  Union[str, Any]=True				,       A__   :  Any=False   ):
  '''simple docstring'''





  if config_path is not None:
    lowerCAmelCase_     : Optional[int]     =  WavaVecaConfig.from_pretrained(A__   )
  else:
    lowerCAmelCase_     : Tuple     =  WavaVecaConfig()

  if is_seq_class:
    lowerCAmelCase_     : Any     =  read_txt_into_dict(A__   )
    lowerCAmelCase_     : Dict     =  idalabel
    lowerCAmelCase_     : Union[str, Any]     =  WavaVecaForSequenceClassification(A__   )
    lowerCAmelCase_     : str     =  WavaVecaFeatureExtractor(
        feature_size=1				,       sampling_rate=1_60_00				,       padding_value=0				,       do_normalize=A__				,       return_attention_mask=A__				,       )
    feature_extractor.save_pretrained(A__   )

  elif is_finetuned:
    if dict_path:
      lowerCAmelCase_     : Optional[Any]     =  Dictionary.load(A__   )

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

      # fairseq has the <pad> and <s> switched
      lowerCAmelCase_     : List[Any]     =  0
      lowerCAmelCase_     : Optional[Any]     =  1
      with open(A__				,       """w"""				,       encoding="""utf-8"""   ) as vocab_handle:
        json.dump(A__				,       A__   )
      lowerCAmelCase_     : str     =  WavaVecaCTCTokenizer(
          A__				,       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=A__				,       )
      lowerCAmelCase_     : Dict     =  True if config.feat_extract_norm == """layer""" else False
      lowerCAmelCase_     : Any     =  WavaVecaFeatureExtractor(
          feature_size=1				,       sampling_rate=1_60_00				,       padding_value=0				,       do_normalize=A__				,       return_attention_mask=A__				,       )
      lowerCAmelCase_     : Optional[Any]     =  WavaVecaProcessor(feature_extractor=A__				,       tokenizer=A__   )
      processor.save_pretrained(A__   )

    lowerCAmelCase_     : Any     =  WavaVecaForCTC(A__   )
  else:
    lowerCAmelCase_     : Dict     =  WavaVecaForPreTraining(A__   )

  if is_finetuned or is_seq_class:
    lowerCAmelCase_,				lowerCAmelCase_,				lowerCAmelCase_     : List[str]     =  fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [checkpoint_path]				,       arg_overrides={"""data""": """/""".join(dict_path.split("""/"""   )[:-1]   )}   )
  else:
    lowerCAmelCase_     : Dict     =  argparse.Namespace(task="""audio_pretraining"""   )
    lowerCAmelCase_     : Dict     =  fairseq.tasks.setup_task(A__   )

    lowerCAmelCase_,				lowerCAmelCase_,				lowerCAmelCase_     : Any     =  fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]				,       task=A__   )

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

  recursively_load_weights(A__				,       A__				,       not is_finetuned   )

  hf_wavavec.save_pretrained(A__   )


if __name__ == "__main__":
  __A   :    str											= 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"
  )
  parser.add_argument(
      "--is_seq_class",
      action="store_true",
      help="Whether the model to convert is a fine-tuned sequence classification model or not",
  )
  __A   :    List[str]											= parser.parse_args()

  __A   :    Optional[int]											= not args.not_finetuned and not args.is_seq_class
  convert_wavaveca_checkpoint(
      args.checkpoint_path,
      args.pytorch_dump_folder_path,
      args.config_path,
      args.dict_path,
      is_finetuned,
      args.is_seq_class,
  )







'''simple docstring'''






def 							UpperCamelCase_     (       A__   :  bytes   ):
  '''simple docstring'''





  return "".join([hex(A__   )[2:].zfill(2   ).upper() for byte in list(A__   )]   )







def 							UpperCamelCase_     (       A__   :  str   ):
  '''simple docstring'''





  if (len(A__   ) % 2) != 0:
    raise ValueError(
        """Base16 encoded data is invalid:
Data does not have an even number of hex digits."""   )
  # Check the character set - the standard base16 alphabet
  # is uppercase according to RFC3548 section 6
  if not set(A__   ) <= set("""0123456789ABCDEF"""   ):
    raise ValueError(
        """Base16 encoded data is invalid:
Data is not uppercase hex or it contains invalid characters."""   )
  # For every two hexadecimal digits (= a byte), turn it into an integer.
  # Then, string the result together into bytes, and return it.
  return bytes(int(data[i] + data[i + 1]				,       16   ) for i in range(0				,       len(A__   )				,       2   )   )


if __name__ == "__main__":
  import doctest

  doctest.testmod()



