


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

A__						:   List[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__						:   str				=    [
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
    """projector""",
    """classifier""",
]



def      UpperCAmelCase__   (      UpperCAmelCase_       :							List[Any]				)	->						Tuple:
     __lowerCamelCase  :      Tuple														=						{}
     with open(UpperCAmelCase_ ,       'r'				) as file:
          for line_number, line in enumerate(UpperCAmelCase_				):
               __lowerCamelCase  :      Any														=						line.strip()
               if line:
                    __lowerCamelCase  :      Optional[int]														=						line.split()
                    __lowerCamelCase  :      List[Any]														=						line_number
                    __lowerCamelCase  :      int														=						words[0]
                    __lowerCamelCase  :      List[str]														=						value
     return result



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Optional[Any] ,       UpperCAmelCase_       :							Optional[Any] ,       UpperCAmelCase_       :							Dict ,       UpperCAmelCase_       :							Tuple ,       UpperCAmelCase_       :							Any				)	->						Dict:
     for attribute in key.split('.'				):
          __lowerCamelCase  :      Optional[int]														=						getattr(UpperCAmelCase_ ,       UpperCAmelCase_				)

     __lowerCamelCase  :      Any														=						None
     for param_key in PARAM_MAPPING.keys():
          if full_name.endswith(UpperCAmelCase_				):
               __lowerCamelCase  :      Dict														=						PARAM_MAPPING[full_name.split('.'				)[-1]]
               __lowerCamelCase  :      str														=						'param'

     if weight_type is not None and weight_type != "param":
          __lowerCamelCase  :      List[Any]														=						getattr(UpperCAmelCase_ ,       UpperCAmelCase_				).shape
     elif weight_type is not None and weight_type == "param":
          __lowerCamelCase  :      Any														=						hf_pointer
          for attribute in hf_param_name.split('.'				):
               __lowerCamelCase  :      Optional[int]														=						getattr(UpperCAmelCase_ ,       UpperCAmelCase_				)
          __lowerCamelCase  :      List[Any]														=						shape_pointer.shape

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

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



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Dict ,       UpperCAmelCase_       :							Tuple ,       UpperCAmelCase_       :							Optional[Any] ,       UpperCAmelCase_       :							Dict ,       UpperCAmelCase_       :							int				)	->						Dict:
     __lowerCamelCase  :      Dict														=						None
     for param_key in PARAM_MAPPING.keys():
          if full_name.endswith(UpperCAmelCase_				):
               __lowerCamelCase  :      Tuple														=						PARAM_MAPPING[full_name.split('.'				)[-1]]
               __lowerCamelCase  :      Optional[Any]														=						'param'

     if weight_type is not None and weight_type != "param":
          __lowerCamelCase  :      List[str]														=						'.'.join([key, weight_type]				)
     elif weight_type is not None and weight_type == "param":
          __lowerCamelCase  :      str														=						'.'.join([key, hf_param_name]				)
     else:
          __lowerCamelCase  :      Dict														=						key

     __lowerCamelCase  :      str														=						value if 'lm_head' in full_key else value[0]


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



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Any ,       UpperCAmelCase_       :							Optional[int] ,       UpperCAmelCase_       :							str				)	->						Any:
     __lowerCamelCase  :      Optional[int]														=						[]
     __lowerCamelCase  :      Union[str, Any]														=						fairseq_model.state_dict()

     __lowerCamelCase  :      int														=						hf_model.wavaveca.feature_extractor

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

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



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Tuple ,       UpperCAmelCase_       :							str ,       UpperCAmelCase_       :							Tuple ,       UpperCAmelCase_       :							List[str] ,       UpperCAmelCase_       :							Tuple				)	->						Union[str, Any]:
     __lowerCamelCase  :      Tuple														=						full_name.split('conv_layers.'				)[-1]
     __lowerCamelCase  :      str														=						name.split('.'				)
     __lowerCamelCase  :      Optional[int]														=						int(items[0]				)
     __lowerCamelCase  :      List[str]														=						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  :      List[str]														=						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  :      Dict														=						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  :      List[str]														=						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  :      Union[str, Any]														=						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      UpperCAmelCase__   (      UpperCAmelCase_       :							Union[str, Any] ,       UpperCAmelCase_       :							Dict ,       UpperCAmelCase_       :							Union[str, Any]=None ,       UpperCAmelCase_       :							List[str]=None ,       UpperCAmelCase_       :							Union[str, Any]=True ,       UpperCAmelCase_       :							List[str]=False				)	->						Optional[Any]:
     if config_path is not None:
          __lowerCamelCase  :      str														=						WavaVecaConfig.from_pretrained(UpperCAmelCase_				)
     else:
          __lowerCamelCase  :      Any														=						WavaVecaConfig()

     if is_seq_class:
          __lowerCamelCase  :      List[str]														=						read_txt_into_dict(UpperCAmelCase_				)
          __lowerCamelCase  :      List[Any]														=						idalabel
          __lowerCamelCase  :      str														=						WavaVecaForSequenceClassification(UpperCAmelCase_				)
          __lowerCamelCase  :      int														=						WavaVecaFeatureExtractor(
              feature_size=1 ,       sampling_rate=1_60_00 ,       padding_value=0 ,       do_normalize=UpperCAmelCase_ ,       return_attention_mask=UpperCAmelCase_ ,       )
          feature_extractor.save_pretrained(UpperCAmelCase_				)

     elif is_finetuned:
          if dict_path:
               __lowerCamelCase  :      str														=						Dictionary.load(UpperCAmelCase_				)

               # important change bos & pad token id since CTC symbol is <pad> and
               # not <s> as in fairseq
               __lowerCamelCase  :      Any														=						target_dict.pad_index
               __lowerCamelCase  :      int														=						target_dict.bos_index
               __lowerCamelCase  :      int														=						target_dict.eos_index
               __lowerCamelCase  :      Optional[Any]														=						len(target_dict.symbols				)
               __lowerCamelCase  :      Tuple														=						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  :      Tuple														=						target_dict.indices

               # fairseq has the <pad> and <s> switched
               __lowerCamelCase  :      Union[str, Any]														=						0
               __lowerCamelCase  :      Optional[Any]														=						1
               with open(UpperCAmelCase_ ,       'w' ,       encoding='utf-8'				) as vocab_handle:
                    json.dump(UpperCAmelCase_ ,       UpperCAmelCase_				)
               __lowerCamelCase  :      str														=						WavaVecaCTCTokenizer(
                   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  :      List[str]														=						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=UpperCAmelCase_ ,       return_attention_mask=UpperCAmelCase_ ,       )
               __lowerCamelCase  :      int														=						WavaVecaProcessor(feature_extractor=UpperCAmelCase_ ,       tokenizer=UpperCAmelCase_				)
               processor.save_pretrained(UpperCAmelCase_				)

          __lowerCamelCase  :      Any														=						WavaVecaForCTC(UpperCAmelCase_				)
     else:
          __lowerCamelCase  :      List[str]														=						WavaVecaForPreTraining(UpperCAmelCase_				)

     if is_finetuned or is_seq_class:
          __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase  :      Any														=						fairseq.checkpoint_utils.load_model_ensemble_and_task(
              [checkpoint_path] ,       arg_overrides={'data': '/'.join(dict_path.split('/'				)[:-1]				)}				)
     else:
          __lowerCamelCase  :      Union[str, Any]														=						argparse.Namespace(task='audio_pretraining'				)
          __lowerCamelCase  :      Optional[int]														=						fairseq.tasks.setup_task(UpperCAmelCase_				)

          __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase  :      List[str]														=						fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path] ,       task=UpperCAmelCase_				)

     __lowerCamelCase  :      List[str]														=						model[0].eval()

     recursively_load_weights(UpperCAmelCase_ ,       UpperCAmelCase_ ,       not is_finetuned				)

     hf_wavavec.save_pretrained(UpperCAmelCase_				)


if __name__ == "__main__":
      A__						:   Dict				=    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__						:   int				=    parser.parse_args()

      A__						:   str				=    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'''




import os
from collections.abc import Iterator





def 	__lowerCamelCase			(      __lowerCAmelCase       :							str = "."		)		->  Iterator[str]:
   for dir_path, dir_names, filenames in os.walk(__lowerCAmelCase		):
      snake_case 						=   [d for d in dir_names if d != """scripts""" and d[0] not in """._"""]
      for filename in filenames:
         if filename == "__init__.py":
            continue
         if os.path.splitext(__lowerCAmelCase		)[1] in (".py", ".ipynb"):
            yield os.path.join(__lowerCAmelCase      ,    __lowerCAmelCase		).lstrip("""./"""		)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]		)		->  Any:
   return F'''{i * "  "}*''' if i else "\n##"





def 	__lowerCamelCase			(      __lowerCAmelCase       :							str      ,    __lowerCAmelCase       :							str		)		->  str:
   snake_case 						=   old_path.split(os.sep		)
   for i, new_part in enumerate(new_path.split(os.sep		)		):
      if (i + 1 > len(__lowerCAmelCase		) or old_parts[i] != new_part) and new_part:
         print(F'''{md_prefix(__lowerCAmelCase		)} {new_part.replace("_"      ,    " "		).title()}'''		)
   return new_path





def 	__lowerCamelCase			(      __lowerCAmelCase       :							str = "."		)		->  None:
   snake_case 						=   """"""
   for filepath in sorted(good_file_paths(__lowerCAmelCase		)		):
      snake_case						,    snake_case 						=   os.path.split(__lowerCAmelCase		)
      if filepath != old_path:
         snake_case 						=   print_path(__lowerCAmelCase      ,    __lowerCAmelCase		)
      snake_case 						=   (filepath.count(os.sep		) + 1) if filepath else 0
      snake_case 						=   F'''{filepath}/{filename}'''.replace(""" """      ,    """%20"""		)
      snake_case 						=   os.path.splitext(filename.replace("""_"""      ,    """ """		).title()		)[0]
      print(F'''{md_prefix(__lowerCAmelCase		)} [{filename}]({url})'''		)


if __name__ == "__main__":
 print_directory_md(".")

